[CI] heart beating ongoing chunk

Author: lihao-figma

redis/heartbeat.lua

@@ -3,15 +3,30 @@ 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
+-- we're still the owner of the test, check if we need to extend the deadline
 if redis.call('hget', owners_key, test) == worker_queue_key then
-  return redis.call('zadd', zset_key, current_time, test)
+  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
+  return 0
 end
+
+return false

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,32 @@ 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 StandardError => e
+          warn("Failed to send heartbeat for #{test_id.inspect}: #{e.class} - #{e.message}")
+          warn(e.backtrace.join("\n")) if e.backtrace
+          false
+        end
+
         private
 
         attr_reader :index