Lock rows for deletion & wrap async_task in transaction to preserve queue records

Author: jturel

app/services/katello/applicability/scheduler.rb

@@ -8,6 +8,10 @@ class Scheduler
       DRAIN_MUTEX = Mutex.new
       DRAIN_SLEEP_SECONDS = 3
 
+      def self.queue
+        Katello::ApplicableHostQueue
+      end
+
       def self.scheduler_task
         ::ForemanTasks::Task::DynflowTask.running.for_action(Actions::Katello::Applicability::Scheduler).first
       end
@@ -21,32 +25,38 @@ def self.bulk_generate_tasks
       end
 
       def self.drain_loop
-        until (host_ids = ::Katello::ApplicableHostQueue.pop_hosts).empty?
-          begin
-            ForemanTasks.async_task(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: host_ids)
-          rescue => e
-            Rails.logger.error("Error while draining applicability queue #{e}")
-          end
+        catch(:done) do
+          loop do
+            ids = queue.pop_hosts do |host_ids|
+              throw(:done) if host_ids.empty?
 
-          sleep DRAIN_SLEEP_SECONDS unless host_ids.length == ::Katello::ApplicableHostQueue.batch_size # allow some time for the queue to fill
+              ForemanTasks.async_task(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: host_ids)
+            end
+            sleep DRAIN_SLEEP_SECONDS unless ids.length == queue.batch_size # allow some time for the queue to fill
+          end
         end
       end
 
       def self.trigger_drain
         return if DRAIN_MUTEX.locked?
 
-        # synchronize instead of try_lock -> consistency is more important than performance
         DRAIN_MUTEX.synchronize do
-          depth = Katello::ApplicableHostQueue.queue_depth
+          depth = queue.queue_depth
           return if depth == 0
           return if scheduler_task
 
-          if depth < Katello::ApplicableHostQueue.batch_size && bulk_generate_tasks.empty?
-            host_ids = Katello::ApplicableHostQueue.pop_hosts
-            ForemanTasks.async_task(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: host_ids)
+          if depth < queue.batch_size && bulk_generate_tasks.empty?
+            queue.pop_hosts do |host_ids|
+              ForemanTasks.async_task(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: host_ids)
+            end
           else
             # High applicability activity detected
-            trigger_scheduler_task
+            begin
+              trigger_scheduler_task
+            rescue RuntimeError => e
+              # Scheduler already started
+              Rails.logger.error "[applicability] #{e}"
+            end
           end
         end
       end

app/services/katello/applicable_host_queue.rb

@@ -15,12 +15,16 @@ def self.push_hosts(ids)
       ActiveSupport::Notifications.instrument("applicability_push_hosts")
     end
 
-    DELETE_QUERY = "DELETE FROM #{Katello::HostQueueElement.table_name} WHERE id IN (%s) RETURNING host_id".freeze
-
     def self.pop_hosts(amount = self.batch_size)
-      query = HostQueueElement.order(:id).select(:id).limit(amount)
-      result = ActiveRecord::Base.connection.execute(format(DELETE_QUERY, query.to_sql))
-      result.values.flatten
+      HostQueueElement.transaction do
+        elements = HostQueueElement.order(:id).select(:id, :host_id).limit(amount).lock
+
+        host_ids = elements.map(&:host_id)
+        yield(host_ids) if block_given?
+
+        elements.delete_all
+        host_ids
+      end
     end
   end
 end

test/services/katello/applicability/scheduler_test.rb

@@ -3,64 +3,111 @@
 module Katello
   module Applicability
     class SchedulerTest < ActiveSupport::TestCase
-      def queue
-        Katello::ApplicableHostQueue
+      let(:scheduler) { Katello::Applicability::Scheduler }
+      let(:queue) { @queue }
+
+      def setup_queue(queue)
+        @queue = queue
+        scheduler.stubs(:queue).returns(queue)
+      end
+
+      class TestQueue
+        def pop_hosts
+          @ids = @queue.pop
+          yield(@ids)
+          @ids
+        end
+
+        def queue_depth
+          @queue.flatten.length
+        end
+      end
+
+      class EmptyQueue < TestQueue
+        def initialize
+          @queue = []
+          super
+        end
+      end
+
+      class SingleBatchQueue < TestQueue
+        def initialize
+          @queue = [[], [1]]
+          super
+        end
+
+        def batch_size
+          1
+        end
+      end
+
+      class LowVolumeQueue < TestQueue
+        def initialize
+          @queue = [[], [1]]
+          super
+        end
+
+        def batch_size
+          5
+        end
       end
 
       class DrainLoopTest < SchedulerTest
         test "does nothing when queue is empty" do
-          queue.expects(:pop_hosts).returns([])
+          setup_queue EmptyQueue.new
+
           ForemanTasks.expects(:async_task).never
+          scheduler.expects(:sleep).never
 
-          Katello::Applicability::Scheduler.drain_loop
+          scheduler.drain_loop
         end
 
         test "spawns BulkGenerate" do
-          queue.stubs(:pop_hosts).returns([1]).then.returns([])
+          setup_queue SingleBatchQueue.new
+
           ForemanTasks.expects(:async_task).with(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: [1])
-          Katello::Applicability::Scheduler.expects(:sleep)
+          scheduler.expects(:sleep).never
 
-          Katello::Applicability::Scheduler.drain_loop
+          scheduler.drain_loop
         end
 
-        test "doesn't sleep when queue is full" do
-          queue.stubs(:pop_hosts).returns([1, 2, 3]).then.returns([])
-          queue.stubs(:batch_size).returns(3)
-          ForemanTasks.expects(:async_task)
-          Katello::Applicability::Scheduler.expects(:sleep).never
+        test "sleeps when queue is filling" do
+          setup_queue LowVolumeQueue.new
+
+          ForemanTasks.expects(:async_task).with(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: [1])
+          scheduler.expects(:sleep)
 
-          Katello::Applicability::Scheduler.drain_loop
+          scheduler.drain_loop
         end
       end
 
       class TriggerDrainTest < SchedulerTest
         test "spawns BulkGenerate" do
-          queue.expects(:batch_size).returns(2)
-          queue.expects(:queue_depth).returns(1)
-          Katello::Applicability::Scheduler.expects(:bulk_generate_tasks).returns([])
-          queue.expects(:pop_hosts).returns([:fake])
-          ForemanTasks.expects(:async_task).with(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: [:fake])
+          setup_queue LowVolumeQueue.new
+          scheduler.expects(:bulk_generate_tasks).returns([])
+          ForemanTasks.expects(:async_task).with(Actions::Katello::Applicability::Hosts::BulkGenerate, host_ids: [1])
 
           Katello::Applicability::Scheduler.trigger_drain
         end
 
         test "spawns scheduler task" do
-          queue.expects(:batch_size).returns(2)
-          queue.expects(:queue_depth).returns(3)
+          setup_queue LowVolumeQueue.new
+          scheduler.expects(:bulk_generate_tasks).returns([:fake])
+
           ForemanTasks.expects(:async_task).with(Actions::Katello::Applicability::Scheduler)
 
           Katello::Applicability::Scheduler.trigger_drain
         end
 
         test "does nothing when queue is empty" do
-          queue.expects(:queue_depth).returns(0)
+          setup_queue EmptyQueue.new
           ForemanTasks.expects(:async_task).never
 
           Katello::Applicability::Scheduler.trigger_drain
         end
 
         test "spawns nothing when scheduler task is running" do
-          queue.expects(:queue_depth).returns(1)
+          setup_queue LowVolumeQueue.new
           Katello::Applicability::Scheduler.expects(:scheduler_task).returns(stub)
           ForemanTasks.expects(:async_task).never
 

test/services/katello/applicable_host_queue_test.rb

@@ -4,30 +4,30 @@ module Katello
   class ApplicableHostQueueTest < ActiveSupport::TestCase
     def setup
       Setting['applicability_batch_size'] = 50
+
+      # Prevent any test here from triggering the scheduler & draining the queue
+      Katello::Applicability::Scheduler.stubs(:trigger_drain)
     end
 
     def test_pop_nothing
       assert_empty ApplicableHostQueue.pop_hosts
     end
 
     def test_pop_1_host
-      Katello::Applicability::Scheduler.expects(:trigger_drain)
       ApplicableHostQueue.push_hosts([999])
       popped_hosts = ApplicableHostQueue.pop_hosts
 
       assert_equal [999], popped_hosts.sort
     end
 
     def test_pop_5_hosts
-      Katello::Applicability::Scheduler.expects(:trigger_drain).times(5)
       5.times { |i| ApplicableHostQueue.push_hosts([i]) }
       popped_hosts = ApplicableHostQueue.pop_hosts
 
       assert_equal [0, 1, 2, 3, 4], popped_hosts.sort
     end
 
     def test_pop_batch_size_only
-      Katello::Applicability::Scheduler.expects(:trigger_drain).times(5)
       Setting['applicability_batch_size'] = 3
       5.times { |i| ApplicableHostQueue.push_hosts([i]) }
       popped_hosts = ApplicableHostQueue.pop_hosts
@@ -36,16 +36,26 @@ def test_pop_batch_size_only
     end
 
     def test_pop_duplicate_hosts
-      Katello::Applicability::Scheduler.expects(:trigger_drain).times(5)
       5.times { |i| ApplicableHostQueue.push_hosts([i]) }
       5.times { |i| ApplicableHostQueue.push_hosts([i]) }
       popped_hosts = ApplicableHostQueue.pop_hosts
 
       assert_equal [0, 1, 2, 3, 4], popped_hosts.sort
     end
 
+    def test_pop_hosts_yield_error
+      5.times { |i| ApplicableHostQueue.push_hosts([i]) }
+
+      assert_raises(RuntimeError) do
+        ApplicableHostQueue.pop_hosts do
+          fail "Elements are not to be removed from the queue!"
+        end
+      end
+
+      assert_equal 5, HostQueueElement.count
+    end
+
     def test_queue_depth
-      Katello::Applicability::Scheduler.expects(:trigger_drain).times(3)
       3.times { |i| ApplicableHostQueue.push_hosts([i]) }
 
       assert_equal 3, ApplicableHostQueue.queue_depth