From b1d4f0971c3f40da002ac04180d1212fc2fac27b Mon Sep 17 00:00:00 2001
From: "kirby@puppetlabs.com" <mk.kirby@gmail.com>
Date: Thu, 7 May 2020 21:38:42 -0700
Subject: [PATCH] Expire ondemand requests and remove associated instances when
 ondemand_request_ttl has been exceeded for a given request

---
 lib/vmpooler.rb                |   1 +
 lib/vmpooler/pool_manager.rb   |  45 ++++++++++-
 spec/unit/pool_manager_spec.rb | 142 ++++++++++++++++++++++++++++++---
 3 files changed, 171 insertions(+), 17 deletions(-)

diff --git a/lib/vmpooler.rb b/lib/vmpooler.rb
index 799194e..e5932d6 100644
--- a/lib/vmpooler.rb
+++ b/lib/vmpooler.rb
@@ -64,6 +64,7 @@ module Vmpooler
     parsed_config[:config]['vm_checktime'] = string_to_int(ENV['VM_CHECKTIME']) || parsed_config[:config]['vm_checktime'] || 1
     parsed_config[:config]['vm_lifetime'] = string_to_int(ENV['VM_LIFETIME']) || parsed_config[:config]['vm_lifetime'] || 24
     parsed_config[:config]['ready_ttl'] = string_to_int(ENV['READY_TTL']) || parsed_config[:config]['ready_ttl'] || 5
+    parsed_config[:config]['ondemand_request_ttl'] = string_to_int(ENV['ONDEMAND_REQUEST_TTL']) || parsed_config[:config]['ready_ttl'] || 5
     parsed_config[:config]['prefix'] = ENV['PREFIX'] || parsed_config[:config]['prefix'] || ''
 
     parsed_config[:config]['logfile'] = ENV['LOGFILE'] if ENV['LOGFILE']
diff --git a/lib/vmpooler/pool_manager.rb b/lib/vmpooler/pool_manager.rb
index 039bae7..5a0e762 100644
--- a/lib/vmpooler/pool_manager.rb
+++ b/lib/vmpooler/pool_manager.rb
@@ -139,12 +139,18 @@ module Vmpooler
       finish = format('%<time>.2f', time: Time.now - Time.parse(clone_time))
 
       if request_id
-        pool_alias = redis.hget("vmpooler__vm__#{vm}", 'pool_alias')
+        vm_hash = redis.hgetall("vmpooler__vm__#{vm}")
+        if vm_hash['status'] == 'failed'
+          move_vm_queue(pool, vm, 'pending', 'completed', redis, "moved to completed queue. '#{request_id}' could not be filled in time")
+          return nil
+        end
+
         redis.pipelined do
           redis.hset("vmpooler__active__#{pool}", vm, Time.now)
           redis.hset("vmpooler__vm__#{vm}", 'checkout', Time.now)
-          redis.sadd("vmpooler__#{request_id}__#{pool_alias}__#{pool}", vm)
+          redis.sadd("vmpooler__#{request_id}__#{vm_hash['pool_alias']}__#{pool}", vm)
         end
+        puts redis.smembers("vmpooler__#{request_id}__#{pool['alias']}__#{pool}")
         move_vm_queue(pool, vm, 'pending', 'running', redis)
       else
         redis.smove('vmpooler__pending__' + pool, 'vmpooler__ready__' + pool, vm)
@@ -1473,8 +1479,9 @@ module Vmpooler
     end
 
     def check_ondemand_requests_ready(redis)
-      in_progress_requests = redis.zrange('vmpooler__provisioning__processing', 0, -1)
-      in_progress_requests&.each do |request_id|
+      in_progress_requests = redis.zrange('vmpooler__provisioning__processing', 0, -1, with_scores: true)
+      in_progress_requests&.each do |request_id, score|
+        next if request_expired?(request_id, score, redis)
         next unless vms_ready?(request_id, redis)
 
         redis.multi
@@ -1485,6 +1492,36 @@ module Vmpooler
       in_progress_requests.length
     end
 
+    def request_expired?(request_id, score, redis)
+      delta = Time.now.to_i - score.to_i
+      ondemand_request_ttl = $config[:config]['ondemand_request_ttl']
+      return false unless (delta / 60) > ondemand_request_ttl
+      $logger.log('s', "Ondemand request for '#{request_id}' failed to provision all instances within the configured ttl '#{ondemand_request_ttl}'")
+      expiration_ttl = $config[:redis]['data_ttl'].to_i * 60 * 60
+      redis.pipelined do
+        redis.zrem('vmpooler__provisioning__processing', request_id)
+        redis.hset("vmpooler__odrequest__#{request_id}", 'status', 'failed')
+        redis.expire("vmpooler__odrequest__#{request_id}", expiration_ttl)
+      end
+      remove_vms_for_failed_request(request_id, expiration_ttl, redis)
+      true
+    end
+
+    def remove_vms_for_failed_request(request_id, expiration_ttl, redis)
+      request_hash = redis.hgetall("vmpooler__odrequest__#{request_id}")
+      requested_platforms = request_hash['requested'].split(',')
+      requested_platforms.each do |platform|
+        platform_alias, pool, _count = platform.split(':')
+        pools_filled = redis.smembers("vmpooler__#{request_id}__#{platform_alias}__#{pool}")
+        redis.pipelined do
+          pools_filled&.each do |vm|
+            move_vm_queue(pool, vm, 'running', 'completed', redis, "moved to completed queue. '#{request_id}' could not be filled in time")
+          end
+          redis.expire("vmpooler__#{request_id}__#{platform_alias}__#{pool}", expiration_ttl)
+        end
+      end
+    end
+
     def execute!(maxloop = 0, loop_delay = 1)
       $logger.log('d', 'starting vmpooler')
 
diff --git a/spec/unit/pool_manager_spec.rb b/spec/unit/pool_manager_spec.rb
index c555b17..b7d6cc3 100644
--- a/spec/unit/pool_manager_spec.rb
+++ b/spec/unit/pool_manager_spec.rb
@@ -36,6 +36,7 @@ describe 'Pool Manager' do
 :config: {}
 :providers:
   :mock:
+:redis: {}
 :pools:
   - name: '#{pool}'
     size: 1
@@ -312,23 +313,47 @@ EOT
     end
 
     context 'with request_id' do
-      it 'sets the vm as active' do
-        redis_connection_pool.with do |redis|
-          expect(Time).to receive(:now).and_return(current_time).at_least(:once)
-          redis.hset("vmpooler__vm__#{vm}", 'pool_alias', pool)
-          subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
-          expect(redis.hget("vmpooler__active__#{pool}", vm)).to eq(current_time.to_s)
-          expect(redis.hget("vmpooler__vm__#{vm}", 'checkout')).to eq(current_time.to_s)
-          expect(redis.sismember("vmpooler__#{request_id}__#{pool}__#{pool}", vm)).to be true
+      context 'with a pending request' do
+        it 'sets the vm as active' do
+          redis_connection_pool.with do |redis|
+            expect(Time).to receive(:now).and_return(current_time).at_least(:once)
+            redis.hset("vmpooler__vm__#{vm}", 'pool_alias', pool)
+            subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
+            expect(redis.hget("vmpooler__active__#{pool}", vm)).to eq(current_time.to_s)
+            expect(redis.hget("vmpooler__vm__#{vm}", 'checkout')).to eq(current_time.to_s)
+            expect(redis.sismember("vmpooler__#{request_id}__#{pool}__#{pool}", vm)).to be true
+          end
+        end
+
+        it 'logs that the vm is ready for the request' do
+          redis_connection_pool.with do |redis|
+            redis.hset("vmpooler__vm__#{vm}", 'pool_alias', pool)
+            expect(logger).to receive(:log).with('s', "[>] [#{pool}] '#{vm}' is 'ready' for request '#{request_id}'")
+
+            subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
+          end
         end
       end
 
-      it 'logs that the vm is ready for the request' do
-        redis_connection_pool.with do |redis|
-          redis.hset("vmpooler__vm__#{vm}", 'pool_alias', pool)
-          expect(logger).to receive(:log).with('s', "[>] [#{pool}] '#{vm}' is 'ready' for request '#{request_id}'")
+      context 'when the request has been marked as failed' do
+        before(:each) do
+          redis_connection_pool.with do |redis|
+            redis.hset("vmpooler__vm__#{vm}", 'status', 'failed')
+          end
+        end
 
-          subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
+        it 'moves the vm to completed' do
+          redis_connection_pool.with do |redis|
+            expect(subject).to receive(:move_vm_queue).with(pool, vm, 'pending', 'completed', redis, "moved to completed queue. '#{request_id}' could not be filled in time")
+            subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
+          end
+        end
+
+        it 'returns nil' do
+          redis_connection_pool.with do |redis|
+            result = subject.move_pending_vm_to_ready(vm, pool, redis, request_id)
+            expect(result).to be nil
+          end
         end
       end
     end
@@ -4836,6 +4861,11 @@ EOT
   end
 
   describe '#check_ondemand_requests_ready' do
+
+    before(:each) do
+      config[:config]['ondemand_request_ttl'] = 5
+    end
+
     it 'returns 0 when no provisoning requests are in progress' do
       redis_connection_pool.with do |redis|
         result = subject.check_ondemand_requests_ready(redis)
@@ -4877,6 +4907,92 @@ EOT
           end
         end
       end
+
+      context 'when a request has taken too long to be filled' do
+        it 'should return true for request_expired?' do
+          redis_connection_pool.with do |redis|
+            expect(subject).to receive(:request_expired?).with(request_id, Float, redis).and_return(true)
+            subject.check_ondemand_requests_ready(redis)
+          end
+        end
+      end
+    end
+  end
+
+  describe '#request_expired?' do
+    let(:ondemand_request_ttl) { 5 }
+    before(:each) do
+      config[:config]['ondemand_request_ttl'] = ondemand_request_ttl
+    end
+
+    context 'with a request that has taken too long to be filled' do
+      let(:expired_time) { (Time.now - 960).to_i }
+      before(:each) do
+        redis_connection_pool.with do |redis|
+          expect(subject).to receive(:remove_vms_for_failed_request)
+          create_ondemand_processing(request_id, expired_time, redis)
+        end
+      end
+
+      it 'returns true when the request is expired' do
+        redis_connection_pool.with do |redis|
+          result = subject.request_expired?(request_id, expired_time, redis)
+          expect(result).to be true
+        end
+      end
+
+      it 'logs a message that the request has expired' do
+        redis_connection_pool.with do |redis|
+          expect(logger).to receive(:log).with('s', "Ondemand request for '#{request_id}' failed to provision all instances within the configured ttl '#{ondemand_request_ttl}'")
+          subject.request_expired?(request_id, expired_time, redis)
+        end
+      end
+    end
+
+    context 'with a request that has been made within the ttl' do
+      before(:each) do
+        redis_connection_pool.with do |redis|
+          create_ondemand_processing(request_id, current_time, redis)
+        end
+      end
+
+      it 'should return false' do
+        redis_connection_pool.with do |redis|
+          result = subject.request_expired?(request_id, current_time, redis)
+          expect(result).to be false
+        end
+      end
+    end
+  end
+
+  describe '#remove_vms_for_failed_request)' do
+    let(:expiration_ttl) { 100 * 60 * 60 }
+    let(:platform_alias) { pool }
+    let(:platforms_string) { "#{platform_alias}:#{pool}:3" }
+    context 'with two vms marked as ready for the request' do
+      before(:each) do
+        redis_connection_pool.with do |redis|
+          create_ondemand_request_for_test(request_id, current_time, platforms_string, redis)
+          [vm,"#{vm}2"].each do |v|
+            create_running_vm(pool, v, redis)
+            redis.sadd("vmpooler__#{request_id}__#{platform_alias}__#{pool}", v)
+          end
+        end
+      end
+
+      it 'should remove the associated vms' do
+        redis_connection_pool.with do |redis|
+          expect(subject).to receive(:move_vm_queue).with(pool, String, 'running', 'completed', redis, "moved to completed queue. '#{request_id}' could not be filled in time").twice
+          subject.remove_vms_for_failed_request(request_id, expiration_ttl, redis)
+        end
+      end
+
+      it 'should mark the ready set for expiration' do
+        redis_connection_pool.with do |redis|
+          expect(redis).to receive(:expire).with("vmpooler__#{request_id}__#{platform_alias}__#{pool}", expiration_ttl)
+          subject.remove_vms_for_failed_request(request_id, expiration_ttl, redis)
+        end
+      end
     end
   end