Merge pull request #103 from figma/ckwarteng/recover-ci-queue-initialization

ruby / ci-queue: Recover queue initialization when the master exits

Author: ckwarteng

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

@@ -59,6 +59,9 @@ def wait_for_master(timeout: 120)
           (timeout * 10 + 1).to_i.times do
             return true if queue_initialized?
 
+            yield if block_given?
+            return true if queue_initialized?
+
             sleep 0.1
           end
           raise LostMaster, "The master worker (worker #{master_worker_id}) is still `#{master_status}` after #{timeout} seconds waiting."

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

@@ -2,6 +2,7 @@
 
 require 'ci/queue/static'
 require 'concurrent/set'
+require 'securerandom'
 
 module CI
   module Queue
@@ -37,18 +38,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
-
-          if acquire_master_role?
-            executables = reorder_tests(tests, random: random)
-
-            chunks = executables.select { |e| e.is_a?(CI::Queue::TestChunk) }
-            individual_tests = executables.reject { |e| e.is_a?(CI::Queue::TestChunk) }
-
-            store_chunk_metadata(chunks) if chunks.any?
-
-            all_ids = chunks.map(&:id) + individual_tests.map(&:id)
-            push(all_ids)
-          end
+          @random = random
+          initialize_queue
 
           register_worker_presence
 
@@ -76,7 +67,7 @@ def idle?
         end
 
         def poll
-          wait_for_master
+          wait_for_master(timeout: config.queue_init_timeout) { initialize_queue }
           if master?
             warn "Worker #{worker_id} is the master"
           else
@@ -447,20 +438,47 @@ def try_to_reserve_lost_test
           lost_test
         end
 
+        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) }
+          # An initializer may outlive its setup marker. Give its chunks unique
+          # IDs so staging metadata cannot overwrite another initializer's data.
+          initialization_id = SecureRandom.hex(8)
+          chunks = chunks.map do |chunk|
+            CI::Queue::TestChunk.new(
+              "#{chunk.id}:#{initialization_id}", chunk.suite_name,
+              chunk.test_ids, chunk.estimated_duration, test_count: chunk.test_count
+            )
+          end
+          store_chunk_metadata(chunks)
+          push(chunks.map(&:id) + individual_tests.map(&:id))
+        end
+
+        # The setup marker can expire under a slow master, so only the first
+        # master to finish publishes.
         def push(tests)
           @total = tests.size
+          return unless @master
 
-          if @master
-            redis.multi 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)
+          published = redis.watch(key('master-status')) do
+            if queue_initialized?
+              redis.unwatch
+              nil
+            else
+              redis.multi 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)
+              end
             end
           end
+          @master = false unless published
         rescue *CONNECTION_ERRORS
           raise if @master
         end
@@ -472,7 +490,7 @@ def register
         def acquire_master_role?
           return true if @master
 
-          @master = redis.setnx(key('master-status'), 'setup')
+          @master = redis.set(key('master-status'), 'setup', nx: true, px: setup_timeout_milliseconds)
           if @master
             begin
               redis.set(key('master-worker-id'), worker_id)
@@ -490,6 +508,11 @@ def acquire_master_role?
           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)
@@ -498,9 +521,8 @@ def register_worker_presence
         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
+          # Four commands per chunk, plus two expiry commands: at most 22.
+          batch_size = 5
 
           chunks.each_slice(batch_size) do |chunk_batch|
             redis.multi do |transaction|

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

@@ -38,7 +38,7 @@ def test_chunk_timeout_stored_in_redis_hash
 
     # Verify timeout was stored in test-group-timeout hash
     # Timeout should be: estimated_duration (5000ms = 5s) * 1.1 buffer = 5.5s
-    chunk_timeout = @redis.hget('build:42:test-group-timeout', 'TestA:chunk_0')
+    chunk_timeout = @redis.hget('build:42:test-group-timeout', stored_chunk_id('TestA:chunk_0'))
     refute_nil chunk_timeout
     assert_equal '5.5', chunk_timeout
   end
@@ -54,7 +54,7 @@ def test_chunk_timeout_scales_with_test_count
       worker.populate(small_tests)
     end
 
-    small_timeout = @redis.hget('build:42:test-group-timeout', 'SmallSuite:chunk_0')
+    small_timeout = @redis.hget('build:42:test-group-timeout', stored_chunk_id('SmallSuite:chunk_0'))
     assert_equal '1.1', small_timeout # 1000ms = 1s * 1.1 buffer = 1.1s
 
     @redis.flushdb
@@ -69,7 +69,7 @@ def test_chunk_timeout_scales_with_test_count
       worker.populate(large_tests)
     end
 
-    large_timeout = @redis.hget('build:42:test-group-timeout', 'LargeSuite:chunk_0')
+    large_timeout = @redis.hget('build:42:test-group-timeout', stored_chunk_id('LargeSuite:chunk_0'))
     assert_equal '5.5', large_timeout # 5000ms = 5s * 1.1 buffer = 5.5s
   end
 
@@ -92,9 +92,9 @@ def test_multiple_chunks_stored_with_different_timeouts
     end
 
     # Verify each chunk has correct timeout
-    assert_equal '2.2', @redis.hget('build:42:test-group-timeout', 'TestA:chunk_0') # 2000ms = 2s * 1.1 buffer = 2.2s
-    assert_equal '3.3', @redis.hget('build:42:test-group-timeout', 'TestB:chunk_0') # 3000ms = 3s * 1.1 buffer = 3.3s
-    assert_equal '1.1', @redis.hget('build:42:test-group-timeout', 'TestC:chunk_0') # 1000ms = 1s * 1.1 buffer = 1.1s
+    assert_equal '2.2', @redis.hget('build:42:test-group-timeout', stored_chunk_id('TestA:chunk_0')) # 2000ms = 2s * 1.1 buffer = 2.2s
+    assert_equal '3.3', @redis.hget('build:42:test-group-timeout', stored_chunk_id('TestB:chunk_0')) # 3000ms = 3s * 1.1 buffer = 3.3s
+    assert_equal '1.1', @redis.hget('build:42:test-group-timeout', stored_chunk_id('TestC:chunk_0')) # 1000ms = 1s * 1.1 buffer = 1.1s
   end
 
   def test_timeout_hash_has_ttl
@@ -145,8 +145,8 @@ def test_mixed_chunks_and_tests_only_chunks_have_timeouts
     end
 
     # Chunks should have timeouts
-    assert_equal '2.2', @redis.hget('build:42:test-group-timeout', 'TestA:chunk_0') # 2000ms = 2s * 1.1 buffer = 2.2s
-    assert_equal '3.3', @redis.hget('build:42:test-group-timeout', 'TestC:chunk_0') # 3000ms = 3s * 1.1 buffer = 3.3s
+    assert_equal '2.2', @redis.hget('build:42:test-group-timeout', stored_chunk_id('TestA:chunk_0')) # 2000ms = 2s * 1.1 buffer = 2.2s
+    assert_equal '3.3', @redis.hget('build:42:test-group-timeout', stored_chunk_id('TestC:chunk_0')) # 3000ms = 3s * 1.1 buffer = 3.3s
 
     # Individual test should not
     assert_nil @redis.hget('build:42:test-group-timeout', 'TestB#test_1')
@@ -235,7 +235,7 @@ def test_chunk_not_marked_lost_before_dynamic_timeout
 
     # Reserve the chunk with worker1
     reserved_id = worker1.send(:try_to_reserve_test)
-    assert_equal 'TestSuite:chunk_0', reserved_id
+    assert_match /\ATestSuite:chunk_0:[0-9a-f]{16}\z/, reserved_id
 
     # Wait 1 second (less than 2.5s timeout)
     sleep 1
@@ -296,14 +296,18 @@ def test_batching_with_many_chunks
 
     # Verify all chunks have timeouts stored despite batching
     chunks.each do |chunk|
-      timeout = @redis.hget('build:42:test-group-timeout', chunk.id)
+      timeout = @redis.hget('build:42:test-group-timeout', stored_chunk_id(chunk.id))
       refute_nil timeout, "Chunk #{chunk.id} should have timeout stored"
       assert_equal '1.1', timeout # 1000ms = 1s * 1.1 buffer = 1.1s
     end
   end
 
   private
 
+  def stored_chunk_id(original_id)
+    @redis.smembers('build:42:chunks').find { |id| id.start_with?("#{original_id}:") }
+  end
+
   def create_mock_tests(test_ids)
     test_ids.map { |id| MockTest.new(id) }
   end

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

@@ -36,7 +36,7 @@ def test_populate_stores_chunk_metadata_in_redis
     end
 
     # Verify chunk metadata was stored
-    chunk_data = @redis.get('build:42:chunk:TestA:chunk_0')
+    chunk_data = @redis.get("build:42:chunk:#{stored_chunk_id('TestA:chunk_0')}")
     refute_nil chunk_data
 
     parsed = JSON.parse(chunk_data)
@@ -100,10 +100,10 @@ def test_resolve_chunk
     end
 
     # Manually store and resolve chunk
-    resolved = @worker.send(:resolve_chunk, 'TestA:chunk_0')
+    resolved = @worker.send(:resolve_chunk, stored_chunk_id('TestA:chunk_0'))
 
     assert_instance_of CI::Queue::Redis::Worker::ResolvedChunk, resolved
-    assert_equal 'TestA:chunk_0', resolved.id
+    assert_match /\ATestA:chunk_0(?::[0-9a-f]{16})?\z/, resolved.id
     assert_equal 'TestA', resolved.suite_name
     assert_equal 3, resolved.tests.size
     assert_equal test_ids, resolved.tests.map(&:id)
@@ -120,10 +120,10 @@ def test_resolve_chunk_with_subset_of_tests
       @worker.populate(tests)
     end
 
-    resolved = @worker.send(:resolve_chunk, 'TestA:chunk_0')
+    resolved = @worker.send(:resolve_chunk, stored_chunk_id('TestA:chunk_0'))
 
     assert_instance_of CI::Queue::Redis::Worker::ResolvedChunk, resolved
-    assert_equal 'TestA:chunk_0', resolved.id
+    assert_match /\ATestA:chunk_0(?::[0-9a-f]{16})?\z/, resolved.id
     assert_equal 2, resolved.tests.size
     assert_equal ['TestA#test_1', 'TestA#test_3'], resolved.tests.map(&:id)
   end
@@ -140,7 +140,7 @@ def test_resolved_chunk_interface
 
     resolved = CI::Queue::Redis::Worker::ResolvedChunk.new(chunk, tests)
 
-    assert_equal 'TestA:chunk_0', resolved.id
+    assert_match /\ATestA:chunk_0(?::[0-9a-f]{16})?\z/, resolved.id
     assert_equal 'TestA', resolved.suite_name
     assert_equal 2, resolved.size
     assert resolved.chunk?
@@ -195,7 +195,7 @@ def test_populate_with_mixed_chunks_and_tests
 
     # Check that both chunk and individual test IDs are in queue
     queue_items = @redis.lrange('build:42:queue', 0, -1)
-    assert_includes queue_items, 'TestA:chunk_0'
+    assert_includes queue_items, stored_chunk_id('TestA:chunk_0')
     assert_includes queue_items, 'TestB#test_1'
   end
 
@@ -210,12 +210,12 @@ def test_chunk_metadata_has_ttl
       @worker.populate(tests)
     end
 
-    ttl = @redis.ttl('build:42:chunk:TestA:chunk_0')
+    ttl = @redis.ttl("build:42:chunk:#{stored_chunk_id('TestA:chunk_0')}")
     assert ttl > 0, 'Chunk metadata should have TTL set'
   end
 
   def test_populate_with_many_chunks_uses_batching
-    # Create 20 test chunks to verify batching works (batch size is 7)
+    # Create 20 test chunks to verify batching works (batch size is 5)
     tests = (1..20).map { |i| MockTest.new("TestSuite#{i}#test_1") }
     chunks = (1..20).map do |i|
       test_ids = ["TestSuite#{i}#test_1"]
@@ -228,7 +228,7 @@ def test_populate_with_many_chunks_uses_batching
 
     # Verify all chunks were stored despite batching
     chunks.each do |chunk|
-      chunk_data = @redis.get("build:42:chunk:#{chunk.id}")
+      chunk_data = @redis.get("build:42:chunk:#{stored_chunk_id(chunk.id)}")
       refute_nil chunk_data, "Chunk #{chunk.id} should be stored"
 
       parsed = JSON.parse(chunk_data)
@@ -239,7 +239,7 @@ def test_populate_with_many_chunks_uses_batching
     # Verify all chunk IDs are in the chunks set
     stored_chunks = @redis.smembers('build:42:chunks')
     chunks.each do |chunk|
-      assert_includes stored_chunks, chunk.id
+      assert_includes stored_chunks, stored_chunk_id(chunk.id)
     end
   end
 
@@ -257,7 +257,7 @@ def test_worker_history_retry_expands_chunks_and_deduplicates_tests
     end
 
     history_key = 'build:42:worker:1:queue'
-    [chunk.id, tests.last.id, chunk.id, tests[1].id].each do |id|
+    [stored_chunk_id(chunk.id), tests.last.id, stored_chunk_id(chunk.id), tests[1].id].each do |id|
       @redis.lpush(history_key, id)
     end
 
@@ -304,6 +304,10 @@ def test_worker_history_retry_rejects_invalid_chunk_metadata
 
   private
 
+  def stored_chunk_id(original_id)
+    @redis.smembers('build:42:chunks').find { |id| id.start_with?("#{original_id}:") }
+  end
+
   def create_mock_tests(test_ids)
     test_ids.map { |id| MockTest.new(id) }
   end