Merge pull request #99 from figma/dpashkevich/worker-history-retry-selection

ruby: Replay worker history on retries

Author: dpashk-figma

ruby/README.md

@@ -38,6 +38,16 @@ 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`.
 
+To reconstruct process-local artifacts after a distributed worker is retried, replay every test represented by that worker's reservation history:
+
+```bash
+minitest-queue --queue redis://example.com \
+  --retry-selection worker-history \
+  run -Itest test/**/*_test.rb
+```
+
+Worker-history retries require the retry to retain its worker ID and queue build ID. Missing reservation history or suite chunk metadata fails the retry. The replay uses a local retry queue and does not rejoin the shared queue, so surviving workers must drain any remaining work. When centralized reporting is used, restarting every worker before the shared queue is exhausted fails the build rather than accepting an incomplete 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_selection
       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_selection: :failed_tests
       )
         @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_selection = retry_selection
       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

@@ -3,18 +3,27 @@ module CI
   module Queue
     module Redis
       class Retry < Static
-        def initialize(tests, config, redis:)
+        def initialize(tests, config, redis:, require_exhaustion: false)
           @redis = redis
+          @require_exhaustion = require_exhaustion
           super(tests, config)
         end
 
         def build
           @build ||= CI::Queue::Redis::BuildRecord.new(self, redis, config)
         end
 
+        def poll
+          super
+          return unless require_exhaustion
+          return if exhausted?
+
+          raise IncompleteRetry, 'Worker history replay stopped before completion'
+        end
+
         private
 
-        attr_reader :redis
+        attr_reader :redis, :require_exhaustion
       end
     end
   end

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

@@ -145,13 +145,23 @@ def retrying?
           end
         end
 
-        def retry_queue
-          failures = build.failed_tests.to_set
-          log = redis.lrange(key('worker', worker_id, 'queue'), 0, -1)
-          log.select! { |id| failures.include?(id) }
-          log.uniq!
-          log.reverse!
-          Retry.new(log, config, redis: redis)
+        def retry_queue(selection: :failed_tests)
+          reservations = redis.lrange(key('worker', worker_id, 'queue'), 0, -1)
+          test_ids = case selection
+                     when :failed_tests
+                       failed_test_ids(reservations)
+                     when :worker_history
+                       worker_history_test_ids(reservations)
+                     else
+                       raise ArgumentError, "Unknown retry selection: #{selection.inspect}"
+                     end
+
+          Retry.new(
+            test_ids,
+            config,
+            redis: redis,
+            require_exhaustion: selection == :worker_history
+          )
         end
 
         def supervisor
@@ -525,6 +535,31 @@ def chunk_id?(id)
           id.include?(':chunk_')
         end
 
+        def failed_test_ids(reservations)
+          failures = build.failed_tests.to_set
+          reservations.select { |id| failures.include?(id) }.uniq.reverse
+        end
+
+        def worker_history_test_ids(reservations)
+          raise WorkerHistoryError, "Reservation history is missing for worker #{worker_id}" if reservations.empty?
+
+          reservations.reverse.flat_map { |id| expand_reservation(id) }.uniq
+        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 => e
+          raise WorkerHistoryError, "Chunk metadata is invalid for #{id}: #{e.message}"
+        end
+
         def resolve_executable(id)
           # Detect chunk by ID pattern
           if chunk_id?(id)

ruby/lib/ci/queue/static.rb

@@ -22,6 +22,11 @@ def initialize(tests, config)
         @config = config
         @progress = 0
         @total = tests.size
+        @shutdown_required = false
+      end
+
+      def shutdown!
+        @shutdown_required = true
       end
 
       def distributed?
@@ -66,7 +71,7 @@ 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

ruby/lib/minitest/queue/runner.rb

@@ -54,11 +54,15 @@ def run_command
             abort! "The test run is too old and can't be retried"
           end
           reset_counters
-          retry_queue = queue.retry_queue
+          retry_queue = selected_retry_queue
           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_selection == :worker_history
+              puts "Replaying #{retry_queue.total} tests from worker 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
+            Select tests for a retried worker: failed-tests (default) or worker-history.
+          EOS
+          opts.separator ""
+          opts.on('--retry-selection SELECTION', %w[failed-tests worker-history], help) do |selection|
+            queue_config.retry_selection = selection.tr('-', '_').to_sym
+          end
+
           help = <<~EOS
             Defines how many time a single test can be requeued.
             Defaults to 0.
@@ -741,6 +753,19 @@ def retry?
         ENV["BUILDKITE_RETRY_COUNT"].to_i > 0 ||
           ENV["SEMAPHORE_PIPELINE_RERUN"] == "true"
       end
+
+      def selected_retry_queue
+        return queue.retry_queue if queue_config.retry_selection == :failed_tests
+
+        unless queue_config.retry_selection == :worker_history
+          abort! "Unknown retry selection: #{queue_config.retry_selection.inspect}"
+        end
+        abort! 'Worker-history retry selection requires a distributed Redis queue' unless queue.distributed?
+
+        queue.retry_queue(selection: :worker_history)
+      rescue CI::Queue::Redis::WorkerHistoryError => e
+        abort! e.message
+      end
     end
   end
 end

ruby/test/ci/queue/configuration_test.rb

@@ -73,6 +73,12 @@ def test_redis_ttl_defaults
       assert_equal(28_800, config.redis_ttl)
     end
 
+    def test_retry_selection_defaults_to_failed_tests
+      config = Configuration.new
+
+      assert_equal :failed_tests, config.retry_selection
+    end
+
     def test_redis_ttl_from_env
       config = Configuration.from_env(
         "CI_QUEUE_REDIS_TTL" => "14400"

ruby/test/ci/queue/redis/worker_chunk_test.rb

@@ -243,6 +243,65 @@ def test_populate_with_many_chunks_uses_batching
     end
   end
 
+  def test_worker_history_retry_expands_chunks_and_deduplicates_tests
+    tests = create_mock_tests(['TestA#test_1', 'TestA#test_2', 'TestB#test_1'])
+    chunk = CI::Queue::TestChunk.new(
+      'TestA:chunk_0',
+      'TestA',
+      ['TestA#test_1', 'TestA#test_2'],
+      2000.0
+    )
+
+    @worker.stub(:reorder_tests, [chunk, tests.last]) do
+      @worker.populate(tests)
+    end
+
+    history_key = 'build:42:worker:1:queue'
+    [chunk.id, tests.last.id, chunk.id, tests[1].id].each do |id|
+      @redis.lpush(history_key, id)
+    end
+
+    retry_queue = @worker.retry_queue(selection: :worker_history)
+    retry_queue.populate(tests)
+
+    assert_equal ['TestA#test_1', 'TestA#test_2', 'TestB#test_1'], retry_queue.to_a.map(&:id)
+  end
+
+  def test_worker_history_retry_requires_chunk_metadata
+    @redis.lpush('build:42:worker:1:queue', 'TestA:chunk_0')
+
+    error = assert_raises(CI::Queue::Redis::WorkerHistoryError) do
+      @worker.retry_queue(selection: :worker_history)
+    end
+
+    assert_equal 'Chunk metadata is missing for TestA:chunk_0', error.message
+  end
+
+  def test_worker_history_retry_rejects_empty_chunks
+    @redis.set(
+      'build:42:chunk:TestA:chunk_0',
+      CI::Queue::TestChunk.new('TestA:chunk_0', 'TestA', [], 0).to_json
+    )
+    @redis.lpush('build:42:worker:1:queue', 'TestA:chunk_0')
+
+    error = assert_raises(CI::Queue::Redis::WorkerHistoryError) do
+      @worker.retry_queue(selection: :worker_history)
+    end
+
+    assert_equal 'Chunk metadata contains no tests for TestA:chunk_0', error.message
+  end
+
+  def test_worker_history_retry_rejects_invalid_chunk_metadata
+    @redis.set('build:42:chunk:TestA:chunk_0', '{')
+    @redis.lpush('build:42:worker:1:queue', 'TestA:chunk_0')
+
+    error = assert_raises(CI::Queue::Redis::WorkerHistoryError) do
+      @worker.retry_queue(selection: :worker_history)
+    end
+
+    assert_match 'Chunk metadata is invalid for TestA:chunk_0:', error.message
+  end
+
   private
 
   def create_mock_tests(test_ids)

ruby/test/ci/queue/redis_test.rb

@@ -64,6 +64,43 @@ def test_retry_queue_with_all_tests_passing_2
     assert_equal retry_test_order, retry_test_order
   end
 
+  def test_retry_queue_with_worker_history
+    original_order = poll(@queue)
+    retry_queue = populate(@queue.retry_queue(selection: :worker_history))
+
+    assert_equal original_order, poll(retry_queue)
+  end
+
+  def test_worker_history_retry_requires_reservations
+    error = assert_raises(CI::Queue::Redis::WorkerHistoryError) do
+      @queue.retry_queue(selection: :worker_history)
+    end
+
+    assert_equal 'Reservation history is missing for worker 1', error.message
+  end
+
+  def test_worker_history_retry_requires_complete_replay
+    poll(@queue)
+    retry_queue = populate(@queue.retry_queue(selection: :worker_history))
+
+    error = assert_raises(CI::Queue::Redis::IncompleteRetry) do
+      retry_queue.poll do
+        retry_queue.shutdown!
+      end
+    end
+
+    assert_equal 'Worker history replay stopped before completion', error.message
+    refute_predicate retry_queue, :exhausted?
+  end
+
+  def test_retry_queue_rejects_unknown_selection
+    error = assert_raises(ArgumentError) do
+      @queue.retry_queue(selection: :everything)
+    end
+
+    assert_equal 'Unknown retry selection: :everything', error.message
+  end
+
   def test_shutdown
     poll(@queue) do
       @queue.shutdown!