From 97e59974f3cae5422ac2f2ff280b2ecb009d3c40 Mon Sep 17 00:00:00 2001
From: Josh Cooper <josh@puppetlabs.com>
Date: Fri, 8 Jul 2016 16:27:28 -0700
Subject: [PATCH] (QENG-4070) Consistently return 503 if valid pool is empty

There were several problems with how the pooler checked out vms with
respect to empty pools, invalid pools, and aliases:

- If the vmpooler config did not contain any aliases and the caller
requested a vm from an empty pool or a non-existent one, the vmpooler
would error with:

    NoMethodError - undefined method `[]' for nil:NilClass

If the config contained a non-nil alias section, then:

- If the caller requested a vm from an empty pool and either the vm
didn't have an alias or the aliased pool was empty or non-existent, then
the request for that vm would be silently ignored. The vmpooler would
return 200 if the caller asked for multiple vms and the vmpooler was
able to checkout at least one vm.  Otherwise it would return 404.

- Similarly, if the caller requested a vm from a non-existent pool, then
the request was silently ignored.

This commit adds a `pool_names` Set to the config containing all valid
pool names including aliases. This is used to determine whether a
requested template name is valid or not. This is necessary because redis
does not distinguish between empty and non-existent sets, e.g. the
following returns false in both cases:

    backend.exists('vmpooler__ready__' + key)

If the caller requests a vm (single or multiple), and any vm references
an invalid pool name, we immediately return 404. Otherwise, we know the
request is for valid pool names, since the vmpooler requires a restart
to change pool names and counts.

We then attempt to acquire each vm, trying to match on pool name or
failing back to aliased pool name, as was the previous behavior.

The resulting behavior is:

- If the caller asks for at least one vm from an unknown pool, then
don't try to checkout any vms and respond with 404.
- If the caller asks for a vm, and at least one pool is empty, then
respond with 503, returning checked out vms back to the pool.
- Otherwise return 200 with the list of checked out vms.

This commit also makes `alias` optional again.

This commit also re-enables tests that were merged in from master, but
originally commented out due to the bugs described above..
---
 lib/vmpooler.rb                          |   5 +
 lib/vmpooler/api/v1.rb                   |  68 +++++-----
 spec/helpers.rb                          |   5 +
 spec/vmpooler/api/v1/vm_spec.rb          | 152 +++++++++++++++++++++++
 spec/vmpooler/api/v1/vm_template_spec.rb | 135 ++++++++++++++++++++
 5 files changed, 334 insertions(+), 31 deletions(-)

diff --git a/lib/vmpooler.rb b/lib/vmpooler.rb
index 70bb819..7cfacbb 100644
--- a/lib/vmpooler.rb
+++ b/lib/vmpooler.rb
@@ -10,6 +10,7 @@ module Vmpooler
   require 'time'
   require 'timeout'
   require 'yaml'
+  require 'set'
 
   %w( api graphite logger pool_manager vsphere_helper ).each do |lib|
     begin
@@ -35,15 +36,19 @@ module Vmpooler
     parsed_config[:config]['prefix']       ||= ''
 
     # Create an index of pool aliases
+    parsed_config[:pool_names] = Set.new
     parsed_config[:pools].each do |pool|
+      parsed_config[:pool_names] << pool['name']
       if pool['alias']
         if pool['alias'].kind_of?(Array)
           pool['alias'].each do |a|
             parsed_config[:alias] ||= {}
             parsed_config[:alias][a] = pool['name']
+            parsed_config[:pool_names] << a
           end
         elsif pool['alias'].kind_of?(String)
           parsed_config[:alias][pool['alias']] = pool['name']
+          parsed_config[:pool_names] << pool['alias']
         end
       end
     end
diff --git a/lib/vmpooler/api/v1.rb b/lib/vmpooler/api/v1.rb
index e6c1ea7..f9f29b3 100644
--- a/lib/vmpooler/api/v1.rb
+++ b/lib/vmpooler/api/v1.rb
@@ -20,6 +20,10 @@ module Vmpooler
       Vmpooler::API.settings.config[:pools]
     end
 
+    def pool_exists?(template)
+      Vmpooler::API.settings.config[:pool_names].include?(template)
+    end
+
     def need_auth!
       validate_auth(backend)
     end
@@ -28,25 +32,19 @@ module Vmpooler
       validate_token(backend)
     end
 
-    def alias_deref(hash)
-      newhash = {}
+    def fetch_single_vm(template)
+      vm = backend.spop('vmpooler__ready__' + template)
 
-      hash.each do |key, val|
-        if backend.exists('vmpooler__ready__' + key)
-          newhash[key] = val
-        else
-          if Vmpooler::API.settings.config[:alias][key]
-            newkey = Vmpooler::API.settings.config[:alias][key]
-            newhash[newkey] = val
-          end
-        end
+      return [vm, template] if vm
+
+      aliases = Vmpooler::API.settings.config[:alias]
+      if aliases && aliased_template = aliases[template]
+        vm = backend.spop('vmpooler__ready__' + aliased_template)
+
+        return [vm, aliased_template] if vm
       end
 
-      newhash
-    end
-
-    def fetch_single_vm(template)
-      backend.spop('vmpooler__ready__' + template)
+      [nil, nil]
     end
 
     def return_vm_to_ready_state(template, vm)
@@ -83,33 +81,31 @@ module Vmpooler
     end
 
     def atomically_allocate_vms(payload)
-      return false unless payload and !payload.empty?
-
       result = { 'ok' => false }
       failed = false
       vms = []
 
       payload.each do |template, count|
         count.to_i.times do |_i|
-          vm = fetch_single_vm(template)
+          vm, name = fetch_single_vm(template)
           if !vm
             failed = true
             break
           else
-            vms << [ template, vm ]
+            vms << [ name, vm ]
           end
         end
       end
 
       if failed
-        vms.each do |(template, vm)|
-          return_vm_to_ready_state(template, vm)
-          status 503
+        vms.each do |(name, vm)|
+          return_vm_to_ready_state(name, vm)
         end
+        status 503
       else
-        vms.each do |(template, vm)|
-          account_for_starting_vm(template, vm)
-          update_result_hosts(result, template, vm)
+        vms.each do |(name, vm)|
+          account_for_starting_vm(name, vm)
+          update_result_hosts(result, name, vm)
         end
 
         result['ok'] = true
@@ -371,12 +367,13 @@ module Vmpooler
     end
 
     post "#{api_prefix}/vm/?" do
-      jdata = alias_deref(JSON.parse(request.body.read))
       content_type :json
       result = { 'ok' => false }
 
-      if jdata and !jdata.empty?
-        result = atomically_allocate_vms(jdata)
+      payload = JSON.parse(request.body.read)
+
+      if all_templates_valid?(payload)
+        result = atomically_allocate_vms(payload)
       else
         status 404
       end
@@ -395,12 +392,21 @@ module Vmpooler
       payload
     end
 
+    def all_templates_valid?(payload)
+      return false unless payload
+
+      payload.keys.all? do |templates|
+        pool_exists?(templates)
+      end
+    end
+
     post "#{api_prefix}/vm/:template/?" do
-      payload = alias_deref(extract_templates_from_query_params(params[:template]))
       content_type :json
       result = { 'ok' => false }
 
-      if payload and !payload.empty?
+      payload = extract_templates_from_query_params(params[:template])
+
+      if all_templates_valid?(payload)
         result = atomically_allocate_vms(payload)
       else
         status 404
diff --git a/spec/helpers.rb b/spec/helpers.rb
index d28b761..292b9f4 100644
--- a/spec/helpers.rb
+++ b/spec/helpers.rb
@@ -35,6 +35,7 @@ end
 def create_ready_vm(template, name, token = nil)
   create_vm(name, token)
   redis.sadd("vmpooler__ready__#{template}", name)
+  # REMIND: should be __vm__?
   redis.hset("vmpooler_vm_#{name}", "template", template)
 end
 
@@ -73,3 +74,7 @@ def vm_reverted_to_snapshot?(vm, snapshot = nil)
     instance == vm and (snapshot ? (sha == snapshot) : true)
   end
 end
+
+def pool_has_ready_vm?(pool, vm)
+  !!redis.sismember('vmpooler__ready__' + pool, vm)
+end
diff --git a/spec/vmpooler/api/v1/vm_spec.rb b/spec/vmpooler/api/v1/vm_spec.rb
index 6a4a068..5622468 100644
--- a/spec/vmpooler/api/v1/vm_spec.rb
+++ b/spec/vmpooler/api/v1/vm_spec.rb
@@ -32,6 +32,7 @@ describe Vmpooler::API::V1 do
           {'name' => 'pool2', 'size' => 10}
         ],
         alias: { 'poolone' => 'pool1' },
+        pool_names: [ 'pool1', 'pool2', 'poolone' ]
       }
     }
 
@@ -84,6 +85,31 @@ describe Vmpooler::API::V1 do
         expect_json(ok = false, http = 404)
       end
 
+      it 'returns 503 for empty pool when aliases are not defined' do
+        Vmpooler::API.settings.config.delete(:alias)
+        Vmpooler::API.settings.config[:pool_names] = ['pool1', 'pool2']
+
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        post "#{prefix}/vm/pool1"
+        post "#{prefix}/vm/pool1"
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns 503 for empty pool referenced by alias' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        post "#{prefix}/vm/poolone"
+        post "#{prefix}/vm/poolone"
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
       it 'returns multiple VMs' do
         create_ready_vm 'pool1', 'abcdefghijklmnop'
         create_ready_vm 'pool2', 'qrstuvwxyz012345'
@@ -104,6 +130,132 @@ describe Vmpooler::API::V1 do
         expect(last_response.body).to eq(JSON.pretty_generate(expected))
       end
 
+      it 'returns multiple VMs even when multiple instances from the same pool are requested' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+        create_ready_vm 'pool1', '2abcdefghijklmnop'
+        create_ready_vm 'pool2', 'qrstuvwxyz012345'
+
+        post "#{prefix}/vm", '{"pool1":"2","pool2":"1"}'
+
+        expected = {
+          ok: true,
+          pool1: {
+            hostname: [ '1abcdefghijklmnop', '2abcdefghijklmnop' ]
+          },
+          pool2: {
+            hostname: 'qrstuvwxyz012345'
+          }
+        }
+
+        result = JSON.parse(last_response.body)
+        expect(result['ok']).to eq(true)
+        expect(result['pool1']['hostname']).to include('1abcdefghijklmnop', '2abcdefghijklmnop')
+        expect(result['pool2']['hostname']).to eq('qrstuvwxyz012345')
+
+        expect_json(ok = true, http = 200)
+      end
+
+      it 'returns multiple VMs even when multiple instances from multiple pools are requested' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+        create_ready_vm 'pool1', '2abcdefghijklmnop'
+        create_ready_vm 'pool2', '1qrstuvwxyz012345'
+        create_ready_vm 'pool2', '2qrstuvwxyz012345'
+        create_ready_vm 'pool2', '3qrstuvwxyz012345'
+
+        post "#{prefix}/vm", '{"pool1":"2","pool2":"3"}'
+
+        expected = {
+          ok: true,
+          pool1: {
+            hostname: [ '1abcdefghijklmnop', '2abcdefghijklmnop' ]
+          },
+          pool2: {
+            hostname: [ '1qrstuvwxyz012345', '2qrstuvwxyz012345', '3qrstuvwxyz012345' ]
+          }
+        }
+
+        result = JSON.parse(last_response.body)
+        expect(result['ok']).to eq(true)
+        expect(result['pool1']['hostname']).to include('1abcdefghijklmnop', '2abcdefghijklmnop')
+        expect(result['pool2']['hostname']).to include('1qrstuvwxyz012345', '2qrstuvwxyz012345', '3qrstuvwxyz012345')
+
+        expect_json(ok = true, http = 200)
+      end
+
+      it 'fails when not all requested vms can be allocated' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+
+        post "#{prefix}/vm", '{"pool1":"1","pool2":"1"}'
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns any checked out vms to their pools when not all requested vms can be allocated' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+
+        post "#{prefix}/vm", '{"pool1":"1","pool2":"1"}'
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+
+        expect(pool_has_ready_vm?('pool1', '1abcdefghijklmnop')).to eq(true)
+      end
+
+      it 'fails when not all requested vms can be allocated, when requesting multiple instances from a pool' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+
+        post "#{prefix}/vm", '{"pool1":"2","pool2":"1"}'
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns any checked out vms to their pools when not all requested vms can be allocated, when requesting multiple instances from a pool' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+
+        post "#{prefix}/vm", '{"pool1":"2","pool2":"1"}'
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+
+        expect(pool_has_ready_vm?('pool1', '1abcdefghijklmnop')).to eq(true)
+      end
+
+      it 'fails when not all requested vms can be allocated, when requesting multiple instances from multiple pools' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+
+        post "#{prefix}/vm", '{"pool1":"2","pool2":"3"}'
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns any checked out vms to their pools when not all requested vms can be allocated, when requesting multiple instances from multiple pools' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+        create_ready_vm 'pool1', '2abcdefghijklmnop'
+
+        post "#{prefix}/vm", '{"pool1":"2","pool2":"3"}'
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+
+        expect(pool_has_ready_vm?('pool1', '1abcdefghijklmnop')).to eq(true)
+        expect(pool_has_ready_vm?('pool1', '2abcdefghijklmnop')).to eq(true)
+      end
+
       context '(auth not configured)' do
         it 'does not extend VM lifetime if auth token is provided' do
           app.settings.set :config, auth: false
diff --git a/spec/vmpooler/api/v1/vm_template_spec.rb b/spec/vmpooler/api/v1/vm_template_spec.rb
index 28b85c3..3b7f055 100644
--- a/spec/vmpooler/api/v1/vm_template_spec.rb
+++ b/spec/vmpooler/api/v1/vm_template_spec.rb
@@ -32,6 +32,7 @@ describe Vmpooler::API::V1 do
           {'name' => 'pool2', 'size' => 10}
         ],
         alias: { 'poolone' => 'pool1' },
+        pool_names: [ 'pool1', 'pool2', 'poolone' ]
       }
     }
 
@@ -84,6 +85,31 @@ describe Vmpooler::API::V1 do
         expect_json(ok = false, http = 404)
       end
 
+      it 'returns 503 for empty pool when aliases are not defined' do
+        Vmpooler::API.settings.config.delete(:alias)
+        Vmpooler::API.settings.config[:pool_names] = ['pool1', 'pool2']
+
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        post "#{prefix}/vm/pool1"
+        post "#{prefix}/vm/pool1"
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns 503 for empty pool referenced by alias' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        post "#{prefix}/vm/poolone"
+        post "#{prefix}/vm/poolone"
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
       it 'returns multiple VMs' do
         create_ready_vm 'pool1', 'abcdefghijklmnop'
         create_ready_vm 'pool2', 'qrstuvwxyz012345'
@@ -104,6 +130,115 @@ describe Vmpooler::API::V1 do
         expect(last_response.body).to eq(JSON.pretty_generate(expected))
       end
 
+      it 'returns multiple VMs even when multiple instances from multiple pools are requested' do
+        create_ready_vm 'pool1', '1abcdefghijklmnop'
+        create_ready_vm 'pool1', '2abcdefghijklmnop'
+
+        create_ready_vm 'pool2', '1qrstuvwxyz012345'
+        create_ready_vm 'pool2', '2qrstuvwxyz012345'
+        create_ready_vm 'pool2', '3qrstuvwxyz012345'
+
+        post "#{prefix}/vm/pool1+pool1+pool2+pool2+pool2", ''
+
+        expected = {
+          ok: true,
+          pool1: {
+            hostname: [ '1abcdefghijklmnop', '2abcdefghijklmnop' ]
+          },
+          pool2: {
+            hostname: [ '1qrstuvwxyz012345', '2qrstuvwxyz012345', '3qrstuvwxyz012345' ]
+          }
+        }
+
+        result = JSON.parse(last_response.body)
+        expect(result['ok']).to eq(true)
+        expect(result['pool1']['hostname']).to include('1abcdefghijklmnop', '2abcdefghijklmnop')
+        expect(result['pool2']['hostname']).to include('1qrstuvwxyz012345', '2qrstuvwxyz012345', '3qrstuvwxyz012345')
+        expect_json(ok = true, http = 200)
+      end
+
+      it 'fails when not all requested vms can be allocated' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+
+        post "#{prefix}/vm/pool1+pool2", ''
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns any checked out vms to their pools when not all requested vms can be allocated' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+
+        post "#{prefix}/vm/pool1+pool2", ''
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+
+        expect(pool_has_ready_vm?('pool1', 'abcdefghijklmnop')).to eq(true)
+      end
+
+      it 'fails when not all requested vms can be allocated, when requesting multiple instances from a pool' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        create_ready_vm 'pool1', '0123456789012345'
+
+        allow(redis).to receive(:spop).with('vmpooler__ready__pool1').and_return 'abcdefghijklmnop'
+        allow(redis).to receive(:spop).with('vmpooler__ready__pool2').and_return nil
+        allow(redis).to receive(:sadd).with("vmpooler__ready__pool1", "abcdefghijklmnop")
+
+        post "#{prefix}/vm/pool1+pool1+pool2", ''
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns any checked out vms to their pools when not all requested vms can be allocated, when requesting multiple instances from a pool' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        create_ready_vm 'pool1', '0123456789012345'
+
+        post "#{prefix}/vm/pool1+pool1+pool2", ''
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+
+        expect(pool_has_ready_vm?('pool1', 'abcdefghijklmnop')).to eq(true)
+        expect(pool_has_ready_vm?('pool1', '0123456789012345')).to eq(true)
+      end
+
+      it 'fails when not all requested vms can be allocated, when requesting multiple instances from multiple pools' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        create_ready_vm 'pool2', '0123456789012345'
+
+        post "#{prefix}/vm/pool1+pool1+pool2+pool2+pool2", ''
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+      end
+
+      it 'returns any checked out vms to their pools when not all requested vms can be allocated, when requesting multiple instances from multiple pools' do
+        create_ready_vm 'pool1', 'abcdefghijklmnop'
+        create_ready_vm 'pool2', '0123456789012345'
+
+        post "#{prefix}/vm/pool1+pool1+pool2+pool2+pool2", ''
+
+        expected = { ok: false }
+
+        expect(last_response.body).to eq(JSON.pretty_generate(expected))
+        expect_json(ok = false, http = 503)
+
+        expect(pool_has_ready_vm?('pool1', 'abcdefghijklmnop')).to eq(true)
+        expect(pool_has_ready_vm?('pool2', '0123456789012345')).to eq(true)
+      end
+
       context '(auth not configured)' do
         it 'does not extend VM lifetime if auth token is provided' do
           app.settings.set :config, auth: false