ruby: Replay worker history with existing retry queue

ruby/lib/ci/queue/configuration.rb

@@ -15,6 +15,7 @@ class Configuration
       attr_accessor :timing_redis_url
       attr_accessor :write_duration_averages
       attr_accessor :heartbeat_grace_period, :heartbeat_interval
+      attr_accessor :retry_mode
       attr_reader :circuit_breakers
       attr_writer :seed, :build_id
       attr_writer :queue_init_timeout, :report_timeout, :inactive_workers_timeout
@@ -66,7 +67,8 @@ def initialize(
         branch: nil,
         timing_redis_url: nil,
         heartbeat_grace_period: 30,
-        heartbeat_interval: 10
+        heartbeat_interval: 10,
+        retry_mode: :failures
       )
         @build_id = build_id
         @circuit_breakers = [CircuitBreaker::Disabled]
@@ -105,6 +107,7 @@ def initialize(
         @write_duration_averages = false
         @heartbeat_grace_period = heartbeat_grace_period
         @heartbeat_interval = heartbeat_interval
+        @retry_mode = retry_mode
       end
 
       def queue_init_timeout

ruby/lib/ci/queue/redis.rb

@@ -17,7 +17,9 @@ module CI
   module Queue
     module Redis
       Error = Class.new(StandardError)
+      IncompleteRetry = Class.new(Error)
       LostMaster = Class.new(Error)
+      WorkerHistoryError = Class.new(Error)
 
       class << self
 

ruby/lib/ci/queue/redis/retry.rb

@@ -12,6 +12,14 @@ def build
           @build ||= CI::Queue::Redis::BuildRecord.new(self, redis, config)
         end
 
+        def poll
+          super
+          return unless config.retry_mode == :worker_history
+          return if exhausted?
+
+          raise IncompleteRetry, 'Worker history replay stopped before completion'
+        end
+
         private
 
         attr_reader :redis

ruby/lib/ci/queue/redis/worker.rb

@@ -146,6 +146,8 @@ def retrying?
         end
 
         def retry_queue
+          return worker_history_retry_queue if config.retry_mode == :worker_history
+
           failures = build.failed_tests.to_set
           log = redis.lrange(key('worker', worker_id, 'queue'), 0, -1)
           log.select! { |id| failures.include?(id) }
@@ -264,6 +266,25 @@ def heartbeat(test_or_id = nil)
 
         attr_reader :index
 
+        def worker_history_retry_queue
+          reservations = redis.lrange(key('worker', worker_id, 'queue'), 0, -1)
+          if reservations.empty?
+            raise WorkerHistoryError, "Reservation history is missing for worker #{worker_id}"
+          end
+
+          seen = {}
+          test_ids = reservations.reverse_each.each_with_object([]) do |reservation_id, ids|
+            expand_reservation(reservation_id).each do |test_id|
+              next if seen[test_id]
+
+              seen[test_id] = true
+              ids << test_id
+            end
+          end
+
+          Retry.new(test_ids, config, redis: redis)
+        end
+
         # Runs a block while sending periodic heartbeats in a background thread.
         # This prevents other workers from stealing the test while it's being executed.
         def with_heartbeat(test_id)
@@ -525,6 +546,20 @@ def chunk_id?(id)
           id.include?(':chunk_')
         end
 
+        def expand_reservation(id)
+          return [id] unless chunk_id?(id)
+
+          chunk_json = redis.get(key('chunk', id))
+          raise WorkerHistoryError, "Chunk metadata is missing for #{id}" unless chunk_json
+
+          test_ids = CI::Queue::TestChunk.from_json(id, chunk_json).test_ids
+          raise WorkerHistoryError, "Chunk metadata contains no tests for #{id}" if test_ids.empty?
+
+          test_ids
+        rescue JSON::ParserError => error
+          raise WorkerHistoryError, "Chunk metadata is invalid for #{id}: #{error.message}"
+        end
+
         def resolve_executable(id)
           # Detect chunk by ID pattern
           if chunk_id?(id)

ruby/lib/ci/queue/static.rb

@@ -22,6 +22,7 @@ def initialize(tests, config)
         @config = config
         @progress = 0
         @total = tests.size
+        @shutdown_required = false
       end
 
       def distributed?
@@ -66,11 +67,15 @@ def size
       end
 
       def poll
-        while config.circuit_breakers.none?(&:open?) && !max_test_failed? && test = @queue.shift
+        while !@shutdown_required && config.circuit_breakers.none?(&:open?) && !max_test_failed? && test = @queue.shift
           yield index.fetch(test)
         end
       end
 
+      def shutdown!
+        @shutdown_required = true
+      end
+
       def exhausted?
         @queue.empty?
       end

ruby/lib/minitest/queue/runner.rb

@@ -58,7 +58,11 @@ def run_command
           if retry_queue.exhausted?
             puts "The retry queue does not contain any failure, we'll process the main queue instead."
           else
-            puts "Retrying failed tests."
+            if queue_config.retry_mode == :worker_history
+              puts "Replaying this worker's reservation history."
+            else
+              puts "Retrying failed tests."
+            end
             self.queue = retry_queue
           end
         end
@@ -494,6 +498,14 @@ def parser
             queue_config.worker_id = worker_id
           end
 
+          help = <<~EOS
+            Retry behavior: failures (default) or worker-history.
+          EOS
+          opts.separator ""
+          opts.on('--retry-mode MODE', %w[failures worker-history], help) do |mode|
+            queue_config.retry_mode = mode.tr('-', '_').to_sym
+          end
+
           help = <<~EOS
             Defines how many time a single test can be requeued.
             Defaults to 0.