ruby: Simplify LostMaster recovery to an expiring setup marker

Replace the renewable initialization lease with two smaller pieces:

- The setup marker expires after half of queue_init_timeout, so a waiting
  worker can take over when the elected master dies before publishing.
- Publication (queue, chunk metadata, total, ready status) is one
  WATCH-guarded transaction, so if a slow master and its replacement both
  finish, only the first publishes and the other joins that queue.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

Author: ckwarteng

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

@@ -1,9 +1,7 @@
 # frozen_string_literal: true
 
 require 'ci/queue/static'
-require 'concurrent/atomic/event'
 require 'concurrent/set'
-require 'securerandom'
 
 module CI
   module Queue
@@ -19,8 +17,6 @@ class << self
 
       class Worker < Base
         DEFAULT_SLEEP_SECONDS = 0.5
-        InitializationLeaseLost = Class.new(StandardError)
-        private_constant :InitializationLeaseLost
         attr_reader :total
 
         def initialize(redis, config)
@@ -41,8 +37,8 @@ def populate(tests, random: Random.new)
           # All workers need an index of tests to resolve IDs
           @index = tests.map { |t| [t.id, t] }.to_h
           @total = tests.size
-          @ordering_random = random.dup
-          initialize_queue(tests)
+          @random = random
+          initialize_queue
 
           register_worker_presence
 
@@ -70,9 +66,7 @@ def idle?
         end
 
         def poll
-          wait_for_master(timeout: config.queue_init_timeout) do
-            initialize_queue(index.values)
-          end
+          wait_for_master(timeout: config.queue_init_timeout) { initialize_queue }
           if master?
             warn "Worker #{worker_id} is the master"
           else
@@ -270,66 +264,6 @@ def heartbeat(test_or_id = nil)
 
         attr_reader :index
 
-        def initialize_queue(tests)
-          return unless acquire_master_role?
-
-          stop_renewal = Concurrent::Event.new
-          renewal = Thread.new do
-            until stop_renewal.wait(master_lease_milliseconds / 3000.0)
-              begin
-                with_master_lease do |transaction|
-                  transaction.pexpire(key('master-lease'), master_lease_milliseconds)
-                end
-              rescue InitializationLeaseLost
-                break
-              rescue *CONNECTION_ERRORS => e
-                warn "Failed to renew queue initialization lease: #{e.message}"
-              end
-            end
-          end
-
-          executables = reorder_tests(tests, random: @ordering_random.dup)
-          chunks, individual_tests = executables.partition { |e| e.is_a?(CI::Queue::TestChunk) }
-          store_chunk_metadata(chunks) if chunks.any?
-          push(chunks.map(&:id) + individual_tests.map(&:id))
-        rescue InitializationLeaseLost
-          # Another process took over. Join its queue instead of publishing ours.
-          @master = false
-        ensure
-          if renewal
-            stop_renewal.set
-            renewal.join
-            release_initialization_lease
-          end
-        end
-
-        def master_lease_milliseconds
-          # Leave time for a surviving worker to take over within its setup wait.
-          [(config.queue_init_timeout.to_f * 500).ceil, 1].max
-        end
-
-        def with_master_lease
-          redis.watch(key('master-lease'), key('master-status')) do
-            unless redis.get(key('master-lease')) == @master_token && master_status == 'setup'
-              raise InitializationLeaseLost
-            end
-
-            result = redis.multi { |transaction| yield transaction }
-            raise InitializationLeaseLost unless result
-
-            result
-          end
-        end
-
-        def release_initialization_lease
-          with_master_lease do |transaction|
-            transaction.del(key('master-lease'), key('master-status'), key('master-worker-id'))
-          end
-          @master = false
-        rescue InitializationLeaseLost, *CONNECTION_ERRORS
-          # Published queues and replacement owners must remain untouched.
-        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)
@@ -503,21 +437,38 @@ def try_to_reserve_lost_test
           lost_test
         end
 
-        def push(tests)
+        def initialize_queue
+          return unless acquire_master_role?
+
+          executables = reorder_tests(index.values, random: @random)
+          chunks, individual_tests = executables.partition { |e| e.is_a?(CI::Queue::TestChunk) }
+          push(chunks.map(&:id) + individual_tests.map(&:id), chunks: chunks)
+        end
+
+        # The setup marker can expire under a slow master, so only the first
+        # master to finish publishes.
+        def push(tests, chunks: [])
           @total = tests.size
+          return unless @master
 
-          if @master
-            with_master_lease do |transaction|
-              transaction.lpush(key('queue'), tests) unless tests.empty?
-              transaction.set(key('total'), @total)
-              transaction.set(key('master-status'), 'ready')
-
-              transaction.expire(key('queue'), config.redis_ttl)
-              transaction.expire(key('total'), config.redis_ttl)
-              transaction.expire(key('master-status'), config.redis_ttl)
-              transaction.del(key('master-lease'))
+          published = redis.watch(key('master-status')) do
+            if queue_initialized?
+              redis.unwatch
+              nil
+            else
+              redis.multi do |transaction|
+                store_chunk_metadata(transaction, chunks)
+                transaction.lpush(key('queue'), tests) unless tests.empty?
+                transaction.set(key('total'), @total)
+                transaction.set(key('master-status'), 'ready')
+
+                transaction.expire(key('queue'), config.redis_ttl)
+                transaction.expire(key('total'), config.redis_ttl)
+                transaction.expire(key('master-status'), config.redis_ttl)
+              end
             end
           end
+          @master = false unless published
         rescue *CONNECTION_ERRORS
           raise if @master
         end
@@ -527,66 +478,64 @@ def register
         end
 
         def acquire_master_role?
-          @master_token = SecureRandom.hex(16)
-          @master = !!redis.watch(key('master-status'), key('master-lease')) do
-            if %w[ready finished].include?(master_status) || redis.get(key('master-lease'))
-              redis.unwatch
-              false
-            else
-              redis.multi do |transaction|
-                transaction.set(key('master-lease'), @master_token, px: master_lease_milliseconds)
-                transaction.set(key('master-status'), 'setup')
-                transaction.set(key('master-worker-id'), worker_id)
-                transaction.expire(key('master-worker-id'), config.redis_ttl)
-              end
+          return true if @master
+
+          @master = redis.set(key('master-status'), 'setup', nx: true, px: setup_timeout_milliseconds)
+          if @master
+            begin
+              redis.set(key('master-worker-id'), worker_id)
+              redis.expire(key('master-worker-id'), config.redis_ttl)
+              warn "Worker #{worker_id} elected as master"
+            rescue *CONNECTION_ERRORS
+              # If setting master-worker-id fails, we still have master status
+              # Log but don't lose master role
+              warn("Failed to set master-worker-id: #{$!.message}")
             end
           end
-          warn "Worker #{worker_id} elected as master" if @master
           @master
         rescue *CONNECTION_ERRORS
           @master = nil
           false
         end
 
+        def setup_timeout_milliseconds
+          # Leave a waiting worker time to take over within queue_init_timeout.
+          [(config.queue_init_timeout.to_f * 500).ceil, 1].max
+        end
+
         def register_worker_presence
           register
           redis.expire(key('workers'), config.redis_ttl)
         rescue *CONNECTION_ERRORS
           raise if master?
         end
 
-        def store_chunk_metadata(chunks)
-          # Batch operations to avoid exceeding Redis multi operation limits
-          # Each chunk requires 4 commands (set, expire, sadd, hset), so batch conservatively
-          batch_size = 5 # 5 chunks = 20 commands + 2 expires = 22 commands per batch
-
-          chunks.each_slice(batch_size) do |chunk_batch|
-            with_master_lease do |transaction|
-              chunk_batch.each do |chunk|
-                # Store chunk metadata with TTL
-                transaction.set(
-                  key('chunk', chunk.id),
-                  chunk.to_json
-                )
-                transaction.expire(key('chunk', chunk.id), config.redis_ttl)
-
-                # Track all chunks for cleanup
-                transaction.sadd(key('chunks'), chunk.id)
-
-                # Store dynamic timeout for this chunk
-                # Timeout = estimated_duration (in ms) converted to seconds + buffer
-                # estimated_duration is in milliseconds, convert to seconds and add 10% buffer
-                buffer_percent = 10
-                estimated_duration_seconds = chunk.estimated_duration / 1000.0
-                chunk_timeout = (estimated_duration_seconds * (1 + buffer_percent / 100.0)).round(2)
-                # Format to string to avoid floating point precision issues in Redis
-                # Use %g to remove trailing zeros
-                transaction.hset(key('test-group-timeout'), chunk.id, format('%g', chunk_timeout))
-              end
-              transaction.expire(key('chunks'), config.redis_ttl)
-              transaction.expire(key('test-group-timeout'), config.redis_ttl)
-            end
+        def store_chunk_metadata(transaction, chunks)
+          return if chunks.empty?
+
+          chunks.each do |chunk|
+            # Store chunk metadata with TTL
+            transaction.set(
+              key('chunk', chunk.id),
+              chunk.to_json
+            )
+            transaction.expire(key('chunk', chunk.id), config.redis_ttl)
+
+            # Track all chunks for cleanup
+            transaction.sadd(key('chunks'), chunk.id)
+
+            # Store dynamic timeout for this chunk
+            # Timeout = estimated_duration (in ms) converted to seconds + buffer
+            # estimated_duration is in milliseconds, convert to seconds and add 10% buffer
+            buffer_percent = 10
+            estimated_duration_seconds = chunk.estimated_duration / 1000.0
+            chunk_timeout = (estimated_duration_seconds * (1 + buffer_percent / 100.0)).round(2)
+            # Format to string to avoid floating point precision issues in Redis
+            # Use %g to remove trailing zeros
+            transaction.hset(key('test-group-timeout'), chunk.id, format('%g', chunk_timeout))
           end
+          transaction.expire(key('chunks'), config.redis_ttl)
+          transaction.expire(key('test-group-timeout'), config.redis_ttl)
         end
 
         def chunk_id?(id)

ruby/test/ci/queue/redis_initialization_test.rb

@@ -16,56 +16,44 @@ def teardown
   end
 
   def test_surviving_worker_recovers_when_initializer_exits
-    child = fork do
-      initializer = worker(0)
-      initializer.define_singleton_method(:reorder_tests) { |*| exit! 143 }
-      initializer.populate(tests)
-    end
-    _, status = Process.wait2(child)
-
-    assert_equal 143, status.exitstatus
-    assert_equal 'setup', @redis.get(key('master-status'))
-    assert_operator @redis.pttl(key('master-lease')), :>, 0
+    exit_during_setup(worker_id: 0)
 
     survivor = worker(1).populate(tests)
     refute_predicate survivor, :master?
     assert_equal tests.sort, poll(survivor).sort
     assert_predicate survivor, :master?
     assert_equal '1', @redis.get(key('master-worker-id'))
     assert_equal 'ready', @redis.get(key('master-status'))
-    assert_nil @redis.get(key('master-lease'))
   end
 
-  def test_worker_recovers_a_setup_marker_without_a_lease
-    @redis.set(key('master-status'), 'setup')
-    @redis.set(key('master-worker-id'), '0')
-
-    replacement = worker(1).populate(tests)
+  def test_retried_initializer_recovers_its_own_setup_marker
+    exit_during_setup(worker_id: 0)
 
-    assert_predicate replacement, :master?
-    assert_equal tests.sort, poll(replacement).sort
-    assert_equal '1', @redis.get(key('master-worker-id'))
+    retried = worker(0).populate(tests)
+    assert_equal tests.sort, poll(retried).sort
+    assert_predicate retried, :master?
+    assert_equal 'ready', @redis.get(key('master-status'))
   end
 
-  def test_setup_exception_releases_election
+  def test_worker_takes_over_after_setup_exception
     initializer = worker(0)
     initializer.define_singleton_method(:reorder_tests) { |*| raise 'setup failed' }
 
     error = assert_raises(RuntimeError) { initializer.populate(tests) }
     assert_equal 'setup failed', error.message
 
     replacement = worker(1).populate(tests)
-    assert_predicate replacement, :master?
     assert_equal tests.sort, poll(replacement).sort
+    assert_predicate replacement, :master?
   end
 
   def test_late_initializer_cannot_publish_over_replacement
     initializer = worker(0)
     replacement = worker(1)
     redis = @redis
-    lease_key = key('master-lease')
+    status_key = key('master-status')
     initializer.define_singleton_method(:reorder_tests) do |input, **|
-      redis.del(lease_key)
+      redis.del(status_key)
       replacement.populate(input)
       input.reverse
     end
@@ -82,11 +70,11 @@ def test_late_initializer_cannot_overwrite_chunk_metadata
     initializer = worker(0)
     replacement = worker(1, strategy: :suite_bin_packing)
     redis = @redis
-    lease_key = key('master-lease')
+    status_key = key('master-status')
     chunk_key = key('chunk:ATest:chunk_0')
     published_metadata = nil
     initializer.define_singleton_method(:reorder_tests) do |input, **|
-      redis.del(lease_key)
+      redis.del(status_key)
       replacement.populate(input)
       published_metadata = redis.get(chunk_key)
       [CI::Queue::TestChunk.new('ATest:chunk_0', 'ATest', [input.first.id], 9999)]
@@ -99,7 +87,7 @@ def test_late_initializer_cannot_overwrite_chunk_metadata
     assert_equal tests.sort, poll(replacement).flat_map(&:tests).sort
   end
 
-  def test_live_initializer_renews_lease_during_slow_setup
+  def test_waiting_worker_does_not_take_over_a_live_setup
     initializer = worker(0)
     started = Queue.new
     proceed = Queue.new
@@ -110,16 +98,16 @@ def test_live_initializer_renews_lease_during_slow_setup
     end
     thread = Thread.new { initializer.populate(tests) }
     started.pop
-    sleep 0.4 # Longer than the initial 0.25-second lease.
 
     contender = worker(1).populate(tests)
     refute_predicate contender, :master?
-    assert_equal '0', @redis.get(key('master-worker-id'))
+    assert_equal 'setup', @redis.get(key('master-status'))
 
     proceed << true
     assert thread.join(2), 'initializer did not complete'
     assert_predicate initializer, :master?
     assert_equal tests.sort, poll(contender).sort
+    refute_predicate contender, :master?
   ensure
     proceed << true if proceed
     thread&.kill
@@ -128,6 +116,19 @@ def test_live_initializer_renews_lease_during_slow_setup
 
   private
 
+  def exit_during_setup(worker_id:)
+    child = fork do
+      initializer = worker(worker_id)
+      initializer.define_singleton_method(:reorder_tests) { |*| exit! 143 }
+      initializer.populate(tests)
+    end
+    _, status = Process.wait2(child)
+
+    assert_equal 143, status.exitstatus
+    assert_equal 'setup', @redis.get(key('master-status'))
+    assert_operator @redis.pttl(key('master-status')), :>, 0
+  end
+
   def tests
     SharedQueueAssertions::TEST_LIST.dup
   end