minitest-distributed 0.1.2 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: bfef30635f4dd487b913e94e04a1098f27d27278af1f1adbb794a1f91f54e880
-  data.tar.gz: cff09c3e45440e77a7f39dabcdc6faf3a99336cc05362ea1740c810c91bb227b
+  metadata.gz: 076d9467680eff44b28d42436648c5f4878fb5689bca7ddd9365bd05991a5352
+  data.tar.gz: '0068262916e339e61d4997eb478df915be5943d3a6b39e70fb02def9ba0d0dbb'
 SHA512:
-  metadata.gz: b9a05e26c4b0a432d3b7c2d8aec342d141b0fc5b4a1b7a355b6e5e14609a85e834c60fde5d6bf32a85529d9678c0ce19e2a8e9feb1d6da49c761b77289c340b0
-  data.tar.gz: 430c1c0da35af7e4db05580676de9ae4d911ecb0673d99535b07c00a72f600dd71e3e279a127e254bba2be13352579d1867ff376edaf1b91fd241b5b6da0b408
+  metadata.gz: 5d71ead8d7f352d9d628ec6682d2367804c63362e1e02e658ae3899b55486e1ff1bedd1cc10e23796a3373a0a3d0580b874646c4f7e3ec472675b8b85d88b001
+  data.tar.gz: cafd308bdad9ee0332323e0e826810fee2be91cbb5bdafa74ec8037891e9a6031ffa6f4f1e4777527e0f21784a48502848893661120145749629240164c3c077

data/.rubocop.yml

@@ -11,6 +11,10 @@ AllCops:
   Exclude:
     - minitest-distributed.gemspec
 
+# This cop is broken when using assignments
+Layout/RescueEnsureAlignment:
+  Enabled: false
+
 ##### Sorbet cops
 
 Sorbet:

data/Gemfile

@@ -1,7 +1,7 @@
 # frozen_string_literal: true
 source "https://rubygems.org"
 
-# Specify your gem's dependencies in minitest-stateful.gemspec
+# Specify your gem's dependencies in minitest-distributed.gemspec
 gemspec
 
 gem "rake", "~> 12.0"

data/README.md

@@ -63,8 +63,8 @@ them to fail.
 
 ### Other optional command line arguments
 
-- `--test-timeout=SECONDS` or `ENV[MINITEST_TEST_TIMEOUT]` (default: 30s): the
-  maximum amount a test is allowed to run before it times out. In a distributed
+- `--test-timeout=SECONDS` or `ENV[MINITEST_TEST_TIMEOUT_SECONDS]` (default: 30s):
+  the maximum amount a test is allowed to run before it times out. In a distributed
   system, it's impossible to differentiate between a worker being slow and a
   worker being broken. When the timeout passes, the other workers will assume
   that the worker running the test has crashed, and will attempt to claim this
@@ -92,24 +92,40 @@ other tests.
 
 ## Development
 
-After checking out the repo, run `bin/setup` to install dependencies. Then,
-run `rake test` to run the tests. You can also run `bin/console` for an
-interactive prompt that will allow you to experiment.
+To bootstrap a local development environment:
 
-To install this gem onto your local machine, run `bundle exec rake install`.
-To release a new version, update the version number in `version.rb`, and then
-run `bundle exec rake release`, which will create a git tag for the version,
-push git commits and tags, and push the `.gem` file to
-[rubygems.org](https://rubygems.org).
+- Run `bin/setup` to install dependencies.
+- Start a Redis server by running `redis-server`, assuming you have Redis
+  installed locally and the binary is on your `PATH`. Alternatively, you can
+  set the `REDIS_URL` environment variable to point to a Redis instance running
+  elsewhere.
+- Now, run `bin/rake test` to run the tests, and verify everything is working.
+- You can also run `bin/console` for an interactive prompt that will allow you
+  to experiment.
+
+### Releasing a new version
+
+- To install this gem onto your local machine, run `bin/rake install`.
+- Only people at Shopify can release a new version to
+  [rubygems.org](https://rubygems.org). To do so, update the `VERSION` constant
+  in `version.rb`, and merge to master. Shipit will take care of building the
+  `.gem` bundle, and pushing it to rubygems.org.
 
 ## Contributing
 
-Bug reports and pull requests are welcome on GitHub at https://github.com/Shopify/minitest-distributed. This project is intended to be a safe, welcoming space for collaboration, and contributors are expected to adhere to the [code of conduct](https://github.com/Shopify/minitest-distributed/blob/master/CODE_OF_CONDUCT.md).
+Bug reports and pull requests are welcome on GitHub at
+https://github.com/Shopify/minitest-distributed. This project is intended to
+be a safe, welcoming space for collaboration, and contributors are expected to
+adhere to the [code of
+conduct](https://github.com/Shopify/minitest-distributed/blob/master/CODE_OF_CONDUCT.md).
 
 ## License
 
-The gem is available as open source under the terms of the [MIT License](https://opensource.org/licenses/MIT).
+The gem is available as open source under the terms of the [MIT
+License](https://opensource.org/licenses/MIT).
 
 ## Code of Conduct
 
-Everyone interacting in the Minitest::Stateful project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/Shopify/minitest-distributed/blob/master/CODE_OF_CONDUCT.md).
+Everyone interacting in the `minitest-distributed` project's codebases, issue
+trackers, chat rooms and mailing lists is expected to follow the [code of
+conduct](https://github.com/Shopify/minitest-distributed/blob/master/CODE_OF_CONDUCT.md).

data/bin/setup

@@ -4,5 +4,3 @@ IFS=$'\n\t'
 set -vx
 
 bundle install
-
-# Do any other automated setup that you need to do here

data/lib/minitest/distributed/configuration.rb

@@ -8,8 +8,8 @@ module Minitest
   module Distributed
     class Configuration < T::Struct
       DEFAULT_BATCH_SIZE = 10
-      DEFAULT_MAX_ATTEMPTS = 3
-      DEFAULT_TEST_TIMEOUT = 30.0 # seconds
+      DEFAULT_MAX_ATTEMPTS = 1
+      DEFAULT_TEST_TIMEOUT_SECONDS = 30.0 # seconds
 
       class << self
         extend T::Sig
@@ -20,11 +20,54 @@ module Minitest
             coordinator_uri: URI(env['MINITEST_COORDINATOR'] || 'memory:'),
             run_id: env['MINITEST_RUN_ID'] || SecureRandom.uuid,
             worker_id: env['MINITEST_WORKER_ID'] || SecureRandom.uuid,
-            test_timeout: Float(env['MINITEST_TEST_TIMEOUT'] || DEFAULT_TEST_TIMEOUT),
+            test_timeout_seconds: Float(env['MINITEST_TEST_TIMEOUT_SECONDS'] || DEFAULT_TEST_TIMEOUT_SECONDS),
             test_batch_size: Integer(env['MINITEST_TEST_BATCH_SIZE'] || DEFAULT_BATCH_SIZE),
             max_attempts: Integer(env['MINITEST_MAX_ATTEMPTS'] || DEFAULT_MAX_ATTEMPTS),
+            max_failures: (max_failures_env = env['MINITEST_MAX_FAILURES']) ? Integer(max_failures_env) : nil,
           )
         end
+
+        sig { params(opts: OptionParser).returns(T.attached_class) }
+        def from_command_line_options(opts)
+          configuration = from_env
+
+          opts.on('--coordinator=URI', "The URI pointing to the coordinator") do |uri|
+            configuration.coordinator_uri = URI.parse(uri)
+          end
+
+          opts.on('--test-timeout=TIMEOUT', "The maximum run time for a single test in seconds") do |timeout|
+            configuration.test_timeout_seconds = Float(timeout)
+          end
+
+          opts.on('--max-attempts=ATTEMPTS', "The maximum number of attempts to run a test") do |attempts|
+            configuration.max_attempts = Integer(attempts)
+          end
+
+          opts.on('--test-batch-size=NUMBER', "The number of tests to process per batch") do |batch_size|
+            configuration.test_batch_size = Integer(batch_size)
+          end
+
+          opts.on('--max-failures=FAILURES', "The maximum allowed failure before aborting a run") do |failures|
+            configuration.max_failures = Integer(failures)
+          end
+
+          opts.on('--run-id=ID', "The ID for this run shared between coordinated workers") do |id|
+            configuration.run_id = id
+          end
+
+          opts.on('--worker-id=ID', "The unique ID for this worker") do |id|
+            configuration.worker_id = id
+          end
+
+          opts.on(
+            '--[no-]retry-failures', "Retry failed and errored tests from a previous run attempt " \
+            "with the same run ID (default: enabled)"
+          ) do |enabled|
+            configuration.retry_failures = enabled
+          end
+
+          configuration
+        end
       end
 
       extend T::Sig
@@ -32,9 +75,11 @@ module Minitest
       prop :coordinator_uri, URI::Generic, default: URI('memory:')
       prop :run_id, String, factory: -> { SecureRandom.uuid }
       prop :worker_id, String, factory: -> { SecureRandom.uuid }
-      prop :test_timeout, Float, default: DEFAULT_TEST_TIMEOUT
+      prop :test_timeout_seconds, Float, default: DEFAULT_TEST_TIMEOUT_SECONDS
       prop :test_batch_size, Integer, default: DEFAULT_BATCH_SIZE
       prop :max_attempts, Integer, default: DEFAULT_MAX_ATTEMPTS
+      prop :max_failures, T.nilable(Integer)
+      prop :retry_failures, T::Boolean, default: true
 
       sig { returns(Coordinators::CoordinatorInterface) }
       def coordinator

data/lib/minitest/distributed/coordinators/coordinator_interface.rb

@@ -18,6 +18,9 @@ module Minitest
         sig { abstract.returns(ResultAggregate) }
         def combined_results; end
 
+        sig { abstract.returns(T::Boolean) }
+        def aborted?; end
+
         sig { abstract.params(test_selector: TestSelector).void }
         def produce(test_selector:); end
 

data/lib/minitest/distributed/coordinators/memory_coordinator.rb

@@ -25,7 +25,8 @@ module Minitest
 
           @leader = T.let(Mutex.new, Mutex)
           @queue = T.let(Queue.new, Queue)
-          @local_results = T.let(ResultAggregate.new, ResultAggregate)
+          @local_results = T.let(ResultAggregate.new(max_failures: configuration.max_failures), ResultAggregate)
+          @aborted = T.let(false, T::Boolean)
         end
 
         sig { override.params(reporter: Minitest::CompositeReporter, options: T::Hash[Symbol, T.untyped]).void }
@@ -33,6 +34,11 @@ module Minitest
           # No need for any additional reporters
         end
 
+        sig { override.returns(T::Boolean) }
+        def aborted?
+          @aborted
+        end
+
         sig { override.params(test_selector: TestSelector).void }
         def produce(test_selector:)
           if @leader.try_lock
@@ -41,24 +47,38 @@ module Minitest
             if tests.empty?
               queue.close
             else
-              tests.each { |test| queue << test }
+              tests.each do |runnable|
+                queue << EnqueuedRunnable.new(
+                  class_name: T.must(runnable.class.name),
+                  method_name: runnable.name,
+                  test_timeout_seconds: configuration.test_timeout_seconds,
+                  max_attempts: configuration.max_attempts,
+                )
+              end
             end
           end
         end
 
         sig { override.params(reporter: AbstractReporter).void }
         def consume(reporter:)
-          until queue.empty? && queue.closed?
-            enqueued_runnable = queue.pop
+          until queue.closed?
+            enqueued_runnable = T.let(queue.pop, EnqueuedRunnable)
+
             reporter.prerecord(enqueued_runnable.runnable_class, enqueued_runnable.method_name)
-            result = enqueued_runnable.run
 
-            local_results.update_with_result(result)
-            local_results.acks += 1
+            enqueued_result = enqueued_runnable.run do |initial_result|
+              if ResultType.of(initial_result) == ResultType::Requeued
+                queue << enqueued_runnable.next_attempt
+              end
+              EnqueuedRunnable::Result::Commit.success
+            end
 
-            reporter.record(result)
+            reporter.record(enqueued_result.committed_result)
+            local_results.update_with_result(enqueued_result)
 
-            queue.close if local_results.completed?
+            # We abort a run if we reach the maximum number of failures
+            queue.close if combined_results.abort?
+            queue.close if combined_results.complete?
           end
         end
       end

data/lib/minitest/distributed/coordinators/redis_coordinator.rb

@@ -15,11 +15,19 @@ module Minitest
       # to the stream.
       #
       # AFter that, all workers will start consuming from the stream. They will first
-      # try to claim stale entries from other workers (determined by the `test_timeout`
-      # option), and process them tp to a maxumim of `max_attempts` attempts. Then,
+      # try to claim stale entries from other workers (determined by the `test_timeout_seconds`
+      # option), and process them up to a maximum of `max_attempts` attempts. Then,
       # they will consume tests from the stream, run them, and ack them. This is done
       # in batches to reduce load on Redis.
       #
+      # Retrying failed tests (up to `max_attempts` times) uses the same mechanism.
+      # When a test fails, and we haven't exhausted the maximum number of attempts, we
+      # do not ACK the result with Redis. The means that another worker will eventually
+      # claim the test, and run it again. However, in this case we don't want to slow
+      # things down unnecessarily. When a test fails and we want to retry it, we add the
+      # test to the `retry_set` in Redis. When other worker sees that a test is in this
+      # set, it can immediately claim the test, rather than waiting for the timeout.
+      #
       # Finally, when we have acked the same number of tests as we populated into the
       # queue, the run is considered finished. The first worker to detect this will
       # remove the consumergroup and the associated stream from Redis.
@@ -47,7 +55,10 @@ module Minitest
         attr_reader :local_results
 
         sig { returns(T::Set[EnqueuedRunnable]) }
-        attr_reader :reclaimed_tests
+        attr_reader :reclaimed_timeout_tests
+
+        sig { returns(T::Set[EnqueuedRunnable]) }
+        attr_reader :reclaimed_failed_tests
 
         sig { params(configuration: Configuration).void }
         def initialize(configuration:)
@@ -58,7 +69,9 @@ module Minitest
           @group_name = T.let('minitest-distributed', String)
           @local_results = T.let(ResultAggregate.new, ResultAggregate)
           @combined_results = T.let(nil, T.nilable(ResultAggregate))
-          @reclaimed_tests = T.let(Set.new, T::Set[EnqueuedRunnable])
+          @reclaimed_timeout_tests = T.let(Set.new, T::Set[EnqueuedRunnable])
+          @reclaimed_failed_tests = T.let(Set.new, T::Set[EnqueuedRunnable])
+          @aborted = T.let(false, T::Boolean)
         end
 
         sig { override.params(reporter: Minitest::CompositeReporter, options: T::Hash[Symbol, T.untyped]).void }
@@ -70,41 +83,51 @@ module Minitest
         def combined_results
           @combined_results ||= begin
             stats_as_string = redis.mget(key('runs'), key('assertions'), key('passes'),
-              key('failures'), key('errors'), key('skips'), key('reruns'), key('acks'), key('size'))
+              key('failures'), key('errors'), key('skips'), key('requeues'), key('discards'),
+              key('acks'), key('size'))
 
             ResultAggregate.new(
+              max_failures: configuration.max_failures,
+
               runs: Integer(stats_as_string.fetch(0) || 0),
               assertions: Integer(stats_as_string.fetch(1) || 0),
               passes: Integer(stats_as_string.fetch(2) || 0),
               failures: Integer(stats_as_string.fetch(3) || 0),
               errors: Integer(stats_as_string.fetch(4) || 0),
               skips: Integer(stats_as_string.fetch(5) || 0),
-              reruns: Integer(stats_as_string.fetch(6) || 0),
-              acks: Integer(stats_as_string.fetch(7) || 0),
+              requeues: Integer(stats_as_string.fetch(6) || 0),
+              discards: Integer(stats_as_string.fetch(7) || 0),
+              acks: Integer(stats_as_string.fetch(8) || 0),
 
-              # In the case where we have no build szie number published yet, we initialize
+              # In the case where we have no build size number published yet, we initialize
               # thesize of the test suite to be arbitrarity large, to make sure it is
               # higher than the number of acks, so the run is not consider completed yet.
-              size: Integer(stats_as_string.fetch(8) || 2_147_483_647),
+              size: Integer(stats_as_string.fetch(9) || 2_147_483_647),
             )
           end
         end
 
+        sig { override.returns(T::Boolean) }
+        def aborted?
+          @aborted
+        end
+
         sig { override.params(test_selector: TestSelector).void }
         def produce(test_selector:)
           # Whoever ends up creating the consumer group will act as leader,
           # and publish the list of tests to the stream.
 
-          begin
+          initial_attempt = begin
             # When using `redis.multi`, the second DEL command gets executed even if the initial GROUP
             # fails. This is bad, because only the leader should be issuing the DEL command.
             # When using EVAL and a Lua script, the script aborts after the first XGROUP command
             # fails, and the DEL never gets executed for followers.
-            redis.evalsha(
+            keys_deleted = redis.evalsha(
               register_consumergroup_script,
               keys: [stream_key, key('size'), key('acks')],
               argv: [group_name],
             )
+            keys_deleted == 0
 
           rescue Redis::CommandError => ce
             if ce.message.include?('BUSYGROUP')
@@ -118,38 +141,67 @@ module Minitest
             end
           end
 
-          run_attempt, previous_failures, previous_errors, _deleted = redis.multi do
-            redis.incr(key('attempt'))
-            redis.lrange(key('failure_list'), 0, -1)
-            redis.lrange(key('error_list'), 0, -1)
-            redis.del(key('failure_list'), key('error_list'))
-          end
-
-          tests = if run_attempt == 1
+          tests = T.let([], T::Array[Minitest::Runnable])
+          tests = if initial_attempt
             # If this is the first attempt for this run ID, we will schedule the full
             # test suite as returned by the test selector to run.
-            test_selector.tests
-          else
-            # For subsequent attempts, we check the list of previous failures and
-            # errors, and only schedule to re-run those tests. This allows for faster
-            # retries of potentially flaky tests.
-            (previous_failures + previous_errors).map do |test_to_retry|
-              EnqueuedRunnable.from_hash!(Marshal.load(test_to_retry))
+
+            tests_from_selector = test_selector.tests
+            adjust_combined_results(ResultAggregate.new(size: tests_from_selector.size))
+            tests_from_selector
+
+          elsif configuration.retry_failures
+            # Before starting a retry attempt, we first check if the previous attempt
+            # was aborted before it was completed. If this is the case, we cannot use
+            # retry mode, and should immediately fail the attempt.
+            if combined_results.abort?
+              # We mark this run as aborted, which causes this worker to not be successful.
+              @aborted = true
+
+              # We still publish an empty size run to Redis, so if there are any followers,
+              # they will wind down normally. Only the leader will exit
+              # with a non-zero exit status and fail the build; any follower will
+              # exit with status 0.
+              adjust_combined_results(ResultAggregate.new(size: 0))
+              T.let([], T::Array[Minitest::Runnable])
+            else
+              previous_failures, previous_errors, _deleted = redis.multi do
+                redis.lrange(list_key(ResultType::Failed.serialize), 0, -1)
+                redis.lrange(list_key(ResultType::Error.serialize), 0, -1)
+                redis.del(list_key(ResultType::Failed.serialize), list_key(ResultType::Error.serialize))
+              end
+
+              # We set the `size` key to the number of tests we are planning to schedule.
+              # We also adjust the number of failures and errors back to 0.
+              # We set the number of requeues to the number of tests that failed, so the
+              # run statistics will reflect that we retried some failed test.
+              #
+              # However, normally requeues are not acked, as we expect the test to be acked
+              # by another worker later. This makes the test loop think iot is already done.
+              # To prevent this, we initialize the number of acks negatively, so it evens out
+              # in the statistics.
+              total_failures = previous_failures.length + previous_errors.length
+              adjust_combined_results(ResultAggregate.new(
+                size: total_failures,
+                failures: -previous_failures.length,
+                errors: -previous_errors.length,
+                requeues: total_failures,
+              ))
+
+              # For subsequent attempts, we check the list of previous failures and
+              # errors, and only schedule to re-run those tests. This allows for faster
+              # retries of potentially flaky tests.
+              test_identifiers_to_retry = T.let(previous_failures + previous_errors, T::Array[String])
+              test_identifiers_to_retry.map { |identifier| DefinedRunnable.from_identifier(identifier) }
             end
+          else
+            adjust_combined_results(ResultAggregate.new(size: 0))
+            T.let([], T::Array[Minitest::Runnable])
           end
 
-          # We set the `size` key to the number of tests we are planning to schedule.
-          # This will allow workers to tell when the run is done. We also adjust the
-          # number of failures and errors in case of a retry run.
-          adjust_combined_results(ResultAggregate.new(
-            size: tests.size,
-            failures: -previous_failures.length,
-            errors: -previous_errors.length,
-            reruns: previous_failures.length + previous_errors.length,
-          ))
-
-          # TODO: break this up in batches.
-          tests.each { |test| redis.xadd(stream_key, test.serialize) }
+          redis.pipelined do
+            tests.each { |test| redis.xadd(stream_key, class_name: T.must(test.class.name), method_name: test.name) }
+          end
         end
 
         sig { override.params(reporter: AbstractReporter).void }
@@ -158,26 +210,29 @@ module Minitest
           loop do
             # First, see if there are any pending tests from other workers to claim.
             stale_runnables = claim_stale_runnables
-            stale_processed = process_batch(stale_runnables, reporter)
+            process_batch(stale_runnables, reporter)
 
-            # Finally, try to process a regular batch of messages
+            # Then, try to process a regular batch of messages
             fresh_runnables = claim_fresh_runnables(block: exponential_backoff)
-            fresh_processed = process_batch(fresh_runnables, reporter)
+            process_batch(fresh_runnables, reporter)
 
             # If we have acked the same amount of tests as we were supposed to, the run
             # is complete and we can exit our loop. Generally, only one worker will detect
             # this condition. The pther workers will quit their consumer loop because the
             # consumergroup will be deleted by the first worker, and their Redis commands
             # will start to fail - see the rescue block below.
-            break if combined_results.completed?
+            break if combined_results.complete?
+
+            # We also abort a run if we reach the maximum number of failures
+            break if combined_results.abort?
 
             # To make sure we don't end up in a busy loop overwhelming Redis with commands
             # when there is no work to do, we increase the blocking time exponentially,
-            # and reset it to the initial value if we processed any messages
-            if stale_processed > 0 || fresh_processed > 0
-              exponential_backoff = INITIAL_BACKOFF
-            else
+            # and reset it to the initial value if we processed any tests.
+            if stale_runnables.empty? && fresh_runnables.empty?
               exponential_backoff <<= 1
+            else
+              exponential_backoff = INITIAL_BACKOFF
             end
           end
 
@@ -203,28 +258,20 @@ module Minitest
           @redis ||= Redis.new(url: configuration.coordinator_uri)
         end
 
-        sig { returns(String) }
-        def ack_batch_script
-          @ack_batch_script = T.let(@ack_batch_script, T.nilable(String))
-          @ack_batch_script ||= redis.script(:load, <<~LUA)
-            local acked_ids, acked, i = {}, 0, 2
-            while ARGV[i] do
-              if redis.call('XACK', KEYS[1], ARGV[1], ARGV[i]) > 0 then
-                acked = acked + 1
-                acked_ids[acked] = ARGV[i]
-              end
-              i = i + 1
-            end
-            return acked_ids
-          LUA
-        end
-
         sig { returns(String) }
         def register_consumergroup_script
           @register_consumergroup_script = T.let(@register_consumergroup_script, T.nilable(String))
           @register_consumergroup_script ||= redis.script(:load, <<~LUA)
+            -- Try to create the consumergroup. This will raise an error if the
+            -- consumergroup has already been registered by somebody else, which
+            -- means another worker will be acting as leader.
+            -- In that case, the next Redis DEL call will not be executed.
             redis.call('XGROUP', 'CREATE', KEYS[1], ARGV[1], '0', 'MKSTREAM')
-            redis.call('DEL', KEYS[2], KEYS[3])
+
+            -- The leader should reset the size and acks key for this run attempt.
+            -- We return the number of keys that were deleted, which can be used to
+            -- determine whether this was the first attempt for this run or not.
+            return redis.call('DEL', KEYS[2], KEYS[3])
           LUA
         end
 
@@ -232,51 +279,119 @@ module Minitest
         def claim_fresh_runnables(block:)
           result = redis.xreadgroup(group_name, configuration.worker_id, stream_key, '>',
             block: block, count: configuration.test_batch_size)
-          EnqueuedRunnable.from_redis_stream_claim(result.fetch(stream_key, []))
+          EnqueuedRunnable.from_redis_stream_claim(result.fetch(stream_key, []), configuration: configuration)
+        end
+
+        sig do
+          params(
+            pending_messages: T::Hash[String, PendingExecution],
+            max_idle_time_ms: Integer,
+          ).returns(T::Array[EnqueuedRunnable])
+        end
+        def xclaim_messages(pending_messages, max_idle_time_ms:)
+          return [] if pending_messages.empty?
+          claimed = redis.xclaim(stream_key, group_name, configuration.worker_id,
+            max_idle_time_ms, pending_messages.keys)
+
+          EnqueuedRunnable.from_redis_stream_claim(claimed, pending_messages, configuration: configuration)
         end
 
         sig { returns(T::Array[EnqueuedRunnable]) }
         def claim_stale_runnables
-          # When we have to reclaim stale tests, those test are potentially too slow
-          # to run inside the test timeout. We only claim one test at a time in order
-          # to prevent the exact same batch from being too slow on repeated attempts,
-          # which would cause us to mark all the tests in that batch as failed.
-          #
-          # This has the side effect that for a retried test, the test timeout
-          # will be TEST_TIMEOUT * BATCH_SIZE in practice. This gives us a higher
-          # likelihood that the test will pass if the batch size > 1.
-          pending = redis.xpending(stream_key, group_name, '-', '+', 1)
-
-          # Every test is allowed to take test_timeout milliseconds. Because we process tests in
-          # batches, they should never be pending for TEST_TIMEOUT * BATCH_SIZE milliseconds.
+          # Every test is allowed to take test_timeout_seconds. Because we process tests in
+          # batches, they should never be pending for TEST_TIMEOUT_SECONDS * BATCH_SIZE seconds.
           # So, only try to claim messages older than that, with a bit of jitter.
-          max_idle_time = Integer(configuration.test_timeout * configuration.test_batch_size * 1000)
-          max_idle_time_with_jitter = max_idle_time * rand(1.0...1.2)
-          to_claim = pending.each_with_object({}) do |message, hash|
-            if message['elapsed'] > max_idle_time_with_jitter
-              hash[message.fetch('entry_id')] = message
+          max_idle_time_ms = Integer(configuration.test_timeout_seconds * configuration.test_batch_size * 1000)
+          max_idle_time_ms_with_jitter = max_idle_time_ms * rand(1.0...1.2)
+
+          # Find all the pending messages to see if we want to attenpt to claim some.
+          pending = redis.xpending(stream_key, group_name, '-', '+', configuration.test_batch_size)
+          return [] if pending.empty?
+
+          active_consumers = Set[configuration.worker_id]
+
+          stale_messages = {}
+          active_messages = {}
+          pending.each do |msg|
+            message = PendingExecution.from_xpending(msg)
+            if message.elapsed_time_ms < max_idle_time_ms_with_jitter
+              active_consumers << message.worker_id
+              active_messages[message.entry_id] = message
+            else
+              stale_messages[message.entry_id] = message
             end
           end
 
-          if to_claim.empty?
-            []
-          else
-            claimed = redis.xclaim(stream_key, group_name, configuration.worker_id, max_idle_time, to_claim.keys)
-            enqueued_runnables = EnqueuedRunnable.from_redis_stream_claim(claimed)
-            enqueued_runnables.each do |er|
-              # `count` will be set to the current attempt of a different worker that has timed out.
-              # The attempt we are going to try will be the next one, so add one.
-              attempt = to_claim.fetch(er.execution_id).fetch('count') + 1
-              if attempt > configuration.max_attempts
-                # If we exhaust our attempts, we will mark the test to immediately fail when it will be run next.
-                mark_runnable_to_fail_immediately(er)
-              else
-                reclaimed_tests << er
+          # If we only have evidence of one active consumer based on the pending message,
+          # we will query Redis for all consumers to make sure we have full data.
+          # We can skip this if we already know that there is more than one active one.
+          if active_consumers.size == 1
+            begin
+              redis.xinfo('consumers', stream_key, group_name).each do |consumer|
+                if consumer.fetch('idle') < max_idle_time_ms
+                  active_consumers << consumer.fetch('name')
+                end
               end
+            rescue Redis::CommandError
+              # This command can fail, specifically during the cleanup phase at the end
+              # of a build, when another worker has removed the stream key already.
             end
+          end
+
+          # Now, see if we want to claim any stale messages. If we are the only active
+          # consumer, we want to claim our own messages as well as messgaes from other
+          # (stale) consumers. If there are multiple active consumers, we are going to
+          # let another consumer claim our own messages.
+          if active_consumers.size > 1
+            stale_messages.reject! { |_key, message| message.worker_id == configuration.worker_id }
+          end
+
+          unless stale_messages.empty?
+            # When we have to reclaim stale tests, those test are potentially too slow
+            # to run inside the test timeout. We only claim one timed out test at a time in order
+            # to prevent the exact same batch from being too slow on repeated attempts,
+            # which would cause us to mark all the tests in that batch as failed.
+            #
+            # This has the side effect that for a retried test, the test timeout
+            # will be TEST_TIMEOUT_SECONDS * BATCH_SIZE in practice. This gives us a higher
+            # likelihood that the test will pass if the batch size > 1.
+            stale_messages = stale_messages.slice(stale_messages.keys.first)
+
+            enqueued_runnables = xclaim_messages(stale_messages, max_idle_time_ms: max_idle_time_ms)
+            reclaimed_timeout_tests.merge(enqueued_runnables)
+            return enqueued_runnables
+          end
+
+          # Now, see if we want to claim any failed tests to retry. Again, if we are the only
+          # active consumer, we want to claim our own messages as well as messgaes from other
+          # (stale) consumers. If there are multiple active consumers, we are going to let
+          # another consumer claim our own messages.
+          if active_consumers.size > 1
+            active_messages.reject! { |_key, message| message.worker_id == configuration.worker_id }
+          end
 
-            enqueued_runnables
+          # For all the active messages, we can check whether they are marked for a retry by
+          # trying to remove the test from the retry set set in Redis. Only one worker will be
+          # able to remove the entry from the set, so only one worker will end up trying to
+          # claim the test for the next attempt.
+          #
+          # We use `redis.multi` so we only need one round-trip for the entire list. Note that
+          # this is not an atomic operation with the XCLAIM call. This is OK, because the retry
+          # set is only there to speed things up and prevent us from having to wait for the test
+          # timeout. If the worker crashes between removing an item from the retry setm the test
+          # will eventually be picked up by another worker.
+          messages_in_retry_set = {}
+          redis.multi do
+            active_messages.each do |key, message|
+              messages_in_retry_set[key] = redis.srem(key('retry_set'), message.attempt_id)
+            end
           end
+
+          # Now, we only select the messages that were on the retry set, and try to claim them.
+          active_messages.keep_if { |key, _value| messages_in_retry_set.fetch(key).value }
+          enqueued_runnables = xclaim_messages(active_messages, max_idle_time_ms: 0)
+          reclaimed_failed_tests.merge(enqueued_runnables)
+          enqueued_runnables
         end
 
         sig { void }
@@ -288,18 +403,6 @@ module Minitest
           # so we can assume that all the Redis cleanup was completed.
         end
 
-        sig { params(er: EnqueuedRunnable).void }
-        def mark_runnable_to_fail_immediately(er)
-          assertion = Minitest::Assertion.new(<<~EOM.chomp)
-            This test takes too long to run (> #{configuration.test_timeout}s).
-
-            We have tried running this test #{configuration.max_attempts} on different workers, but every time the worker has not reported back a result within #{configuration.test_timeout}ms.
-            Try to make the test faster, or increase the test timeout.
-          EOM
-          assertion.set_backtrace(caller)
-          er.canned_failure = assertion
-        end
-
         sig { params(results: ResultAggregate).void }
         def adjust_combined_results(results)
           updated = redis.multi do
@@ -309,14 +412,16 @@ module Minitest
             redis.incrby(key('failures'), results.failures)
             redis.incrby(key('errors'), results.errors)
             redis.incrby(key('skips'), results.skips)
-            redis.incrby(key('reruns'), results.reruns)
+            redis.incrby(key('requeues'), results.requeues)
+            redis.incrby(key('discards'), results.discards)
             redis.incrby(key('acks'), results.acks)
             redis.incrby(key('size'), results.size)
           end
 
-          @combined_results = ResultAggregate.new(runs: updated[0], assertions: updated[1], passes: updated[2],
-            failures: updated[3], errors: updated[4], skips: updated[5], reruns: updated[6],
-            acks: updated[7], size: updated[8])
+          @combined_results = ResultAggregate.new(max_failures: configuration.max_failures,
+            runs: updated[0], assertions: updated[1], passes: updated[2],
+            failures: updated[3], errors: updated[4], skips: updated[5], requeues: updated[6], discards: updated[7],
+            acks: updated[8], size: updated[9])
         end
 
         sig { params(name: String).returns(String) }
@@ -324,59 +429,56 @@ module Minitest
           "minitest/#{configuration.run_id}/#{name}"
         end
 
-        sig { params(batch: T::Array[EnqueuedRunnable], reporter: AbstractReporter).returns(Integer) }
+        sig { params(name: String).returns(String) }
+        def list_key(name)
+          key("#{name}_list")
+        end
+
+        sig { params(batch: T::Array[EnqueuedRunnable], reporter: AbstractReporter).void }
         def process_batch(batch, reporter)
-          to_be_acked = {}
+          return 0 if batch.empty?
+
+          local_results.size += batch.size
+
+          runnable_results = T.let([], T::Array[EnqueuedRunnable::Result])
+          redis.multi do
+            batch.each do |enqueued_runnable|
+              # Fulfill the reporter contract by calling `prerecord` before we run the test.
+              reporter.prerecord(enqueued_runnable.runnable_class, enqueued_runnable.method_name)
+
+              # Actually run the test!
+              runnable_results << enqueued_runnable.run do |initial_result|
+                if ResultType.of(initial_result) == ResultType::Requeued
+                  sadd_future = redis.sadd(key('retry_set'), enqueued_runnable.attempt_id)
+                  EnqueuedRunnable::Result::Commit.new { sadd_future.value }
+                else
+                  xack_future = redis.xack(stream_key, group_name, enqueued_runnable.entry_id)
+                  EnqueuedRunnable::Result::Commit.new { xack_future.value == 1 }
+                end
+              end
+            end
+          end
+
+          batch_result_aggregate = ResultAggregate.new
+          runnable_results.each do |runnable_result|
+            # Complete the reporter contract by calling `record` with the result.
+            reporter.record(runnable_result.committed_result)
 
-          batch.each do |enqueued_runnable|
-            local_results.size += 1
-            reporter.prerecord(enqueued_runnable.runnable_class, enqueued_runnable.method_name)
-            result = enqueued_runnable.run
+            # Update statistics.
+            batch_result_aggregate.update_with_result(runnable_result)
+            local_results.update_with_result(runnable_result)
 
-            case (result_type = ResultType.of(result))
-            when ResultType::Passed
+            case (result_type = ResultType.of(runnable_result.committed_result))
+            when ResultType::Skipped, ResultType::Failed, ResultType::Error
+              redis.lpush(list_key(result_type.serialize), runnable_result.enqueued_runnable.identifier)
+            when ResultType::Passed, ResultType::Requeued, ResultType::Discarded
               # noop
-            when ResultType::Skipped
-              redis.lpush(key('skip_list'), Marshal.dump(enqueued_runnable.serialize))
-            when ResultType::Failed
-              redis.lpush(key('failure_list'), Marshal.dump(enqueued_runnable.serialize))
-            when ResultType::Error
-              redis.lpush(key('error_list'), Marshal.dump(enqueued_runnable.serialize))
             else
               T.absurd(result_type)
             end
-
-            local_results.update_with_result(result)
-            to_be_acked[enqueued_runnable.execution_id] = result
-          end
-
-          return 0 if to_be_acked.empty?
-
-          acked = redis.evalsha(
-            ack_batch_script,
-            keys: [stream_key],
-            argv: [group_name] + to_be_acked.keys
-          )
-
-          batch_results = ResultAggregate.new(acks: acked.length)
-          acked.each do |execution_id|
-            acked_result = to_be_acked.delete(execution_id)
-            reporter.record(acked_result)
-            batch_results.update_with_result(acked_result)
-          end
-
-          to_be_acked.each do |_execution_id, unacked_result|
-            # TODO: use custom assertion class.
-            discard_assertion = Minitest::Skip.new("The test result was discarded, " \
-              "because the test has been claimed another worker.")
-            discard_assertion.set_backtrace(caller)
-            unacked_result.failures = [discard_assertion]
-            reporter.record(unacked_result)
           end
 
-          adjust_combined_results(batch_results)
-          local_results.acks += acked.length
-          acked.length
+          adjust_combined_results(batch_result_aggregate)
         end
 
         INITIAL_BACKOFF = 10 # milliseconds