11# frozen_string_literal: true
22
33require 'ci/queue/static'
4- require 'concurrent/atomic/event'
54require 'concurrent/set'
6- require 'securerandom'
75
86module CI
97 module Queue
@@ -19,8 +17,6 @@ class << self
1917
2018 class Worker < Base
2119 DEFAULT_SLEEP_SECONDS = 0.5
22- InitializationLeaseLost = Class . new ( StandardError )
23- private_constant :InitializationLeaseLost
2420 attr_reader :total
2521
2622 def initialize ( redis , config )
@@ -41,8 +37,8 @@ def populate(tests, random: Random.new)
4137 # All workers need an index of tests to resolve IDs
4238 @index = tests . map { |t | [ t . id , t ] } . to_h
4339 @total = tests . size
44- @ordering_random = random . dup
45- initialize_queue ( tests )
40+ @random = random
41+ initialize_queue
4642
4743 register_worker_presence
4844
@@ -70,9 +66,9 @@ def idle?
7066 end
7167
7268 def poll
73- wait_for_master ( timeout : config . queue_init_timeout ) do
74- initialize_queue ( index . values )
75- end
69+ # If the elected worker dies before publishing, its setup marker
70+ # expires and a waiting worker initializes the queue instead.
71+ wait_for_master ( timeout : config . queue_init_timeout ) { initialize_queue }
7672 if master?
7773 warn "Worker #{ worker_id } is the master"
7874 else
@@ -270,66 +266,6 @@ def heartbeat(test_or_id = nil)
270266
271267 attr_reader :index
272268
273- def initialize_queue ( tests )
274- return unless acquire_master_role?
275-
276- stop_renewal = Concurrent ::Event . new
277- renewal = Thread . new do
278- until stop_renewal . wait ( master_lease_milliseconds / 3000.0 )
279- begin
280- with_master_lease do |transaction |
281- transaction . pexpire ( key ( 'master-lease' ) , master_lease_milliseconds )
282- end
283- rescue InitializationLeaseLost
284- break
285- rescue *CONNECTION_ERRORS => e
286- warn "Failed to renew queue initialization lease: #{ e . message } "
287- end
288- end
289- end
290-
291- executables = reorder_tests ( tests , random : @ordering_random . dup )
292- chunks , individual_tests = executables . partition { |e | e . is_a? ( CI ::Queue ::TestChunk ) }
293- store_chunk_metadata ( chunks ) if chunks . any?
294- push ( chunks . map ( &:id ) + individual_tests . map ( &:id ) )
295- rescue InitializationLeaseLost
296- # Another process took over. Join its queue instead of publishing ours.
297- @master = false
298- ensure
299- if renewal
300- stop_renewal . set
301- renewal . join
302- release_initialization_lease
303- end
304- end
305-
306- def master_lease_milliseconds
307- # Leave time for a surviving worker to take over within its setup wait.
308- [ ( config . queue_init_timeout . to_f * 500 ) . ceil , 1 ] . max
309- end
310-
311- def with_master_lease
312- redis . watch ( key ( 'master-lease' ) , key ( 'master-status' ) ) do
313- unless redis . get ( key ( 'master-lease' ) ) == @master_token && master_status == 'setup'
314- raise InitializationLeaseLost
315- end
316-
317- result = redis . multi { |transaction | yield transaction }
318- raise InitializationLeaseLost unless result
319-
320- result
321- end
322- end
323-
324- def release_initialization_lease
325- with_master_lease do |transaction |
326- transaction . del ( key ( 'master-lease' ) , key ( 'master-status' ) , key ( 'master-worker-id' ) )
327- end
328- @master = false
329- rescue InitializationLeaseLost , *CONNECTION_ERRORS
330- # Published queues and replacement owners must remain untouched.
331- end
332-
333269 # Runs a block while sending periodic heartbeats in a background thread.
334270 # This prevents other workers from stealing the test while it's being executed.
335271 def with_heartbeat ( test_id )
@@ -503,21 +439,40 @@ def try_to_reserve_lost_test
503439 lost_test
504440 end
505441
506- def push ( tests )
442+ def initialize_queue
443+ return unless acquire_master_role?
444+
445+ executables = reorder_tests ( index . values , random : @random )
446+ chunks , individual_tests = executables . partition { |e | e . is_a? ( CI ::Queue ::TestChunk ) }
447+ push ( chunks . map ( &:id ) + individual_tests . map ( &:id ) , chunks : chunks )
448+ end
449+
450+ # Publishes the queue, its chunk metadata, and the ready status in one
451+ # transaction. The setup marker can expire while a slow master is still
452+ # ordering tests, so WATCH ensures only the first master to finish
453+ # publishes; any other master joins that queue instead.
454+ def push ( tests , chunks : [ ] )
507455 @total = tests . size
456+ return unless @master
508457
509- if @master
510- with_master_lease do |transaction |
511- transaction . lpush ( key ( 'queue' ) , tests ) unless tests . empty?
512- transaction . set ( key ( 'total' ) , @total )
513- transaction . set ( key ( 'master-status' ) , 'ready' )
514-
515- transaction . expire ( key ( 'queue' ) , config . redis_ttl )
516- transaction . expire ( key ( 'total' ) , config . redis_ttl )
517- transaction . expire ( key ( 'master-status' ) , config . redis_ttl )
518- transaction . del ( key ( 'master-lease' ) )
458+ published = redis . watch ( key ( 'master-status' ) ) do
459+ if queue_initialized?
460+ redis . unwatch
461+ nil
462+ else
463+ redis . multi do |transaction |
464+ store_chunk_metadata ( transaction , chunks )
465+ transaction . lpush ( key ( 'queue' ) , tests ) unless tests . empty?
466+ transaction . set ( key ( 'total' ) , @total )
467+ transaction . set ( key ( 'master-status' ) , 'ready' )
468+
469+ transaction . expire ( key ( 'queue' ) , config . redis_ttl )
470+ transaction . expire ( key ( 'total' ) , config . redis_ttl )
471+ transaction . expire ( key ( 'master-status' ) , config . redis_ttl )
472+ end
519473 end
520474 end
475+ @master = false unless published
521476 rescue *CONNECTION_ERRORS
522477 raise if @master
523478 end
@@ -527,66 +482,65 @@ def register
527482 end
528483
529484 def acquire_master_role?
530- @master_token = SecureRandom . hex ( 16 )
531- @master = !! redis . watch ( key ( 'master-status' ) , key ( 'master-lease' ) ) do
532- if %w[ ready finished ] . include? ( master_status ) || redis . get ( key ( 'master-lease' ) )
533- redis . unwatch
534- false
535- else
536- redis . multi do | transaction |
537- transaction . set ( key ( 'master-lease' ) , @master_token , px : master_lease_milliseconds )
538- transaction . set ( key ( 'master-status' ) , 'setup' )
539- transaction . set ( key ( ' master-worker-id' ) , worker_id )
540- transaction . expire ( key ( ' master-worker-id' ) , config . redis_ttl )
541- end
485+ return true if @master
486+
487+ @master = redis . set ( key ( 'master-status' ) , 'setup' , nx : true , px : setup_timeout_milliseconds )
488+ if @master
489+ begin
490+ redis . set ( key ( 'master-worker-id' ) , worker_id )
491+ redis . expire ( key ( 'master-worker-id' ) , config . redis_ttl )
492+ warn "Worker #{ worker_id } elected as master"
493+ rescue * CONNECTION_ERRORS
494+ # If setting master-worker-id fails, we still have master status
495+ # Log but don't lose master role
496+ warn ( "Failed to set master-worker-id: #{ $! . message } " )
542497 end
543498 end
544- warn "Worker #{ worker_id } elected as master" if @master
545499 @master
546500 rescue *CONNECTION_ERRORS
547501 @master = nil
548502 false
549503 end
550504
505+ def setup_timeout_milliseconds
506+ # Expire early enough for a waiting worker to take over and still
507+ # publish within its own queue_init_timeout.
508+ [ ( config . queue_init_timeout . to_f * 500 ) . ceil , 1 ] . max
509+ end
510+
551511 def register_worker_presence
552512 register
553513 redis . expire ( key ( 'workers' ) , config . redis_ttl )
554514 rescue *CONNECTION_ERRORS
555515 raise if master?
556516 end
557517
558- def store_chunk_metadata ( chunks )
559- # Batch operations to avoid exceeding Redis multi operation limits
560- # Each chunk requires 4 commands (set, expire, sadd, hset), so batch conservatively
561- batch_size = 5 # 5 chunks = 20 commands + 2 expires = 22 commands per batch
562-
563- chunks . each_slice ( batch_size ) do |chunk_batch |
564- with_master_lease do |transaction |
565- chunk_batch . each do |chunk |
566- # Store chunk metadata with TTL
567- transaction . set (
568- key ( 'chunk' , chunk . id ) ,
569- chunk . to_json
570- )
571- transaction . expire ( key ( 'chunk' , chunk . id ) , config . redis_ttl )
572-
573- # Track all chunks for cleanup
574- transaction . sadd ( key ( 'chunks' ) , chunk . id )
575-
576- # Store dynamic timeout for this chunk
577- # Timeout = estimated_duration (in ms) converted to seconds + buffer
578- # estimated_duration is in milliseconds, convert to seconds and add 10% buffer
579- buffer_percent = 10
580- estimated_duration_seconds = chunk . estimated_duration / 1000.0
581- chunk_timeout = ( estimated_duration_seconds * ( 1 + buffer_percent / 100.0 ) ) . round ( 2 )
582- # Format to string to avoid floating point precision issues in Redis
583- # Use %g to remove trailing zeros
584- transaction . hset ( key ( 'test-group-timeout' ) , chunk . id , format ( '%g' , chunk_timeout ) )
585- end
586- transaction . expire ( key ( 'chunks' ) , config . redis_ttl )
587- transaction . expire ( key ( 'test-group-timeout' ) , config . redis_ttl )
588- end
518+ def store_chunk_metadata ( transaction , chunks )
519+ return if chunks . empty?
520+
521+ chunks . each do |chunk |
522+ # Store chunk metadata with TTL
523+ transaction . set (
524+ key ( 'chunk' , chunk . id ) ,
525+ chunk . to_json
526+ )
527+ transaction . expire ( key ( 'chunk' , chunk . id ) , config . redis_ttl )
528+
529+ # Track all chunks for cleanup
530+ transaction . sadd ( key ( 'chunks' ) , chunk . id )
531+
532+ # Store dynamic timeout for this chunk
533+ # Timeout = estimated_duration (in ms) converted to seconds + buffer
534+ # estimated_duration is in milliseconds, convert to seconds and add 10% buffer
535+ buffer_percent = 10
536+ estimated_duration_seconds = chunk . estimated_duration / 1000.0
537+ chunk_timeout = ( estimated_duration_seconds * ( 1 + buffer_percent / 100.0 ) ) . round ( 2 )
538+ # Format to string to avoid floating point precision issues in Redis
539+ # Use %g to remove trailing zeros
540+ transaction . hset ( key ( 'test-group-timeout' ) , chunk . id , format ( '%g' , chunk_timeout ) )
589541 end
542+ transaction . expire ( key ( 'chunks' ) , config . redis_ttl )
543+ transaction . expire ( key ( 'test-group-timeout' ) , config . redis_ttl )
590544 end
591545
592546 def chunk_id? ( id )
0 commit comments