Add lifecycle hooks for Executor

.rubocop.yml

@@ -21,4 +21,4 @@ Naming/FileName:
 
 Style/BlockComments:
   Exclude:
-    - 'spec/spec_helper.rb'
+    - 'spec/spec_helper.rb'

CHANGELOG.md

@@ -1,6 +1,8 @@
 Unreleased Changes
 ------------------
 
+* Feature - Add lifecycle hooks for Executor.
+
 0.1.0 (2024-11-16)
 ------------------
 

lib/aws-activejob-sqs.rb

@@ -29,6 +29,10 @@ def self.configure
       def self.fifo?(queue_url)
         queue_url.end_with?('.fifo')
       end
+
+      def self.on_worker_stop(...)
+        Executor.on_stop(...)
+      end
     end
   end
 end

lib/aws/active_job/sqs/executor.rb

@@ -15,6 +15,20 @@ class Executor
           fallback_policy: :abort # Concurrent::RejectedExecutionError must be handled
         }.freeze
 
+        class << self
+          def on_stop(&block)
+            lifecycle_hooks[:stop] << block
+          end
+
+          def lifecycle_hooks
+            @lifecycle_hooks ||= Hash.new { |h, k| h[k] = [] }
+          end
+
+          def clear_hooks
+            @lifecycle_hooks = nil
+          end
+        end
+
         def initialize(options = {})
           @executor = Concurrent::ThreadPoolExecutor.new(DEFAULTS.merge(options))
           @retry_standard_errors = options[:retry_standard_errors]
@@ -32,6 +46,7 @@ def execute(message)
         end
 
         def shutdown(timeout = nil)
+          run_hooks_for(:stop)
           @executor.shutdown
           clean_shutdown = @executor.wait_for_termination(timeout)
           if clean_shutdown
@@ -71,6 +86,12 @@ def post_task(message)
             @task_complete.set
           end
         end
+
+        def run_hooks_for(event_name)
+          return unless (hooks = self.class.lifecycle_hooks[event_name])
+
+          hooks.each(&:call)
+        end
       end
     end
   end

spec/aws/active_job/sqs/executor_spec.rb

@@ -82,6 +82,29 @@ module SQS
             expect(tp).to receive(:wait_for_termination).with(5).and_return true
             executor.shutdown(5)
           end
+
+          describe 'when lifecycle hooks are registered' do
+            let(:hook) { double }
+
+            before do
+              allow(hook).to receive(:call)
+            end
+
+            after do
+              Executor.clear_hooks
+            end
+
+            it 'executs hook when shutdown' do
+              Aws::ActiveJob::SQS.on_worker_stop do
+                hook.call
+              end
+              executor = Executor.new
+
+              executor.shutdown
+
+              expect(hook).to have_received(:call)
+            end
+          end
         end
       end
     end

spec/spec_helper.rb

@@ -76,6 +76,12 @@ class TestJobWithDedupKeys < TestJob
   # triggering implicit auto-inclusion in groups with matching metadata.
   config.shared_context_metadata_behavior = :apply_to_host_groups
 
+  # Run specs in random order to surface order dependencies. If you find an
+  # order dependency and want to debug it, you can fix the order by providing
+  # the seed, which is printed after each run.
+  #     --seed 1234
+  config.order = :random
+
 # The settings below are suggested to provide a good initial experience
 # with RSpec, but feel free to customize to your heart's content.
 =begin
@@ -115,12 +121,6 @@ class TestJobWithDedupKeys < TestJob
   # particularly slow.
   config.profile_examples = 10
 
-  # Run specs in random order to surface order dependencies. If you find an
-  # order dependency and want to debug it, you can fix the order by providing
-  # the seed, which is printed after each run.
-  #     --seed 1234
-  config.order = :random
-
   # Seed global randomization in this process using the `--seed` CLI option.
   # Setting this allows you to use `--seed` to deterministically reproduce
   # test failures related to randomization by passing the same `--seed` value