Expose worker history retry mode

Author: dpashk-figma

ruby/README.md

@@ -38,6 +38,19 @@ 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 recover artifacts lost with a retried distributed worker, replay every test
+reserved by that worker before resuming the shared queue:
+
+```bash
+minitest-queue --queue redis://example.com \
+  --retry-mode worker-history \
+  --recovery-manifest log/ci-queue-recovery.json \
+  run -Itest test/**/*_test.rb
+```
+
+This mode requires the retry to retain its worker ID and queue build ID. It
+fails when the worker's reservation history or suite chunk metadata is missing.
+
 
 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 || 0,
           )
         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/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,47 @@
+# frozen_string_literal: true
+
+require 'fileutils'
+require 'json'
+require 'minitest/reporters'
+
+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 recovery? && !queue.replay_completed?
+
+        FileUtils.mkdir_p(File.dirname(path))
+        File.write(path, JSON.pretty_generate(manifest))
+      end
+
+      private
+
+      attr_reader :config, :path, :queue
+
+      def recovery?
+        queue.is_a?(CI::Queue::WorkerHistoryRecovery)
+      end
+
+      def manifest
+        {
+          schema_version: 1,
+          worker_id: config.worker_id.to_s,
+          retry_count: config.retry_count,
+          history_items: recovery? ? queue.history_items : 0,
+          replayed_tests: recovery? ? queue.replayed_tests : 0,
+          resumed_shared_queue: recovery? && queue.resumed_shared_queue?,
+          replay_completed: !recovery? || queue.replay_completed?
+        }
+      end
+    end
+  end
+end

ruby/lib/minitest/queue/runner.rb

@@ -54,12 +54,16 @@ def run_command
             abort! "The test run is too old and can't be retried"
           end
           reset_counters
-          retry_queue = queue.retry_queue
-          if retry_queue.exhausted?
-            puts "The retry queue does not contain any failure, we'll process the main queue instead."
+          if worker_history_retry?
+            recover_worker_history
           else
-            puts "Retrying failed tests."
-            self.queue = retry_queue
+            retry_queue = queue.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."
+              self.queue = retry_queue
+            end
           end
         end
 
@@ -81,6 +85,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 +335,22 @@ def reset_counters
         queue.build.reset_stats(BuildStatusRecorder::COUNTERS)
       end
 
+      def recover_worker_history
+        unless queue.distributed? && queue.respond_to?(:worker_history)
+          abort! 'Worker history recovery requires a distributed Redis queue'
+        end
+
+        history = queue.worker_history
+        puts "Replaying #{history.test_ids.size} tests from #{history.history_items} worker reservations."
+        self.queue = CI::Queue::WorkerHistoryRecovery.new(queue, history)
+      rescue CI::Queue::Redis::WorkerHistoryError => error
+        abort! error.message
+      end
+
+      def worker_history_retry?
+        queue_config.retry_mode == :worker_history
+      end
+
       def populate_queue
         Minitest.queue.populate(Minitest.loaded_tests, random: ordering_seed, &:id)
       end
@@ -494,6 +521,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,20 @@ 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
+    end
+
+    def test_retry_defaults
+      config = Configuration.new
+
+      assert_equal :failures, config.retry_mode
+      assert_equal 0, config.retry_count
+      assert_nil config.recovery_manifest
     end
 
     def test_travis_defaults

ruby/test/integration/minitest_redis_test.rb

@@ -220,6 +220,50 @@ def test_retry_success
       assert_equal 'All tests were ran already', output
     end
 
+    def test_worker_history_retry_writes_recovery_manifest
+      Dir.mktmpdir do |directory|
+        manifest_path = File.join(directory, 'recovery.json')
+        run_worker_history_worker(manifest_path, retry_count: 0)
+
+        first_manifest = JSON.parse(File.read(manifest_path))
+        assert_equal 0, first_manifest['retry_count']
+        assert_equal 0, first_manifest['history_items']
+        assert_equal 0, first_manifest['replayed_tests']
+        refute first_manifest['resumed_shared_queue']
+        assert first_manifest['replay_completed']
+
+        out, = run_worker_history_worker(manifest_path, retry_count: 1)
+
+        assert_includes out, 'Replaying 100 tests from 100 worker reservations.'
+        retry_manifest = JSON.parse(File.read(manifest_path))
+        assert_equal 1, retry_manifest['schema_version']
+        assert_equal '1', retry_manifest['worker_id']
+        assert_equal 1, retry_manifest['retry_count']
+        assert_equal 100, retry_manifest['history_items']
+        assert_equal 100, retry_manifest['replayed_tests']
+        assert retry_manifest['resumed_shared_queue']
+        assert retry_manifest['replay_completed']
+      end
+    end
+
+    def test_worker_history_retry_fails_without_worker_reservations
+      run_worker_history_worker(nil, retry_count: 0, build_id: 'missing-history', worker_id: '1')
+
+      Dir.mktmpdir do |directory|
+        manifest_path = File.join(directory, 'recovery.json')
+        out, = run_worker_history_worker(
+          manifest_path,
+          retry_count: 1,
+          build_id: 'missing-history',
+          worker_id: '2'
+        )
+
+        refute_predicate $?, :success?
+        assert_includes out, 'Reservation history is missing for worker 2'
+        refute_path_exists manifest_path
+      end
+    end
+
     def test_retry_fails_when_test_run_is_expired
       out, err = capture_subprocess_io do
         system(
@@ -788,6 +832,28 @@ def test_utf8_tests_and_marshal
 
     private
 
+    def run_worker_history_worker(manifest_path, retry_count:, build_id: 'worker-history', worker_id: '1')
+      args = [
+        @exe, 'run',
+        '--queue', @redis_url,
+        '--seed', 'foobar',
+        '--build', build_id,
+        '--worker', worker_id,
+        '--timeout', '1',
+        '--retry-mode', 'worker-history'
+      ]
+      args.push('--recovery-manifest', manifest_path) if manifest_path
+      args.push('-Itest', 'test/passing_test.rb')
+
+      capture_subprocess_io do
+        system(
+          { 'BUILDKITE_RETRY_COUNT' => retry_count.to_s },
+          *args,
+          chdir: 'test/fixtures/'
+        )
+      end
+    end
+
     def normalize_xml(output)
       freeze_xml_timing(rewrite_paths(output))
     end

ruby/test/minitest/queue/recovery_reporter_test.rb

@@ -0,0 +1,23 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+module Minitest
+  module Queue
+    class RecoveryReporterTest < Minitest::Test
+      def test_does_not_write_manifest_for_incomplete_run
+        Dir.mktmpdir do |directory|
+          path = File.join(directory, 'recovery.json')
+          queue = Struct.new(:exhausted?).new(false)
+          config = CI::Queue::Configuration.new(worker_id: '17')
+          reporter = RecoveryReporter.new(path: path, queue: queue, config: config)
+
+          reporter.start
+          reporter.report
+
+          refute_path_exists path
+        end
+      end
+    end
+  end
+end