Implement a way to run the async supervisor as a separate process

rosa · rosa · commit 9b49ce07fe78 · 2024-08-31T13:15:39.000+02:00
Still running all its supervised processes as threads, but now can be
run separately from Puma.
diff --git a/lib/puma/plugin/solid_queue.rb b/lib/puma/plugin/solid_queue.rb
@@ -40,7 +40,7 @@ def start_forked(launcher)
     end
 
     def start_async(launcher)
-      launcher.events.on_booted { @solid_queue_supervisor = SolidQueue::Supervisor.start(mode: :async) }
+      launcher.events.on_booted { @solid_queue_supervisor = SolidQueue::Supervisor.start(mode: :async, sidecar: true) }
       launcher.events.on_stopped { solid_queue_supervisor.stop }
       launcher.events.on_restart { solid_queue_supervisor.stop; solid_queue_supervisor.start }
     end
diff --git a/lib/solid_queue/supervisor.rb b/lib/solid_queue/supervisor.rb
@@ -2,23 +2,23 @@
 
 module SolidQueue
   class Supervisor < Processes::Base
-    include Maintenance
+    include Maintenance, Signals, Pidfiled
 
     class << self
-      def start(mode: :fork, load_configuration_from: nil)
+      def start(mode: "fork", load_configuration_from: nil, **options)
         SolidQueue.supervisor = true
         configuration = Configuration.new(mode: mode, load_from: load_configuration_from)
 
         if configuration.configured_processes.any?
-          klass = mode == :fork ? ForkSupervisor : AsyncSupervisor
-          klass.new(configuration).tap(&:start)
+          klass = mode.to_s.inquiry.fork? ? ForkSupervisor : AsyncSupervisor
+          klass.new(configuration, **options).tap(&:start)
         else
           abort "No workers or processed configured. Exiting..."
         end
       end
     end
 
-    def initialize(configuration)
+    def initialize(configuration, **options)
       @configuration = configuration
       super
     end
diff --git a/lib/solid_queue/supervisor/async_supervisor.rb b/lib/solid_queue/supervisor/async_supervisor.rb
@@ -2,8 +2,12 @@
 
 module SolidQueue
   class Supervisor::AsyncSupervisor < Supervisor
-    def initialize(*)
+    skip_callback :boot, :before, :register_signal_handlers, if: :sidecar?
+
+    def initialize(*, sidecar: false)
       super
+
+      @sidecar = sidecar
       @threads = Concurrent::Map.new
     end
 
@@ -22,6 +26,10 @@ def stop
     private
       attr_reader :threads
 
+      def sidecar?
+        @sidecar
+      end
+
       def start_process(configured_process)
         process_instance = configured_process.instantiate.tap do |instance|
           instance.supervised_by process
@@ -32,6 +40,19 @@ def start_process(configured_process)
         threads[process_instance.name] = process_instance
       end
 
+      def supervise
+        unless sidecar?
+          loop do
+            break if stopped?
+
+            procline "supervising #{threads.keys.join(", ")}"
+            process_signal_queue
+
+            interruptible_sleep(10.second) unless stopped?
+          end
+        end
+      end
+
       def stop_threads
         stop_threads = threads.values.map do |thr|
           Thread.new { thr.stop }
@@ -40,6 +61,21 @@ def stop_threads
         stop_threads.each { |thr| thr.join(SolidQueue.shutdown_timeout) }
       end
 
+      def terminate_gracefully
+        SolidQueue.instrument(:graceful_termination, process_id: process_id, supervisor_pid: ::Process.pid, supervised_processes: threads.keys) do |payload|
+          unless all_threads_terminated?
+            payload[:shutdown_timeout_exceeded] = true
+            terminate_immediately
+          end
+        end
+      end
+
+      def terminate_immediately
+        SolidQueue.instrument(:immediate_termination, process_id: process_id, supervisor_pid: ::Process.pid, supervised_processes: threads.keys) do
+          exit!
+        end
+      end
+
       def all_threads_terminated?
         threads.values.none?(&:alive?)
       end
diff --git a/lib/solid_queue/supervisor/fork_supervisor.rb b/lib/solid_queue/supervisor/fork_supervisor.rb
@@ -2,8 +2,6 @@
 
 module SolidQueue
   class Supervisor::ForkSupervisor < Supervisor
-    include Signals, Pidfiled
-
     def initialize(*)
       super
 
diff --git a/test/dummy/app/jobs/store_result_job.rb b/test/dummy/app/jobs/store_result_job.rb
@@ -1,12 +1,12 @@
 class StoreResultJob < ApplicationJob
   queue_as :background
 
-  def perform(value, status: :completed, pause: nil, exception: nil, exit: nil)
+  def perform(value, status: :completed, pause: nil, exception: nil, exit_value: nil)
     result = JobResult.create!(queue_name: queue_name, status: "started", value: value)
 
     sleep(pause) if pause
     raise exception.new if exception
-    exit! if exit
+    exit!(exit_value) if exit_value
 
     result.update!(status: status)
   end
diff --git a/test/integration/async_processes_lifecycle_test.rb b/test/integration/async_processes_lifecycle_test.rb
@@ -0,0 +1,250 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class AsyncProcessesLifecycleTest < ActiveSupport::TestCase
+  self.use_transactional_tests = false
+
+  setup do
+    config_as_hash = { workers: [ { queues: :background }, { queues: :default, threads: 5 } ], dispatchers: [] }
+    @pid = run_supervisor_as_fork(load_configuration_from: config_as_hash, mode: :async)
+
+    wait_for_registered_processes(3, timeout: 3.second)
+    assert_registered_workers_for(:background, :default, supervisor_pid: @pid)
+  end
+
+  teardown do
+    terminate_process(@pid) if process_exists?(@pid)
+  end
+
+  test "enqueue jobs in multiple queues" do
+    6.times { |i| enqueue_store_result_job("job_#{i}") }
+    6.times { |i| enqueue_store_result_job("job_#{i}", :default) }
+
+    wait_for_jobs_to_finish_for(2.seconds)
+
+    assert_equal 12, JobResult.count
+    6.times { |i| assert_completed_job_results("job_#{i}", :background) }
+    6.times { |i| assert_completed_job_results("job_#{i}", :default) }
+
+    terminate_process(@pid)
+    assert_clean_termination
+  end
+
+  test "kill supervisor while there are jobs in-flight" do
+    no_pause = enqueue_store_result_job("no pause")
+    pause = enqueue_store_result_job("pause", pause: 0.2.seconds)
+
+    signal_process(@pid, :KILL, wait: 0.15.seconds)
+    wait_for_jobs_to_finish_for(2.seconds)
+    wait_for_registered_processes(1, timeout: 3.second)
+
+    assert_not process_exists?(@pid)
+
+    assert_completed_job_results("no pause")
+    assert_job_status(no_pause, :finished)
+
+    # Nothing had the chance to finish orderly
+    assert_registered_supervisor
+    assert_registered_workers_for(:background, :default, supervisor_pid: @pid)
+    assert_started_job_result("pause")
+    assert_claimed_jobs
+  end
+
+  test "term supervisor multiple times" do
+    5.times do
+      signal_process(@pid, :TERM, wait: 0.1.second)
+    end
+
+    sleep(1.second)
+    assert_clean_termination
+  end
+
+  test "quit supervisor while there are jobs in-flight" do
+    no_pause = enqueue_store_result_job("no pause")
+    pause = enqueue_store_result_job("pause", pause: 1.second)
+
+    signal_process(@pid, :QUIT, wait: 0.4.second)
+    wait_for_jobs_to_finish_for(2.seconds)
+
+    wait_while_with_timeout(2.seconds) { process_exists?(@pid) }
+    assert_not process_exists?(@pid)
+
+    assert_no_unfinished_jobs
+    assert_clean_termination
+  end
+
+  test "term supervisor while there are jobs in-flight" do
+    no_pause = enqueue_store_result_job("no pause")
+    pause = enqueue_store_result_job("pause", pause: 0.2.seconds)
+
+    signal_process(@pid, :TERM, wait: 0.3.second)
+    wait_for_jobs_to_finish_for(3.seconds)
+
+    assert_completed_job_results("no pause")
+    assert_completed_job_results("pause")
+
+    assert_job_status(no_pause, :finished)
+    assert_job_status(pause, :finished)
+
+    wait_for_process_termination_with_timeout(@pid, timeout: 1.second)
+    assert_clean_termination
+  end
+
+  test "int supervisor while there are jobs in-flight" do
+    no_pause = enqueue_store_result_job("no pause")
+    pause = enqueue_store_result_job("pause", pause: 0.2.seconds)
+
+    signal_process(@pid, :INT, wait: 0.3.second)
+    wait_for_jobs_to_finish_for(2.second)
+
+    assert_completed_job_results("no pause")
+    assert_completed_job_results("pause")
+
+    assert_job_status(no_pause, :finished)
+    assert_job_status(pause, :finished)
+
+    wait_for_process_termination_with_timeout(@pid, timeout: 1.second)
+    assert_clean_termination
+  end
+
+  test "term supervisor exceeding timeout while there are jobs in-flight" do
+    no_pause = enqueue_store_result_job("no pause")
+    pause = enqueue_store_result_job("pause", pause: SolidQueue.shutdown_timeout + 10.second)
+
+    signal_process(@pid, :TERM, wait: 0.5.second)
+
+    sleep(SolidQueue.shutdown_timeout + 0.5.second)
+
+    assert_completed_job_results("no pause")
+    assert_job_status(no_pause, :finished)
+
+    # This job was left claimed as the worker was shutdown without
+    # a chance to terminate orderly
+    assert_started_job_result("pause")
+    assert_job_status(pause, :claimed)
+
+    # Now wait until the supervisor finishes for real, which will complete the cleanup
+    wait_for_process_termination_with_timeout(@pid, timeout: 1.second)
+    assert_clean_termination
+  end
+
+  test "process some jobs that raise errors" do
+    2.times { enqueue_store_result_job("no error", :background) }
+    2.times { enqueue_store_result_job("no error", :default) }
+    error1 = enqueue_store_result_job("error", :background, exception: RuntimeError)
+    enqueue_store_result_job("no error", :background, pause: 0.03)
+    error2 = enqueue_store_result_job("error", :background, exception: RuntimeError, pause: 0.05)
+    2.times { enqueue_store_result_job("no error", :default, pause: 0.01) }
+    error3 = enqueue_store_result_job("error", :default, exception: RuntimeError)
+
+    wait_for_jobs_to_finish_for(2.second, except: [ error1, error2, error3 ])
+
+    assert_completed_job_results("no error", :background, 3)
+    assert_completed_job_results("no error", :default, 4)
+
+    wait_while_with_timeout(1.second) { SolidQueue::FailedExecution.count < 3 }
+    [ error1, error2, error3 ].each do |job|
+      assert_job_status(job, :failed)
+    end
+
+    terminate_process(@pid)
+    assert_clean_termination
+  end
+
+  test "process a job that exits" do
+    2.times do
+      enqueue_store_result_job("no exit", :background)
+      enqueue_store_result_job("no exit", :default)
+    end
+    paused_no_exit = enqueue_store_result_job("paused no exit", :default, pause: 0.5)
+    exit_job = enqueue_store_result_job("exit", :background, exit_value: 9, pause: 0.2)
+    pause_job = enqueue_store_result_job("exit", :background, pause: 0.3)
+
+    2.times { enqueue_store_result_job("no exit", :background) }
+
+    wait_for_jobs_to_finish_for(3.seconds, except: [ exit_job, pause_job, paused_no_exit ])
+
+    assert_completed_job_results("no exit", :default, 2)
+    assert_completed_job_results("no exit", :background, 4)
+
+    # Everything exited because of the exit job, leaving all jobs that ran
+    # after it claimed
+    [ exit_job, pause_job, paused_no_exit ].each do |job|
+      assert_job_status(job, :claimed)
+    end
+
+    wait_for_process_termination_with_timeout(@pid, exitstatus: 9)
+    assert_not process_exists?(@pid)
+
+    # Starting a supervisor again will clean things up
+    # Claimed jobs will be marked as failed and processes will be pruned
+    # Simulate time passing to expire heartbeats
+    SolidQueue::Process.update_all(last_heartbeat_at: 1.hour.ago)
+    @pid = run_supervisor_as_fork(mode: :async)
+    sleep(10)
+
+    [ exit_job, pause_job, paused_no_exit ].each do |job|
+      assert_job_status(job, :failed)
+    end
+
+    terminate_process(@pid)
+    assert_clean_termination
+  end
+
+
+  private
+    def assert_clean_termination
+      wait_for_registered_processes 0, timeout: 0.2.second
+      assert_no_registered_processes
+      assert_no_claimed_jobs
+      assert_not process_exists?(@pid)
+    end
+
+    def assert_registered_workers_for(*queues, supervisor_pid: nil)
+      workers = find_processes_registered_as("Worker")
+      registered_queues = workers.map { |process| process.metadata["queues"] }.compact
+      assert_equal queues.map(&:to_s).sort, registered_queues.sort
+      if supervisor_pid
+        assert_equal [ supervisor_pid ], workers.map { |process| process.supervisor.pid }.uniq
+      end
+    end
+
+    def assert_registered_supervisor
+      processes = find_processes_registered_as("Supervisor(async)")
+      assert_equal 1, processes.count
+      assert_equal @pid, processes.first.pid
+    end
+
+    def assert_no_registered_workers
+      assert_empty find_processes_registered_as("Worker").to_a
+    end
+
+    def enqueue_store_result_job(value, queue_name = :background, **options)
+      StoreResultJob.set(queue: queue_name).perform_later(value, **options)
+    end
+
+    def assert_completed_job_results(value, queue_name = :background, count = 1)
+      skip_active_record_query_cache do
+        assert_equal count, JobResult.where(queue_name: queue_name, status: "completed", value: value).count
+      end
+    end
+
+    def assert_started_job_result(value, queue_name = :background, count = 1)
+      skip_active_record_query_cache do
+        assert_equal count, JobResult.where(queue_name: queue_name, status: "started", value: value).count
+      end
+    end
+
+    def assert_job_status(active_job, status)
+      # Make sure we skip AR query cache. Otherwise the queries done here
+      # might be cached and since we haven't done any non-SELECT queries
+      # after they were cached on the connection used in the test, the cache
+      # will still apply, even though the data returned by the cached queries
+      # might have been deleted in the forked processes.
+      skip_active_record_query_cache do
+        job = SolidQueue::Job.find_by(active_job_id: active_job.job_id)
+        assert job.public_send("#{status}?")
+      end
+    end
+end
diff --git a/test/integration/forked_processes_lifecycle_test.rb b/test/integration/forked_processes_lifecycle_test.rb
@@ -170,7 +170,7 @@ class ForkedProcessesLifecycleTest < ActiveSupport::TestCase
       enqueue_store_result_job("no exit", :default)
     end
     enqueue_store_result_job("paused no exit", :default, pause: 0.5)
-    exit_job = enqueue_store_result_job("exit", :background, exit: true, pause: 0.2)
+    exit_job = enqueue_store_result_job("exit", :background, exit_value: 1, pause: 0.2)
     pause_job = enqueue_store_result_job("exit", :background, pause: 0.3)
 
     2.times { enqueue_store_result_job("no exit", :background) }
diff --git a/test/test_helpers/jobs_test_helper.rb b/test/test_helpers/jobs_test_helper.rb
@@ -28,4 +28,10 @@ def assert_no_claimed_jobs
       assert SolidQueue::ClaimedExecution.none?
     end
   end
+
+  def assert_claimed_jobs(count = 1)
+    skip_active_record_query_cache do
+      assert_equal count, SolidQueue::ClaimedExecution.count
+    end
+  end
 end
diff --git a/test/test_helpers/processes_test_helper.rb b/test/test_helpers/processes_test_helper.rb
diff --git a/test/unit/async_supervisor_test.rb b/test/unit/async_supervisor_test.rb
