Add check to prevent same worker from picking up requeued tests

Author: danielahrnsbrak

redis/acknowledge.lua

@@ -2,13 +2,15 @@ local zset_key = KEYS[1]
 local processed_key = KEYS[2]
 local owners_key = KEYS[3]
 local error_reports_key = KEYS[4]
+local requeued_by_key = KEYS[5]
 
 local entry = ARGV[1]
 local test_id = ARGV[2]
 local error = ARGV[3]
 local ttl = ARGV[4]
 redis.call('zrem', zset_key, entry)
 redis.call('hdel', owners_key, entry)  -- Doesn't matter if it was reclaimed by another workers
+redis.call('hdel', requeued_by_key, entry)
 local acknowledged = redis.call('sadd', processed_key, test_id) == 1
 
 if acknowledged and error ~= "" then

redis/requeue.lua

@@ -5,12 +5,14 @@ local zset_key = KEYS[4]
 local worker_queue_key = KEYS[5]
 local owners_key = KEYS[6]
 local error_reports_key = KEYS[7]
+local requeued_by_key = KEYS[8]
 
 local max_requeues = tonumber(ARGV[1])
 local global_max_requeues = tonumber(ARGV[2])
 local entry = ARGV[3]
 local test_id = ARGV[4]
 local offset = ARGV[5]
+local ttl = tonumber(ARGV[6])
 
 if redis.call('hget', owners_key, entry) == worker_queue_key then
    redis.call('hdel', owners_key, entry)
@@ -42,6 +44,11 @@ else
   redis.call('lpush', queue_key, entry)
 end
 
+redis.call('hset', requeued_by_key, entry, worker_queue_key)
+if ttl and ttl > 0 then
+  redis.call('expire', requeued_by_key, ttl)
+end
+
 redis.call('zrem', zset_key, entry)
 
 return true

redis/reserve.lua

@@ -3,15 +3,54 @@ local zset_key = KEYS[2]
 local processed_key = KEYS[3]
 local worker_queue_key = KEYS[4]
 local owners_key = KEYS[5]
+local requeued_by_key = KEYS[6]
+local workers_key = KEYS[7]
 
 local current_time = ARGV[1]
+local defer_offset = tonumber(ARGV[2]) or 0
+local max_skip_attempts = 4
 
-local test = redis.call('rpop', queue_key)
-if test then
-  redis.call('zadd', zset_key, current_time, test)
-  redis.call('lpush', worker_queue_key, test)
-  redis.call('hset', owners_key, test, worker_queue_key)
-  return test
-else
-  return nil
+local function insert_with_offset(test)
+  local pivot = redis.call('lrange', queue_key, -1 - defer_offset, 0 - defer_offset)[1]
+  if pivot then
+    redis.call('linsert', queue_key, 'BEFORE', pivot, test)
+  else
+    redis.call('lpush', queue_key, test)
+  end
 end
+
+for attempt = 1, max_skip_attempts do
+  local test = redis.call('rpop', queue_key)
+  if not test then
+    return nil
+  end
+
+  local requeued_by = redis.call('hget', requeued_by_key, test)
+  if requeued_by == worker_queue_key then
+    -- If this build only has one worker, allow immediate self-pickup.
+    if redis.call('scard', workers_key) <= 1 then
+      redis.call('hdel', requeued_by_key, test)
+      redis.call('zadd', zset_key, current_time, test)
+      redis.call('lpush', worker_queue_key, test)
+      redis.call('hset', owners_key, test, worker_queue_key)
+      return test
+    end
+
+    insert_with_offset(test)
+
+    -- If this worker only finds its own requeued tests, defer once by returning nil,
+    -- then allow pickup on a subsequent reserve attempt.
+    if attempt == max_skip_attempts then
+      redis.call('hdel', requeued_by_key, test)
+      return nil
+    end
+  else
+    redis.call('hdel', requeued_by_key, test)
+    redis.call('zadd', zset_key, current_time, test)
+    redis.call('lpush', worker_queue_key, test)
+    redis.call('hset', owners_key, test, worker_queue_key)
+    return test
+  end
+end
+
+return nil

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

@@ -183,7 +183,7 @@ def acknowledge(test_key, error: nil, pipeline: redis)
           unreserve_entry(test_id)
           eval_script(
             :acknowledge,
-            keys: [key('running'), key('processed'), key('owners'), key('error-reports')],
+            keys: [key('running'), key('processed'), key('owners'), key('error-reports'), key('requeued-by')],
             argv: [entry, test_id, error.to_s, config.redis_ttl],
             pipeline: pipeline,
           ) == 1
@@ -206,8 +206,9 @@ def requeue(test, offset: Redis.requeue_offset)
               key('worker', worker_id, 'queue'),
               key('owners'),
               key('error-reports'),
+              key('requeued-by'),
             ],
-            argv: [config.max_requeues, global_max_requeues, entry, test_id, offset],
+            argv: [config.max_requeues, global_max_requeues, entry, test_id, offset, config.redis_ttl],
           ) == 1
 
           unless requeued
@@ -372,8 +373,10 @@ def try_to_reserve_test
               key('processed'),
               key('worker', worker_id, 'queue'),
               key('owners'),
+              key('requeued-by'),
+              key('workers'),
             ],
-            argv: [CI::Queue.time_now.to_f],
+            argv: [CI::Queue.time_now.to_f, Redis.requeue_offset],
           )
         end
 

ruby/lib/minitest/queue/local_requeue_reporter.rb

@@ -46,6 +46,17 @@ def message_for(test)
       def result_line
         "#{super}, #{requeues} requeues"
       end
+
+      def location(exception)
+        backtrace = exception.backtrace
+        return super if backtrace && !backtrace.empty?
+
+        nested_exception = exception.respond_to?(:error) ? exception.error : nil
+        nested_backtrace = nested_exception&.backtrace
+        return super(nested_exception) if nested_backtrace && !nested_backtrace.empty?
+
+        'unknown'
+      end
     end
   end
 end

ruby/lib/minitest/queue/test_data.rb

@@ -137,11 +137,13 @@ def relative_path_for(path)
       def error_location(exception)
         @error_location ||= begin
           last_before_assertion = ''
-          exception.backtrace.reverse_each do |s|
+          backtrace_for(exception).reverse_each do |s|
             break if s =~ /in .(assert|refute|flunk|pass|fail|raise|must|wont)/
 
             last_before_assertion = s
           end
+          return ['unknown', 0] if last_before_assertion.empty?
+
           path = last_before_assertion.sub(/:in .*$/, '')
           # the path includes the linenumber at the end,
           # which is seperated by a :
@@ -151,6 +153,17 @@ def error_location(exception)
           [result.first, result.last.to_i]
         end
       end
+
+      def backtrace_for(exception)
+        backtrace = exception.backtrace
+        return backtrace if backtrace && !backtrace.empty?
+
+        nested_exception = exception.respond_to?(:error) ? exception.error : nil
+        nested_backtrace = nested_exception&.backtrace
+        return nested_backtrace if nested_backtrace && !nested_backtrace.empty?
+
+        []
+      end
     end
   end
 end

ruby/test/ci/queue/redis_test.rb

@@ -304,6 +304,28 @@ def test_streaming_timeout_raises_lost_master
     end
   end
 
+  def test_reserve_defers_own_requeued_test_once
+    queue = worker(1, populate: false, build_id: 'self-requeue-script')
+    queue.send(:register)
+    entry = "ATest#test_foo#{DELIMITER}/tmp/a_test.rb"
+    queue_key = queue.send(:key, 'queue')
+    requeued_by_key = queue.send(:key, 'requeued-by')
+    worker_queue_key = queue.send(:key, 'worker', queue.config.worker_id, 'queue')
+    workers_key = queue.send(:key, 'workers')
+
+    @redis.lpush(queue_key, entry)
+    @redis.hset(requeued_by_key, entry, worker_queue_key)
+    @redis.sadd(workers_key, '2')
+
+    first_try = queue.send(:try_to_reserve_test)
+    assert_nil first_try
+    assert_equal [entry], @redis.lrange(queue_key, 0, -1)
+    assert_nil @redis.hget(requeued_by_key, entry)
+
+    second_try = queue.send(:try_to_reserve_test)
+    assert_equal entry, second_try
+  end
+
   def test_heartbeat_uses_test_id_for_processed_check
     queue = worker(1, populate: false)
     entry = "ATest#test_foo#{DELIMITER}/tmp/a_test.rb"
@@ -428,6 +450,87 @@ def test_worker_profiles_aggregates_multiple_workers
     assert_equal 'non-leader', profiles['2']['role']
   end
 
+  def test_worker_does_not_pick_up_its_own_requeued_test_when_others_are_available
+    @redis.flushdb
+
+    test_list = TEST_LIST.first(3)
+    w1 = worker(1, tests: test_list, build_id: 'self-requeue', timeout: 10, max_requeues: 1, requeue_tolerance: 1.0)
+    w2 = worker(2, populate: false, build_id: 'self-requeue', timeout: 10, max_requeues: 1, requeue_tolerance: 1.0)
+    w3 = worker(3, populate: false, build_id: 'self-requeue', timeout: 10, max_requeues: 1, requeue_tolerance: 1.0)
+    w2.send(:register)
+    w3.send(:register)
+
+    id_for = ->(test) { test.respond_to?(:id) ? test.id : CI::Queue::QueueEntry.test_id(test) }
+
+    requeued_test_id = nil
+    picked_up_requeue = {}
+    worker_two_reserved = false
+    worker_three_reserved = false
+    release_other_workers = false
+
+    mon = Monitor.new
+    cond = mon.new_cond
+
+    threads = [
+      Thread.new do
+        w2.poll do |test|
+          test_id = id_for.call(test)
+          mon.synchronize do
+            worker_two_reserved = true
+            picked_up_requeue['2'] = true if test_id == requeued_test_id
+            cond.broadcast
+            cond.wait_until { release_other_workers }
+          end
+          w2.acknowledge(test_id)
+        end
+      end,
+      Thread.new do
+        w3.poll do |test|
+          test_id = id_for.call(test)
+          mon.synchronize do
+            worker_three_reserved = true
+            picked_up_requeue['3'] = true if test_id == requeued_test_id
+            cond.broadcast
+            cond.wait_until { release_other_workers }
+          end
+          w3.acknowledge(test_id)
+        end
+      end,
+    ]
+
+    mon.synchronize do
+      cond.wait_until { worker_two_reserved && worker_three_reserved }
+    end
+
+    worker_one_picked_its_own_requeue = false
+    first_test = true
+
+    w1.poll do |test|
+      test_id = id_for.call(test)
+
+      if first_test
+        first_test = false
+        requeued_test_id = test_id
+        w1.report_failure!
+        assert_equal true, w1.requeue(test)
+        mon.synchronize do
+          release_other_workers = true
+          cond.broadcast
+        end
+      else
+        worker_one_picked_its_own_requeue = true if test_id == requeued_test_id
+        w1.acknowledge(test_id)
+      end
+    end
+
+    threads.each { |t| t.join(5) }
+
+    assert_equal false, worker_one_picked_its_own_requeue
+    assert_equal true, picked_up_requeue.values.any?
+  ensure
+    threads&.each(&:kill)
+  end
+
   private
 
   def shuffled_test_list

ruby/test/minitest/queue/local_requeue_reporter_test.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+require 'test_helper'
+
+module Minitest::Queue
+  class LocalRequeueReporterTest < Minitest::Test
+    include ReporterTestHelper
+
+    def setup
+      @reporter = LocalRequeueReporter.new(verbose: true)
+    end
+
+    def test_message_for_requeued_failure_without_backtrace
+      test = result('test_foo', requeued: true)
+      test.failure.failure.set_backtrace(nil)
+
+      message = @reporter.send(:message_for, test)
+
+      assert_includes message, '[unknown]'
+      assert_includes message, 'Failed'
+    end
+  end
+end

ruby/test/minitest/queue/test_data_test.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+require 'test_helper'
+
+module Minitest::Queue
+  class TestDataTest < Minitest::Test
+    include ReporterTestHelper
+
+    def test_error_location_without_backtrace
+      failure = Minitest::Assertion.new('Assertion failed')
+      test = result('test_foo', failure: failure)
+      failure.set_backtrace(nil)
+
+      data = TestData.new(
+        test: test,
+        index: 0,
+        namespace: 'namespace',
+        base_path: Minitest::Queue.project_root,
+      ).to_h
+
+      assert_equal 'unknown', data[:error_file_path]
+      assert_equal 0, data[:error_file_number]
+    end
+
+    def test_error_location_uses_nested_exception_backtrace
+      error = StandardError.new('boom')
+      error.set_backtrace([
+        "#{Minitest::Queue.project_root}/test/nested_error_test.rb:42:in `boom'",
+      ])
+      failure = Minitest::UnexpectedError.new(error)
+      failure.set_backtrace(nil)
+      test = result('test_foo', failure: failure)
+
+      data = TestData.new(
+        test: test,
+        index: 0,
+        namespace: 'namespace',
+        base_path: Minitest::Queue.project_root,
+      ).to_h
+
+      assert_equal 'test/nested_error_test.rb', data[:error_file_path].to_s
+      assert_equal 42, data[:error_file_number]
+    end
+  end
+end