Merge pull request #50 from figma/ebarajas/ruby-support-dynamic-deadline

ruby: add support for dynamic deadlines when running suites

Author: ebarajas

redis/reserve.lua

@@ -7,11 +7,17 @@ local test_group_timeout_key = KEYS[6]
 
 local current_time = ARGV[1]
 local use_dynamic_deadline = ARGV[2] == "true"
+local default_timeout = ARGV[3] or 0
 
 local test = redis.call('rpop', queue_key)
 if test then
   if use_dynamic_deadline then
     local dynamic_timeout = redis.call('hget', test_group_timeout_key, test)
+    if not dynamic_timeout then
+      dynamic_timeout = default_timeout
+    else
+      dynamic_timeout = tonumber(dynamic_timeout)
+    end
     redis.call('zadd', zset_key, current_time + dynamic_timeout, test)
   else
     redis.call('zadd', zset_key, current_time, test)

redis/reserve_lost.lua

@@ -4,9 +4,10 @@ local worker_queue_key = KEYS[3]
 local owners_key = KEYS[4]
 local test_group_timeout_key = KEYS[5]
 
-local current_time = ARGV[1]
-local timeout = ARGV[2]
+local current_time = tonumber(ARGV[1])
+local timeout = tonumber(ARGV[2])
 local use_dynamic_deadline = ARGV[3] == "true"
+local default_timeout = tonumber(ARGV[4]) or 0
 
 local lost_tests
 if use_dynamic_deadline then
@@ -19,9 +20,14 @@ for _, test in ipairs(lost_tests) do
   if redis.call('sismember', processed_key, test) == 0 then
     if use_dynamic_deadline then
       local dynamic_timeout = redis.call('hget', test_group_timeout_key, test)
+      if not dynamic_timeout or dynamic_timeout == "" then
+        dynamic_timeout = default_timeout
+      else
+        dynamic_timeout = tonumber(dynamic_timeout)
+      end
       redis.call('zadd', zset_key, current_time + dynamic_timeout, test)
     else
-      redis.call('zadd', zset_key, current_time, test)
+      redis.call('zadd', zset_key, current_time + timeout, test)
     end
     redis.call('lpush', worker_queue_key, test)
     redis.call('hset', owners_key, test, worker_queue_key) -- Take ownership

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

@@ -219,8 +219,9 @@ def try_to_reserve_test
               key('processed'),
               key('worker', worker_id, 'queue'),
               key('owners'),
+              key('test-group-timeout'),
             ],
-            argv: [CI::Queue.time_now.to_f],
+            argv: [CI::Queue.time_now.to_f, 'true', config.timeout],
           )
         end
 
@@ -232,8 +233,9 @@ def try_to_reserve_lost_test
               key('completed'),
               key('worker', worker_id, 'queue'),
               key('owners'),
+              key('test-group-timeout'),
             ],
-            argv: [CI::Queue.time_now.to_f, timeout],
+            argv: [CI::Queue.time_now.to_f, timeout, 'true', config.timeout],
           )
 
           if lost_test.nil? && idle?
@@ -277,8 +279,8 @@ def register
 
         def store_chunk_metadata(chunks)
           # Batch operations to avoid exceeding Redis multi operation limits
-          # Each chunk requires 3 commands (set, expire, sadd), so batch conservatively
-          batch_size = 7 # 7 chunks = 21 commands + 1 expire = 22 commands per batch
+          # 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|
             redis.multi do |transaction|
@@ -292,8 +294,14 @@ def store_chunk_metadata(chunks)
 
                 # Track all chunks for cleanup
                 transaction.sadd(key('chunks'), chunk.id)
+
+                # Store dynamic timeout for this chunk
+                # Timeout = default_timeout * number_of_tests
+                chunk_timeout = config.timeout * chunk.test_count
+                transaction.hset(key('test-group-timeout'), chunk.id, chunk_timeout)
               end
               transaction.expire(key('chunks'), config.redis_ttl)
+              transaction.expire(key('test-group-timeout'), config.redis_ttl)
             end
           end
         end

ruby/lib/ci/queue/strategy/suite_bin_packing.rb

@@ -62,8 +62,9 @@ def create_chunks_for_suite(suite_name, suite_tests, max_duration, buffer_percen
           # If suite fits in max duration, create full_suite chunk
           if total_duration <= max_duration
             chunk_id = "#{suite_name}:full_suite"
-            # Don't store test_ids - worker will resolve from index
-            return [TestChunk.new(chunk_id, suite_name, :full_suite, [], total_duration)]
+            # Don't store test_ids in Redis - worker will resolve from index
+            # But pass test_count for timeout calculation
+            return [TestChunk.new(chunk_id, suite_name, :full_suite, [], total_duration, test_count: suite_tests.size)]
           end
 
           # Suite too large - split into partial_suite chunks

ruby/lib/ci/queue/test_chunk.rb

@@ -5,15 +5,15 @@ module Queue
     class TestChunk
       include Comparable
 
-      attr_reader :id, :suite_name, :type, :test_ids, :estimated_duration
+      attr_reader :id, :suite_name, :type, :test_ids, :estimated_duration, :test_count
 
-      # type can be :full_suite or :partial_suite
-      def initialize(id, suite_name, type, test_ids, estimated_duration = 0)
+      def initialize(id, suite_name, type, test_ids, estimated_duration = 0, test_count: nil)
         @id = id
         @suite_name = suite_name
         @type = type
         @test_ids = test_ids.freeze
         @estimated_duration = estimated_duration
+        @test_count = test_count || test_ids.size
       end
 
       def full_suite?
@@ -38,7 +38,8 @@ def to_json(*args)
         data = {
           type: type.to_s,
           suite_name: suite_name,
-          estimated_duration: estimated_duration
+          estimated_duration: estimated_duration,
+          test_count: test_count
         }
 
         # Only include test_ids for partial suites
@@ -55,7 +56,8 @@ def self.from_json(chunk_id, json_string)
           data['suite_name'],
           data['type'].to_sym,
           data['test_ids'] || [], # Empty for full_suite
-          data['estimated_duration']
+          data['estimated_duration'],
+          test_count: data['test_count']
         )
       end
     end