[CI] heart beating ongoing chunk

Author: lihao-figma

redis/heartbeat.lua

@@ -3,15 +3,44 @@ local processed_key = KEYS[2]
 local owners_key = KEYS[3]
 local worker_queue_key = KEYS[4]
 
-local current_time = ARGV[1]
+local current_time = tonumber(ARGV[1])
 local test = ARGV[2]
 
 -- already processed, we do not need to bump the timestamp
 if redis.call('sismember', processed_key, test) == 1 then
   return false
 end
 
--- we're still the owner of the test, we can bump the timestamp
-if redis.call('hget', owners_key, test) == worker_queue_key then
-  return redis.call('zadd', zset_key, current_time, test)
+-- we're still the owner of the test, check if we need to extend the deadline
+local owner_value = redis.call('hget', owners_key, test)
+if owner_value then
+  -- Parse owner value: format is "worker_queue_key|heartbeat_timestamp"
+  local pipe_pos = string.find(owner_value, "|")
+  if not pipe_pos then
+    return false
+  end
+  local stored_worker_key = string.sub(owner_value, 1, pipe_pos - 1)
+  
+  if stored_worker_key == worker_queue_key then
+    -- Always update last heartbeat timestamp in owners hash
+    local new_owner_value = worker_queue_key .. "|" .. current_time
+    redis.call('hset', owners_key, test, new_owner_value)
+    
+    local deadline = redis.call('zscore', zset_key, test)
+    if deadline then
+      deadline = tonumber(deadline)
+      -- Only extend if deadline is within 1 minute of expiring
+      if deadline - 60 < current_time then
+        -- Extend by 2 minutes
+        local new_deadline = current_time + 120
+        redis.call('zadd', zset_key, new_deadline, test)
+        -- Return old deadline and new deadline
+        return {deadline, new_deadline}
+      end
+    end
+    -- No extension needed, but heartbeat was recorded
+    return 0
+  end
 end
+
+return false

redis/release.lua

@@ -2,13 +2,19 @@ local zset_key = KEYS[1]
 local worker_queue_key = KEYS[2]
 local owners_key = KEYS[3]
 
--- owned_tests = {"SomeTest", "worker:1", "SomeOtherTest", "worker:2", ...}
+-- owned_tests = {"SomeTest", "worker:1|1234567890", "SomeOtherTest", "worker:2|1234567891", ...}
 local owned_tests = redis.call('hgetall', owners_key)
 for index, owner_or_test in ipairs(owned_tests) do
-  if owner_or_test == worker_queue_key then -- If we owned a test
-    local test = owned_tests[index - 1]
-    redis.call('zadd', zset_key, "0", test) -- We expire the lease immediately
-    return nil
+  -- Parse owner value: format is "worker_queue_key|heartbeat_timestamp"
+  local pipe_pos = string.find(owner_or_test, "|")
+  if pipe_pos then
+    local stored_worker_key = string.sub(owner_or_test, 1, pipe_pos - 1)
+    
+    if stored_worker_key == worker_queue_key then -- If we owned a test
+      local test = owned_tests[index - 1]
+      redis.call('zadd', zset_key, "0", test) -- We expire the lease immediately
+      return nil
+    end
   end
 end
 

redis/requeue.lua

@@ -10,8 +10,17 @@ local global_max_requeues = tonumber(ARGV[2])
 local test = ARGV[3]
 local offset = ARGV[4]
 
-if redis.call('hget', owners_key, test) == worker_queue_key then
-   redis.call('hdel', owners_key, test)
+local owner_value = redis.call('hget', owners_key, test)
+if owner_value then
+  -- Parse owner value: format is "worker_queue_key|heartbeat_timestamp"
+  local pipe_pos = string.find(owner_value, "|")
+  if pipe_pos then
+    local stored_worker_key = string.sub(owner_value, 1, pipe_pos - 1)
+    
+    if stored_worker_key == worker_queue_key then
+      redis.call('hdel', owners_key, test)
+    end
+  end
 end
 
 if redis.call('sismember', processed_key, test) == 1 then

redis/reserve.lua

@@ -23,7 +23,9 @@ if test then
     redis.call('zadd', zset_key, current_time, test)
   end
   redis.call('lpush', worker_queue_key, test)
-  redis.call('hset', owners_key, test, worker_queue_key)
+  -- Store owner with initial heartbeat timestamp (current_time)
+  local owner_value = worker_queue_key .. "|" .. current_time
+  redis.call('hset', owners_key, test, owner_value)
   return test
 else
   return nil

redis/reserve_lost.lua

@@ -18,6 +18,24 @@ end
 
 for _, test in ipairs(lost_tests) do
   if redis.call('sismember', processed_key, test) == 0 then
+    -- Check if test has been heartbeated recently (within 2 minutes)
+    local owner_value = redis.call('hget', owners_key, test)
+    if owner_value then
+      -- Parse owner value: format is "worker_queue_key|heartbeat_timestamp"
+      local pipe_pos = string.find(owner_value, "|")
+      if pipe_pos then
+        local last_heartbeat_str = string.sub(owner_value, pipe_pos + 1)
+        local last_heartbeat = tonumber(last_heartbeat_str)
+        if last_heartbeat then
+          local heartbeat_age = current_time - last_heartbeat
+          -- If heartbeat is less than 2 minutes old, skip this test
+          if heartbeat_age < 120 then
+            goto continue
+          end
+        end
+      end
+    end
+    
     if use_dynamic_deadline then
       local dynamic_timeout = redis.call('hget', test_group_timeout_key, test)
       if not dynamic_timeout or dynamic_timeout == "" then
@@ -30,8 +48,11 @@ for _, test in ipairs(lost_tests) do
       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
+    -- Store owner with initial heartbeat timestamp (current_time)
+    local new_owner_value = worker_queue_key .. "|" .. current_time
+    redis.call('hset', owners_key, test, new_owner_value) -- Take ownership
     return test
+    ::continue::
   end
 end
 

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

@@ -122,6 +122,14 @@ def load_script(script)
           @scripts_cache[script] ||= redis.script(:load, read_script(script))
         end
 
+        def ensure_connection_and_script(script)
+          # Pre-initialize Redis connection and script in current thread context
+          # This ensures background threads use the same initialized connection
+          load_script(script)
+          # Ping Redis to ensure connection is established
+          redis.ping
+        end
+
         def read_script(name)
           ::File.read(::File.join(CI::Queue::DEV_SCRIPTS_ROOT, "#{name}.lua"))
         rescue SystemCallError

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

@@ -87,7 +87,32 @@ def poll
               executable = resolve_executable(id)
 
               if executable
-                yield executable
+                # Pre-initialize Redis connection and heartbeat script in main thread
+                # to ensure background thread uses the same initialized context
+                ensure_connection_and_script(:heartbeat)
+
+                # Start heartbeat thread to extend lease while executing
+                heartbeat_thread = nil
+                begin
+                  heartbeat_thread = Thread.new do
+                    heartbeat_interval = 10 # Send heartbeat every 10 seconds
+                    loop do
+                      sleep heartbeat_interval
+                      break if Thread.current[:stop]
+
+                      heartbeat(id)
+                    end
+                  end
+                  heartbeat_thread[:stop] = false
+
+                  yield executable
+                ensure
+                  # Stop heartbeat thread when execution completes
+                  if heartbeat_thread
+                    heartbeat_thread[:stop] = true
+                    heartbeat_thread.join(1) # Wait up to 1 second for thread to finish
+                  end
+                end
               else
                 warn("Warning: Could not resolve executable for ID #{id.inspect}. Acknowledging to remove from queue.")
                 acknowledge(id)
@@ -209,6 +234,31 @@ def release!
           nil
         end
 
+        def heartbeat(test_id)
+          current_time = CI::Queue.time_now.to_f
+          result = eval_script(
+            :heartbeat,
+            keys: [
+              key('running'),
+              key('processed'),
+              key('owners'),
+              key('worker', worker_id, 'queue')
+            ],
+            argv: [current_time, test_id]
+          )
+          if result.is_a?(Array) && result.size == 2
+            old_deadline = result[0]
+            new_deadline = result[1]
+            old_deadline_readable = Time.at(old_deadline).strftime('%Y-%m-%d %H:%M:%S')
+            new_deadline_readable = Time.at(new_deadline).strftime('%Y-%m-%d %H:%M:%S')
+            warn("[heartbeat] Extended deadline for #{test_id.inspect} from #{old_deadline_readable} (#{old_deadline}) to #{new_deadline_readable} (#{new_deadline})")
+          end
+          result
+        rescue *CONNECTION_ERRORS => e
+          warn("Failed to send heartbeat for #{test_id.inspect}: #{e.class} - #{e.message}")
+          false
+        end
+
         private
 
         attr_reader :index