ruby: Replay worker history with retry queue

Author: yfang-figma

ruby/README.md

@@ -38,6 +38,23 @@ minitest-queue --queue redis://example.com run -Itest test/**/*_test.rb
 
 Additionally you can configure the requeue settings (see main README) with `--max-requeues` and `--requeue-tolerance`.
 
+#### Replaying a retried worker for coverage
+
+By default, a retried worker runs only tests recorded as failed. For process-local
+artifacts such as SimpleCov output, use worker-history mode to replay every test
+previously reserved by that worker:
+
+```bash
+minitest-queue --queue redis://example.com run \
+  --retry-mode worker-history \
+  --recovery-manifest tmp/ci-queue-recovery.json \
+  -Itest test/**/*_test.rb
+```
+
+The replay uses the existing local retry queue and exits after the worker history
+is exhausted. It does not claim additional work from the shared Redis queue. The
+manifest is written atomically only after a complete run.
+
 
 If you'd like to centralize the error reporting you can do so with:
 

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, :recovery_manifest, :retry_count
       attr_reader :circuit_breakers
       attr_writer :seed, :build_id
       attr_writer :queue_init_timeout, :report_timeout, :inactive_workers_timeout
@@ -30,6 +31,7 @@ def from_env(env)
             redis_ttl: env['CI_QUEUE_REDIS_TTL']&.to_i ||  8 * 60 * 60,
             known_flaky_tests: load_known_flaky_tests(env['CI_QUEUE_KNOWN_FLAKY_TESTS']),
             branch: env['BUILDKITE_BRANCH'],
+            retry_count: env['BUILDKITE_RETRY_COUNT'].to_i,
           )
         end
 
@@ -66,7 +68,10 @@ def initialize(
         branch: nil,
         timing_redis_url: nil,
         heartbeat_grace_period: 30,
-        heartbeat_interval: 10
+        heartbeat_interval: 10,
+        retry_mode: :failures,
+        recovery_manifest: nil,
+        retry_count: 0
       )
         @build_id = build_id
         @circuit_breakers = [CircuitBreaker::Disabled]
@@ -105,6 +110,9 @@ def initialize(
         @write_duration_averages = false
         @heartbeat_grace_period = heartbeat_grace_period
         @heartbeat_interval = heartbeat_interval
+        @retry_mode = retry_mode
+        @recovery_manifest = recovery_manifest
+        @retry_count = retry_count
       end
 
       def queue_init_timeout

ruby/lib/ci/queue/redis.rb

@@ -18,6 +18,7 @@ module Queue
     module Redis
       Error = Class.new(StandardError)
       LostMaster = Class.new(Error)
+      WorkerHistoryError = Class.new(Error)
 
       class << self
 

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

@@ -3,11 +3,30 @@ module CI
   module Queue
     module Redis
       class Retry < Static
-        def initialize(tests, config, redis:)
+        attr_reader :history_items, :replayed_tests
+
+        def initialize(tests, config, redis:, history_items: 0, worker_history: false)
           @redis = redis
+          @history_items = history_items
+          @replayed_tests = worker_history ? tests.size : 0
+          @worker_history = worker_history
+          @worker_history_complete = false
           super(tests, config)
         end
 
+        def worker_history?
+          @worker_history
+        end
+
+        def worker_history_complete?
+          @worker_history_complete
+        end
+
+        def poll(&block)
+          super
+          @worker_history_complete = exhausted? if worker_history?
+        end
+
         def build
           @build ||= CI::Queue::Redis::BuildRecord.new(self, redis, config)
         end

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

@@ -145,7 +145,10 @@ def retrying?
           end
         end
 
-        def retry_queue
+        def retry_queue(scope: :failures)
+          return worker_history_retry_queue if scope == :worker_history
+          raise ArgumentError, "Unknown retry scope: #{scope}" unless scope == :failures
+
           failures = build.failed_tests.to_set
           log = redis.lrange(key('worker', worker_id, 'queue'), 0, -1)
           log.select! { |id| failures.include?(id) }
@@ -264,6 +267,31 @@ 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,
+            history_items: reservations.size,
+            worker_history: true
+          )
+        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 +553,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.rb

@@ -15,6 +15,7 @@
 require 'minitest/queue/grind_reporter'
 require 'minitest/queue/test_time_recorder'
 require 'minitest/queue/test_time_reporter'
+require 'minitest/queue/recovery_reporter'
 
 module Minitest
   class Requeue < Skip

ruby/lib/minitest/queue/recovery_reporter.rb

@@ -0,0 +1,56 @@
+# frozen_string_literal: true
+
+require 'fileutils'
+require 'json'
+require 'minitest/reporters'
+require 'tempfile'
+
+module Minitest
+  module Queue
+    class RecoveryReporter < Minitest::Reporters::BaseReporter
+      def initialize(path:, queue:, config:)
+        super({})
+        @path = path
+        @queue = queue
+        @config = config
+      end
+
+      def report
+        super
+        return unless queue.exhausted?
+        return if worker_history_retry? && !queue.worker_history_complete?
+
+        directory = File.dirname(path)
+        FileUtils.mkdir_p(directory)
+        temporary = Tempfile.new(['recovery', '.json'], directory)
+        temporary.write(JSON.pretty_generate(manifest))
+        temporary.flush
+        temporary.fsync
+        temporary.close
+        File.rename(temporary.path, path)
+      ensure
+        temporary&.close!
+      end
+
+      private
+
+      attr_reader :config, :path, :queue
+
+      def worker_history_retry?
+        queue.respond_to?(:worker_history?) && queue.worker_history?
+      end
+
+      def manifest
+        {
+          schema_version: 1,
+          worker_id: config.worker_id.to_s,
+          retry_count: config.retry_count,
+          history_items: worker_history_retry? ? queue.history_items : 0,
+          replayed_tests: worker_history_retry? ? queue.replayed_tests : 0,
+          resumed_shared_queue: false,
+          replay_completed: !worker_history_retry? || queue.worker_history_complete?
+        }
+      end
+    end
+  end
+end

ruby/lib/minitest/queue/runner.rb

@@ -4,6 +4,7 @@
 require 'minitest/queue'
 require 'ci/queue'
 require 'digest/md5'
+require 'fileutils'
 require 'minitest/reporters/bisect_reporter'
 require 'minitest/reporters/statsd_reporter'
 
@@ -49,16 +50,21 @@ def retry_command
 
       def run_command
         require_worker_id!
+        FileUtils.rm_f(queue_config.recovery_manifest) if queue_config.recovery_manifest
         if queue.retrying? || retry?
           if queue.expired?
             abort! "The test run is too old and can't be retried"
           end
           reset_counters
-          retry_queue = queue.retry_queue
+          retry_queue = retry_queue_for_retry
           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 worker_history_retry?
+              puts "Replaying #{retry_queue.replayed_tests} tests from #{retry_queue.history_items} worker reservations."
+            else
+              puts "Retrying failed tests."
+            end
             self.queue = retry_queue
           end
         end
@@ -81,6 +87,13 @@ def run_command
         if queue_config.statsd_endpoint
           reporters << Minitest::Reporters::StatsdReporter.new(statsd_endpoint: queue_config.statsd_endpoint)
         end
+        if queue_config.recovery_manifest
+          reporters << RecoveryReporter.new(
+            path: queue_config.recovery_manifest,
+            queue: queue,
+            config: queue_config
+          )
+        end
         Minitest.queue_reporters = reporters
 
         trap('TERM') { Minitest.queue.shutdown! }
@@ -324,6 +337,22 @@ def reset_counters
         queue.build.reset_stats(BuildStatusRecorder::COUNTERS)
       end
 
+      def worker_history_retry?
+        queue_config.retry_mode == :worker_history
+      end
+
+      def retry_queue_for_retry
+        if worker_history_retry?
+          abort! 'Worker history recovery requires a distributed Redis queue' unless queue.distributed?
+
+          queue.retry_queue(scope: :worker_history)
+        else
+          queue.retry_queue
+        end
+      rescue CI::Queue::Redis::WorkerHistoryError => error
+        abort! error.message
+      end
+
       def populate_queue
         Minitest.queue.populate(Minitest.loaded_tests, random: ordering_seed, &:id)
       end
@@ -494,6 +523,22 @@ 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
+            Write worker-history recovery details to a JSON file after a completed run.
+          EOS
+          opts.separator ""
+          opts.on('--recovery-manifest PATH', help) do |path|
+            queue_config.recovery_manifest = path
+          end
+
           help = <<~EOS
             Defines how many time a single test can be requeued.
             Defaults to 0.

ruby/test/ci/queue/configuration_test.rb

@@ -33,10 +33,13 @@ def test_buildkite_defaults
         'BUILDKITE_BUILD_ID' => '9e08ef3c-d6e6-4a86-91dd-577ce5205b8e',
         'BUILDKITE_PARALLEL_JOB' => '12',
         'BUILDKITE_COMMIT' => 'faa647bbb8168a77cf338e7488c3f8445c3e6554',
+        'BUILDKITE_RETRY_COUNT' => '2',
       )
       assert_equal '9e08ef3c-d6e6-4a86-91dd-577ce5205b8e', config.build_id
       assert_equal '12', config.worker_id
       assert_equal 'faa647bbb8168a77cf338e7488c3f8445c3e6554', config.seed
+      assert_equal 2, config.retry_count
+      assert_equal :failures, config.retry_mode
     end
 
     def test_travis_defaults