ruby_reactor 0.5.1 → 0.5.2

data/lib/ruby_reactor/executor.rb

@@ -7,10 +7,13 @@ require_relative "executor/retry_manager"
 require_relative "executor/compensation_manager"
 require_relative "executor/result_handler"
 require_relative "executor/step_executor"
+require_relative "executor/ordered_lock_support"
 
 module RubyReactor
   # rubocop:disable Metrics/ClassLength
   class Executor
+    include OrderedLockSupport
+
     attr_reader :reactor_class, :context, :dependency_graph, :compensation_manager, :retry_manager, :result_handler,
                 :step_executor, :result, :middlewares
 
@@ -41,6 +44,8 @@ module RubyReactor
       @result = nil
       @acquired_lock = nil
       @acquired_semaphore = nil
+      @contention_snooze = false
+      @skip_context_persist = false
     end
 
     def self.resolve_middlewares(reactor_class)
@@ -71,9 +76,13 @@ module RubyReactor
       middlewares.on(:start_reactor, reactor_class.name, context.inputs, @context)
       completed = false
 
-      if (skipped = check_period_gate)
+      enter_ordered_lock_scope
+      # short_circuit_result covers both the strict ordered-lock chain skip
+      # and the already-marked period bucket.
+      short = short_circuit_result
+      if short
         completed = true
-        return finalize_skipped(skipped)
+        return short_circuit!(short)
       end
 
       # Validate inputs BEFORE consuming a rate-limit slot or grabbing a
@@ -109,7 +118,9 @@ module RubyReactor
     rescue RubyReactor::Lock::AcquisitionError,
            RubyReactor::Semaphore::AcquisitionError,
            RubyReactor::RateLimit::ExceededError,
-           RubyReactor::RateLimitRegistry::UnknownLimitError => e
+           RubyReactor::RateLimitRegistry::UnknownLimitError,
+           RubyReactor::OrderedLock::WaitError => e
+      @contention_snooze = true
       raise e
     rescue StandardError => e
       @result = @result_handler.handle_execution_error(e)
@@ -118,77 +129,98 @@ module RubyReactor
       @result
     ensure
       release_locks
-      save_context if persist_context?
+      leave_ordered_lock_scope
+      save_context if persist_context? && !skip_context_persist?
 
+      emit_lifecycle_completion(completed)
+    end
+
+    # Contention errors (lock/semaphore/rate-limit/ordered-lock wait) are
+    # expected "try again later" signals, not failures — the worker snoozes
+    # and re-runs. Emitting `failed_reactor` for them floods dashboards with
+    # phantom failures (one per snooze round), so route them to a distinct
+    # `snooze_reactor` event instead.
+    def emit_lifecycle_completion(completed)
       if completed
         middlewares.on(:complete_reactor, reactor_class.name, @result, @context)
+      elsif @contention_snooze
+        middlewares.on(:snooze_reactor, reactor_class.name, $ERROR_INFO, @context)
       else
         middlewares.on(:failed_reactor, reactor_class.name, $ERROR_INFO, @context)
       end
     end
 
-    def resume_execution # rubocop:disable Metrics/MethodLength
+    def resume_execution # rubocop:disable Metrics/MethodLength,Metrics/PerceivedComplexity
       middlewares.on(:start_reactor, reactor_class.name, context.inputs, @context)
       completed = false
+
       # A fresh async reactor run reaches the worker through resume_execution
       # (it never calls execute), so the period and rate-limit gates that live
       # in execute must be applied here too. Genuine resumes (a step already ran
       # or we paused mid-flight, so current_step is set) must NOT re-gate: a
       # paused reactor must not throttle or skip itself on the way back in.
       first_run = first_execution?
-      begin
-        @context.status = :running
 
-        if first_run && (skipped = check_period_gate)
-          completed = true
-          return finalize_skipped(skipped)
-        end
-        check_rate_limit if first_run
+      enter_ordered_lock_scope
+      # ordered-lock skip applies on any run; the period gate only on a fresh
+      # first run (a genuine resume must not skip itself when its own marker
+      # eventually lands).
+      short = ordered_lock_short_circuit
+      short ||= check_period_gate if first_run
+      if short
+        completed = true
+        return short_circuit!(short)
+      end
 
-        acquire_concurrency_primitives
+      @context.status = :running
+      check_rate_limit if first_run
 
-        # Post-lock re-check (see execute) — closes the period race for the
-        # first run of a locked async reactor.
-        if first_run && (skipped = check_period_gate)
-          completed = true
-          return finalize_skipped(skipped)
-        end
+      # Resumes intentionally skip check_rate_limit (a paused run must not
+      # block itself on resume), so acquire lock/semaphore directly rather
+      # than via acquire_locks.
+      acquire_exclusive_lock if @reactor_class.respond_to?(:lock_config) && @reactor_class.lock_config
+      acquire_semaphore if @reactor_class.respond_to?(:semaphore_config) && @reactor_class.semaphore_config
 
-        prepare_for_resume
-        save_context
+      # Post-lock re-check (see execute) — closes the period race for the
+      # first run of a locked async reactor.
+      if first_run && (skipped = check_period_gate)
+        completed = true
+        return finalize_skipped(skipped)
+      end
 
-        @result = if @context.current_step
-                    execute_current_step_and_continue
-                  else
-                    execute_remaining_steps
-                  end
+      prepare_for_resume
+      save_context
 
-        update_context_status(@result)
-        mark_period_on_success(@result)
+      @result = if @context.current_step
+                  execute_current_step_and_continue
+                else
+                  execute_remaining_steps
+                end
 
-        handle_interrupt(@result) if @result.is_a?(RubyReactor::InterruptResult)
-        completed = true
-        @result
-      rescue RubyReactor::Lock::AcquisitionError,
-             RubyReactor::Semaphore::AcquisitionError,
-             RubyReactor::RateLimit::ExceededError,
-             RubyReactor::RateLimitRegistry::UnknownLimitError
-        raise
-      rescue StandardError => e
-        handle_resume_error(e)
-        update_context_status(@result)
-        completed = true
-        @result
-      ensure
-        release_locks
-        save_context
+      update_context_status(@result)
+      mark_period_on_success(@result)
 
-        if completed
-          middlewares.on(:complete_reactor, reactor_class.name, @result, @context)
-        else
-          middlewares.on(:failed_reactor, reactor_class.name, $ERROR_INFO, @context)
-        end
-      end
+      handle_interrupt(@result) if @result.is_a?(RubyReactor::InterruptResult)
+      completed = true
+      @result
+    rescue RubyReactor::Lock::AcquisitionError,
+           RubyReactor::Semaphore::AcquisitionError,
+           RubyReactor::RateLimit::ExceededError,
+           RubyReactor::RateLimitRegistry::UnknownLimitError,
+           RubyReactor::OrderedLock::WaitError => e
+      @contention_snooze = true
+      raise e
+    rescue StandardError => e
+      handle_resume_error(e)
+      update_context_status(@result)
+      completed = true
+      @result
+    ensure
+      release_locks
+      leave_ordered_lock_scope
+      save_context unless skip_context_persist?
+
+      emit_lifecycle_completion(completed)
     end
 
     def undo_all

data/lib/ruby_reactor/ordered_lock.rb

@@ -0,0 +1,158 @@
+# frozen_string_literal: true
+
+module RubyReactor
+  # Strict-ordering primitive. A monotonically increasing nonce is assigned at
+  # enqueue time; the worker can proceed only when its nonce equals
+  # `last_completed + 1`. Otherwise the worker raises {WaitError}, which the
+  # Sidekiq worker rescues and re-snoozes via `perform_in`.
+  #
+  # See `with_ordered_lock` for usage from a reactor.
+  class OrderedLock
+    # Raised by the gate check when the worker's nonce is ahead of
+    # `last_completed + 1`. Carries `retry_after_seconds`, a hint derived from
+    # the poison-pill timeout on the *blocker* nonce.
+    class WaitError < StandardError
+      attr_reader :retry_after_seconds, :key, :nonce, :last_completed
+
+      def initialize(key:, nonce:, last_completed:, retry_after_seconds:)
+        @key = key
+        @nonce = nonce
+        @last_completed = last_completed
+        @retry_after_seconds = retry_after_seconds
+        super("OrderedLock '#{key}' nonce #{nonce} waiting on #{last_completed + 1}")
+      end
+    end
+
+    # Default poison-pill: if the blocker nonce was assigned more than
+    # `poison_pill_timeout` seconds ago and never advanced, the gate treats it
+    # as dead and advances past it. Prevents permanent head-of-line blocking
+    # from a crashed caller that INCRed but never enqueued.
+    DEFAULT_POISON_PILL_TIMEOUT = 600
+
+    # TTL on the Redis counter keys. Bumped on every assign so an active
+    # sequence never expires; only fully-drained ones GC themselves.
+    DEFAULT_TTL = 86_400
+
+    attr_reader :key, :nonce, :epoch, :poison_pill_timeout, :strict
+
+    def initialize(key, nonce: nil, epoch: nil, poison_pill_timeout: DEFAULT_POISON_PILL_TIMEOUT, # rubocop:disable Metrics/ParameterLists
+                   ttl: DEFAULT_TTL, strict: true)
+      @key = key
+      @nonce = nonce
+      @epoch = epoch
+      @poison_pill_timeout = poison_pill_timeout
+      @ttl = ttl
+      @strict = strict
+    end
+
+    # Atomic INCR on the `next` counter. Caller-side; runs during
+    # `Reactor.run` BEFORE `perform_async`. Returns `[nonce, epoch]` — the nonce
+    # we own plus the generation it belongs to (used to fence stale stragglers).
+    def self.assign(key, ttl: DEFAULT_TTL)
+      adapter = RubyReactor.configuration.storage_adapter
+      adapter.ordered_lock_assign(key, ttl: ttl)
+    end
+
+    # Gate check. Returns `:go`, `:drained_go`, `:skip_chain_failed`,
+    # `:stale_batch`, or raises {WaitError}.
+    # - `:go` — proceed to run steps.
+    # - `:drained_go` — the batch fully drained and GC'd while this caller slept.
+    #   A genuine late straggler should run; a Sidekiq redelivery of an
+    #   already-terminal context should be skipped. The executor disambiguates
+    #   via the stored context status.
+    # - `:skip_chain_failed` — only in strict mode: an earlier nonce in this
+    #   sequence terminated with a Failure, so this run is short-circuited
+    #   with `Skipped(reason: :ordered_lock_chain_failed)` without executing.
+    # - `:stale_batch` — this run's epoch no longer matches the key's current
+    #   generation: its batch fully drained and the numbering was reused by a
+    #   newer batch. The run is short-circuited with
+    #   `Skipped(reason: :ordered_lock_stale_batch)` and must not participate.
+    # - `:poison_advance` is collapsed to `:go` from the caller's perspective.
+    def check!
+      raise ArgumentError, "OrderedLock#check! requires a nonce" unless @nonce
+
+      state, retry_after, last_completed, first_failed = adapter.ordered_lock_can_proceed(
+        @key,
+        nonce: @nonce,
+        poison_pill_timeout: @poison_pill_timeout,
+        epoch: @epoch.to_i
+      )
+
+      case state
+      when "go", "poison_advance"
+        chain_failed?(first_failed) ? :skip_chain_failed : :go
+      when "drained_go"
+        # Batch fully drained and GC'd while this caller slept. A genuine late
+        # straggler may run (poison semantics); a redelivery of an
+        # already-terminal context must not. The executor disambiguates via the
+        # stored context status. (fail_key is GC'd here, so no chain check.)
+        :drained_go
+      when "stale"
+        :stale_batch
+      when "wait"
+        raise WaitError.new(
+          key: @key,
+          nonce: @nonce,
+          last_completed: last_completed,
+          retry_after_seconds: retry_after
+        )
+      else
+        raise "Unexpected OrderedLock state: #{state.inspect}"
+      end
+    end
+
+    # Move `last_completed` forward. Idempotent: only the nonce equal to
+    # `last_completed + 1` advances; others are no-ops (the poison-pill path
+    # may have already skipped us).
+    #
+    # Call on terminal status only (success, permanent failure, escalated skip).
+    # Retryable failures must NOT advance — the same nonce keeps owning until
+    # the job either succeeds or exhausts its retry budget.
+    #
+    # `failed:` records this nonce as the chain-failure marker (only the FIRST
+    # failure sticks). In strict mode the marker causes subsequent nonces to
+    # short-circuit with Skipped.
+    def advance!(failed: false)
+      raise ArgumentError, "OrderedLock#advance! requires a nonce" unless @nonce
+
+      adapter.ordered_lock_advance(@key, nonce: @nonce, failed: failed, epoch: @epoch.to_i, ttl: @ttl)
+    end
+
+    # Restamp this nonce's `assigned_at` to "now" while its steps execute, so a
+    # successor does not poison-advance past a blocker that is merely slow (not
+    # dead). Called on an interval by a background heartbeat thread for the
+    # duration of step execution. No-op if the nonce's timer was already deleted
+    # by a terminal advance, or if the batch has gone stale (epoch fence).
+    def heartbeat!
+      return unless @nonce
+
+      adapter.ordered_lock_heartbeat(@key, nonce: @nonce, epoch: @epoch.to_i)
+    end
+
+    # Read-only inspection. `{ next:, last_completed:, in_flight: [...] }`.
+    def self.peek(key)
+      RubyReactor.configuration.storage_adapter.ordered_lock_peek(key)
+    end
+
+    # Manual ops escape hatch — force-advance past a stuck nonce.
+    def self.skip!(key, nonce:)
+      RubyReactor.configuration.storage_adapter.ordered_lock_skip(key, nonce: nonce)
+    end
+
+    # Nuke all counters for a key. Ops only; concurrent enqueues during reset
+    # produce undefined ordering.
+    def self.reset!(key)
+      RubyReactor.configuration.storage_adapter.ordered_lock_reset(key)
+    end
+
+    private
+
+    def chain_failed?(first_failed)
+      @strict && first_failed.to_i.positive? && @nonce > first_failed.to_i
+    end
+
+    def adapter
+      RubyReactor.configuration.storage_adapter
+    end
+  end
+end

data/lib/ruby_reactor/reactor.rb

@@ -102,6 +102,11 @@ module RubyReactor
         return validation_result
       end
 
+      # Assign-at-enqueue: ordered_lock nonce is INCRed atomically here so
+      # the order matches the caller's order, not whichever worker happens to
+      # pick the job up first.
+      assign_ordered_lock_nonce!
+
       if self.class.async? && !@context.inline_async_execution
         # For async reactors, queue a job for the whole reactor
         @context.status = :running
@@ -423,6 +428,42 @@ module RubyReactor
       serialized_context = ContextSerializer.serialize(@context)
       storage.store_context(@context.context_id, serialized_context, reactor_class_name)
     end
+
+    def assign_ordered_lock_nonce!
+      return unless self.class.respond_to?(:ordered_lock_config) && self.class.ordered_lock_config
+      return if @context.private_data[:ordered_lock] || @context.private_data["ordered_lock"]
+
+      config = self.class.ordered_lock_config
+      key = config[:key_proc].call(@context.inputs)
+
+      # Synchronous nested `Reactor.run` of an ordered-lock reactor on the same
+      # key would deadlock: the outer nonce holds the slot, an inner nonce
+      # would never advance until the outer completes — but the outer is
+      # blocked waiting for the inner to return. Mirror the compose behavior:
+      # silently skip nonce assignment (the inner runs without gate/advance)
+      # and log a warning so this isn't invisible.
+      active = Executor::OrderedLockSupport.active_keys
+      if active.include?(key)
+        RubyReactor.configuration.logger.warn(
+          "RubyReactor: nested `Reactor.run` of #{self.class.name || "<anonymous>"} on " \
+          "ordered-lock key '#{key}' from inside another ordered-lock reactor on the same " \
+          "key — nonce assignment skipped, inner run executes without ordering enforcement. " \
+          "Use a different key or move the inner call to a top-level invocation if you need ordering."
+        )
+        return
+      end
+
+      nonce, epoch = RubyReactor::OrderedLock.assign(key, ttl: config[:ttl])
+
+      @context.private_data[:ordered_lock] = {
+        key: key,
+        nonce: nonce,
+        epoch: epoch,
+        poison_pill_timeout: config[:poison_pill_timeout],
+        ttl: config[:ttl],
+        strict: config.fetch(:strict, true)
+      }
+    end
   end
   # rubocop:enable Metrics/ClassLength
 end

data/lib/ruby_reactor/rspec/helpers.rb

@@ -2,7 +2,13 @@
 
 module RubyReactor
   module RSpec
+    # Globally-included helpers. Only methods whose names clearly belong to
+    # RubyReactor's test surface live here (`test_reactor`). Sidekiq-coupled
+    # helpers live in `SidekiqHelpers` and are scoped to `type: :reactor`.
     module Helpers
+      # Build a `TestSubject` around a reactor invocation. Captures the run for
+      # later introspection via matchers; runs the reactor lazily on first
+      # query unless `.run` is called explicitly.
       def test_reactor(reactor_class, inputs, context: {}, async: nil, process_jobs: true)
         TestSubject.new(
           reactor_class: reactor_class,

data/lib/ruby_reactor/rspec/matchers.rb

@@ -415,6 +415,72 @@ module RubyReactor
         end
       end
 
+      # Asserts the last-assigned ordered_lock nonce for a key. Subject is
+      # the user-provided ordered_lock key (without the `ordered_lock:` prefix).
+      #
+      #   expect("orders:42").to have_ordered_lock_next(3)
+      ::RSpec::Matchers.define :have_ordered_lock_next do |expected|
+        match { |key| Matchers.coordination_adapter.ordered_lock_peek(key)[:next] == expected }
+
+        failure_message do |key|
+          state = Matchers.coordination_adapter.ordered_lock_peek(key)
+          "expected ordered_lock '#{key}' next to be #{expected}, got #{state[:next]} " \
+            "(last_completed: #{state[:last_completed]}, in_flight: #{state[:in_flight].inspect})"
+        end
+      end
+
+      # Asserts the last-advanced cursor for an ordered_lock key.
+      #
+      #   expect("orders:42").to have_ordered_lock_last_completed(2)
+      ::RSpec::Matchers.define :have_ordered_lock_last_completed do |expected|
+        match do |key|
+          Matchers.coordination_adapter.ordered_lock_peek(key)[:last_completed] == expected
+        end
+
+        failure_message do |key|
+          state = Matchers.coordination_adapter.ordered_lock_peek(key)
+          "expected ordered_lock '#{key}' last_completed to be #{expected}, " \
+            "got #{state[:last_completed]} (next: #{state[:next]}, in_flight: #{state[:in_flight].inspect})"
+        end
+      end
+
+      # Asserts the exact set of in-flight nonces for an ordered_lock key.
+      # Order-insensitive — the matcher sorts both sides.
+      #
+      #   expect("orders:42").to have_ordered_lock_in_flight(2, 3)
+      ::RSpec::Matchers.define :have_ordered_lock_in_flight do |*expected|
+        match do |key|
+          actual = Matchers.coordination_adapter.ordered_lock_peek(key)[:in_flight].sort
+          actual == expected.flatten.map(&:to_i).sort
+        end
+
+        failure_message do |key|
+          state = Matchers.coordination_adapter.ordered_lock_peek(key)
+          "expected ordered_lock '#{key}' in_flight to be #{expected.flatten.sort.inspect}, " \
+            "got #{state[:in_flight].inspect} (next: #{state[:next]}, last_completed: #{state[:last_completed]})"
+        end
+      end
+
+      # Asserts an ordered_lock key has fully drained — counters GC'd, no
+      # in-flight nonces. After a clean drain `peek` returns all zeros.
+      #
+      #   expect("orders:42").to be_ordered_lock_drained
+      ::RSpec::Matchers.define :be_ordered_lock_drained do
+        match do |key|
+          state = Matchers.coordination_adapter.ordered_lock_peek(key)
+          state[:next].zero? && state[:last_completed].zero? && state[:in_flight].empty?
+        end
+
+        failure_message do |key|
+          state = Matchers.coordination_adapter.ordered_lock_peek(key)
+          "expected ordered_lock '#{key}' to be drained, but state is #{state.inspect}"
+        end
+
+        failure_message_when_negated do |key|
+          "expected ordered_lock '#{key}' not to be drained, but counters are all zero"
+        end
+      end
+
       # Add more matchers as per plan
       # rubocop:enable Metrics/BlockLength
     end

data/lib/ruby_reactor/rspec/sidekiq_helpers.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module RubyReactor
+  module RSpec
+    # Async-job manipulation helpers. Names like `drain_async_jobs` are too
+    # generic to live in the global spec namespace, so this module is only
+    # auto-included into examples tagged `type: :reactor`. Specs that need it
+    # outside that tag should `include RubyReactor::RSpec::SidekiqHelpers`
+    # explicitly.
+    module SidekiqHelpers
+      # Drain every queued async job across all RubyReactor worker classes
+      # until the queues are empty. Recursive — handles jobs that re-enqueue
+      # themselves (e.g. ordered_lock snoozes) and worker chains that queue
+      # additional jobs.
+      def drain_async_jobs(max_iterations: 100)
+        SidekiqHelpers.drain_async_jobs(max_iterations: max_iterations)
+      end
+
+      # All currently-pending async jobs, wrapped in `PendingJob` so callers
+      # can perform individual jobs out-of-order (e.g. to assert
+      # ordered_lock's snooze behavior) without touching Sidekiq internals.
+      def pending_async_jobs
+        SidekiqHelpers.pending_async_jobs
+      end
+
+      PendingJob = Struct.new(:worker_class, :raw) do
+        def perform!
+          worker_class.jobs.delete(raw)
+          worker_class.new.perform(*raw["args"])
+        end
+
+        def args
+          raw["args"]
+        end
+      end
+
+      def self.worker_classes
+        @worker_classes ||= [
+          RubyReactor::SidekiqWorkers::Worker,
+          RubyReactor::SidekiqWorkers::MapElementWorker,
+          RubyReactor::SidekiqWorkers::MapCollectorWorker
+        ]
+      end
+
+      def self.drain_async_jobs(max_iterations: 100)
+        return unless defined?(Sidekiq::Testing)
+
+        max_iterations.times do
+          processed_any = false
+          worker_classes.each do |worker_class|
+            while (job = worker_class.jobs.shift)
+              worker_class.new.perform(*job["args"])
+              processed_any = true
+            end
+          end
+
+          break unless processed_any
+        end
+      end
+
+      def self.pending_async_jobs
+        return [] unless defined?(Sidekiq::Testing)
+
+        worker_classes.flat_map do |worker_class|
+          worker_class.jobs.map { |raw| PendingJob.new(worker_class, raw) }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_reactor/rspec/storage_reset.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+module RubyReactor
+  module RSpec
+    # Test-only `reset!` impls layered onto storage adapters at framework
+    # load time. Kept out of `lib/ruby_reactor/storage/*` so production code
+    # never gains a "wipe everything" entry point.
+    module StorageReset
+      module RedisAdapterReset
+        def reset!
+          @redis.flushdb
+        end
+      end
+
+      def self.install!
+        return if @installed
+
+        ::RubyReactor::Storage::RedisAdapter.prepend(RedisAdapterReset)
+        @installed = true
+      end
+    end
+  end
+end

data/lib/ruby_reactor/rspec/test_subject.rb

@@ -246,6 +246,8 @@ module RubyReactor
                   end
                 end
           RubyReactor::Success.new(val)
+        when "skipped"
+          skipped_result(ctx)
         when "running"
           # Try to determine if it is truly running or if we just missed the completion
           if @process_jobs && defined?(Sidekiq::Testing)
@@ -269,6 +271,17 @@ module RubyReactor
         end
       end
 
+      # A clean halt: either a `with_period` gate or a step returning
+      # `RubyReactor.Skipped(...)`. The sync run already produced the exact
+      # Skipped (reason/step intact) — surface it. For async runs the worker
+      # swallows the return value, so rebuild from the trace.
+      def skipped_result(ctx)
+        return @run_result if @run_result.is_a?(RubyReactor::Skipped)
+
+        entry = ctx.execution_trace.reverse.find { |t| t[:type].to_s == "skipped" }
+        RubyReactor::Skipped.new(reason: entry&.dig(:reason), step_name: entry&.dig(:step))
+      end
+
       def success?
         ensure_executed!
         @reactor_instance.context.status.to_s == "completed"
@@ -419,34 +432,7 @@ module RubyReactor
       def process_pending_jobs
         return unless defined?(Sidekiq::Testing)
 
-        # Loop until no more jobs are being queued
-        # This handles batched map execution where jobs queue more jobs
-        max_iterations = 100
-        iterations = 0
-
-        while iterations < max_iterations
-          iterations += 1
-          jobs_processed = false
-
-          # Known worker classes to check
-          worker_classes = [
-            RubyReactor::SidekiqWorkers::Worker,
-            RubyReactor::SidekiqWorkers::MapElementWorker,
-            RubyReactor::SidekiqWorkers::MapCollectorWorker
-          ]
-
-          worker_classes.each do |worker_class|
-            while worker_class.jobs.any?
-              job = worker_class.jobs.shift
-              worker_class.new.perform(*job["args"])
-              jobs_processed = true
-            end
-          end
-
-          break unless jobs_processed
-        end
-
-        # Final reload
+        SidekiqHelpers.drain_async_jobs
         @reactor_instance = @reactor_class.find(@reactor_instance.context.context_id)
       end