Create a more reliable ttl expiration system using atomic redis worker

Author: ericallam

internal-packages/run-engine/src/engine/index.ts

@@ -75,13 +75,15 @@ import {
   RunEngineOptions,
   TriggerParams,
 } from "./types.js";
+import { ttlWorkerCatalog } from "./ttlWorkerCatalog.js";
 import { workerCatalog } from "./workerCatalog.js";
 import pMap from "p-map";
 
 export class RunEngine {
   private runLockRedis: Redis;
   private runLock: RunLocker;
   private worker: EngineWorker;
+  private ttlWorker: Worker<typeof ttlWorkerCatalog>;
   private logger: Logger;
   private tracer: Tracer;
   private meter: Meter;
@@ -193,7 +195,9 @@ export class RunEngine {
           shardCount: options.queue?.ttlSystem?.shardCount,
           pollIntervalMs: options.queue?.ttlSystem?.pollIntervalMs,
           batchSize: options.queue?.ttlSystem?.batchSize,
-          callback: this.#ttlExpiredCallback.bind(this),
+          workerQueueSuffix: "ttl-worker:{queue:ttl-expiration:}queue",
+          workerItemsSuffix: "ttl-worker:{queue:ttl-expiration:}items",
+          visibilityTimeoutMs: options.queue?.ttlSystem?.visibilityTimeoutMs ?? 30_000,
         },
     });
 
@@ -337,6 +341,31 @@ export class RunEngine {
       waitpointSystem: this.waitpointSystem,
     });
 
+    this.ttlWorker = new Worker({
+      name: "ttl-expiration",
+      redisOptions: {
+        ...options.queue.redis,
+        keyPrefix: `${options.queue.redis.keyPrefix}runqueue:ttl-worker:`,
+      },
+      catalog: ttlWorkerCatalog,
+      concurrency: { limit: 20 },
+      pollIntervalMs: options.worker.pollIntervalMs ?? 1000,
+      immediatePollIntervalMs: options.worker.immediatePollIntervalMs ?? 100,
+      shutdownTimeoutMs: options.worker.shutdownTimeoutMs ?? 10_000,
+      logger: new Logger("RunEngineTtlWorker", options.logLevel ?? "info"),
+      jobs: {
+        expireTtlRun: async ({ payload }) => {
+          await this.ttlSystem.expireRunsBatch([payload.runId]);
+        },
+      },
+    });
+
+    // Start TTL worker whenever TTL system is enabled, so expired runs enqueued by the
+    // Lua script get processed even when the main engine worker is disabled (e.g. in tests).
+    if (options.queue?.ttlSystem && !options.queue.ttlSystem.disabled) {
+      this.ttlWorker.start();
+    }
+
     this.batchSystem = new BatchSystem({
       resources,
       waitpointSystem: this.waitpointSystem,
@@ -1621,6 +1650,7 @@ export class RunEngine {
       //stop the run queue
       await this.runQueue.quit();
       await this.worker.stop();
+      await this.ttlWorker.stop();
       await this.runLock.quit();
 
       // This is just a failsafe
@@ -2229,41 +2259,6 @@ export class RunEngine {
     });
   }
 
-  /**
-   * Callback for the TTL system when runs expire.
-   * Uses the optimized batch method that doesn't require run locks
-   * since the Lua script already atomically claimed these runs.
-   */
-  async #ttlExpiredCallback(
-    runs: Array<{ queueKey: string; runId: string; orgId: string }>
-  ): Promise<void> {
-    if (runs.length === 0) return;
-
-    try {
-      const runIds = runs.map((r) => r.runId);
-      const result = await this.ttlSystem.expireRunsBatch(runIds);
-
-      if (result.expired.length > 0) {
-        this.logger.debug("TTL system expired runs", {
-          expiredCount: result.expired.length,
-          expiredRunIds: result.expired,
-        });
-      }
-
-      if (result.skipped.length > 0) {
-        this.logger.debug("TTL system skipped runs", {
-          skippedCount: result.skipped.length,
-          skipped: result.skipped,
-        });
-      }
-    } catch (error) {
-      this.logger.error("Failed to expire runs via TTL system", {
-        runIds: runs.map((r) => r.runId),
-        error,
-      });
-    }
-  }
-
   /**
    * Applies `defaultMaxTtl` to a run's TTL:
    * - No max configured → pass through as-is.

internal-packages/run-engine/src/engine/ttlWorkerCatalog.ts

@@ -0,0 +1,12 @@
+import { z } from "zod";
+
+export const ttlWorkerCatalog = {
+  expireTtlRun: {
+    schema: z.object({
+      runId: z.string(),
+      orgId: z.string(),
+      queueKey: z.string(),
+    }),
+    visibilityTimeoutMs: 30_000,
+  },
+};

internal-packages/run-engine/src/engine/types.ts

@@ -73,6 +73,8 @@ export type RunEngineOptions = {
       batchSize?: number;
       /** Whether TTL consumers are disabled (default: false) */
       disabled?: boolean;
+      /** Visibility timeout for TTL worker jobs (ms, default: 30000) */
+      visibilityTimeoutMs?: number;
     };
   };
   runLock: {

internal-packages/run-engine/src/run-queue/index.ts

@@ -100,15 +100,15 @@ export type RunQueueOptions = {
     pollIntervalMs?: number;
     /** Max number of runs to expire per poll per shard (default: 100) */
     batchSize?: number;
-    /** Callback to handle expired runs */
-    callback: TtlSystemCallback;
+    /** Key suffix for TTL worker's queue sorted set (relative to RunQueue keyPrefix) */
+    workerQueueSuffix: string;
+    /** Key suffix for TTL worker's items hash (relative to RunQueue keyPrefix) */
+    workerItemsSuffix: string;
+    /** Visibility timeout for TTL worker jobs (ms, default: 30000) */
+    visibilityTimeoutMs?: number;
   };
 };
 
-export interface TtlSystemCallback {
-  (runs: Array<{ queueKey: string; runId: string; orgId: string }>): Promise<void>;
-}
-
 export interface ConcurrencySweeperCallback {
   (runIds: string[]): Promise<Array<{ id: string; orgId: string }>>;
 }
@@ -1289,19 +1289,7 @@ export class RunQueue {
             shard,
             count: expiredRuns.length,
           });
-
-          // Call the callback with expired runs
-          try {
-            await this.options.ttlSystem!.callback(expiredRuns);
-            processedCount += expiredRuns.length;
-          } catch (callbackError) {
-            this.logger.error(`TTL callback failed for shard ${shard}`, {
-              error: callbackError,
-              service: this.name,
-              shard,
-              runCount: expiredRuns.length,
-            });
-          }
+          processedCount += expiredRuns.length;
         }
       }
     } catch (error) {
@@ -1318,24 +1306,36 @@ export class RunQueue {
   }
 
   /**
-   * Atomically expire TTL runs: removes from TTL set AND acknowledges from normal queue.
-   * This prevents race conditions with the normal dequeue system.
+   * Atomically expire TTL runs: removes from TTL set, acknowledges from normal queue,
+   * and enqueues each run to the TTL worker for DB updates.
    */
   async #expireTtlRuns(
     shard: number,
     now: number,
     batchSize: number
   ): Promise<Array<{ queueKey: string; runId: string; orgId: string }>> {
-    const shardCount = this.options.ttlSystem?.shardCount ?? this.shardCount;
+    const ttlSystem = this.options.ttlSystem;
+    if (!ttlSystem) {
+      return [];
+    }
+
+    const shardCount = ttlSystem.shardCount ?? this.shardCount;
     const ttlQueueKey = this.keys.ttlQueueKeyForShard(shard);
+    const keyPrefix = this.options.redis.keyPrefix ?? "";
+    const workerQueueKey = keyPrefix + ttlSystem.workerQueueSuffix;
+    const workerItemsKey = keyPrefix + ttlSystem.workerItemsSuffix;
+    const visibilityTimeoutMs = (ttlSystem.visibilityTimeoutMs ?? 30_000).toString();
 
-    // Atomically get and remove expired runs from TTL set, and ack them from normal queues
+    // Atomically get and remove expired runs from TTL set, ack them from normal queues, and enqueue to TTL worker
     const results = await this.redis.expireTtlRuns(
       ttlQueueKey,
-      this.options.redis.keyPrefix ?? "",
+      keyPrefix,
       now.toString(),
       batchSize.toString(),
-      shardCount.toString()
+      shardCount.toString(),
+      workerQueueKey,
+      workerItemsKey,
+      visibilityTimeoutMs
     );
 
     if (!results || results.length === 0) {
@@ -2587,7 +2587,7 @@ redis.call('SREM', envCurrentDequeuedKey, messageId)
       `,
     });
 
-    // Expire TTL runs - atomically removes from TTL set and acknowledges from normal queue
+    // Expire TTL runs - atomically removes from TTL set, acknowledges from normal queue, and enqueues to TTL worker
     this.redis.defineCommand("expireTtlRuns", {
       numberOfKeys: 1,
       lua: `
@@ -2596,6 +2596,9 @@ local keyPrefix = ARGV[1]
 local currentTime = tonumber(ARGV[2])
 local batchSize = tonumber(ARGV[3])
 local shardCount = tonumber(ARGV[4])
+local workerQueueKey = ARGV[5]
+local workerItemsKey = ARGV[6]
+local visibilityTimeoutMs = tonumber(ARGV[7])
 
 -- Get expired runs from TTL sorted set (score <= currentTime)
 local expiredMembers = redis.call('ZRANGEBYSCORE', ttlQueueKey, '-inf', currentTime, 'LIMIT', 0, batchSize)
@@ -2604,6 +2607,9 @@ if #expiredMembers == 0 then
   return {}
 end
 
+local time = redis.call('TIME')
+local nowMs = tonumber(time[1]) * 1000 + math.floor(tonumber(time[2]) / 1000)
+
 local results = {}
 
 for i, member in ipairs(expiredMembers) do
@@ -2656,6 +2662,16 @@ for i, member in ipairs(expiredMembers) do
       redis.call('SREM', envConcurrencyKey, runId)
       redis.call('SREM', envDequeuedKey, runId)
 
+      -- Enqueue to TTL worker (runId is natural dedup key)
+      local serializedItem = cjson.encode({
+        job = "expireTtlRun",
+        item = { runId = runId, orgId = orgId, queueKey = rawQueueKey },
+        visibilityTimeoutMs = visibilityTimeoutMs,
+        attempt = 0
+      })
+      redis.call('ZADD', workerQueueKey, nowMs, runId)
+      redis.call('HSET', workerItemsKey, runId, serializedItem)
+
       -- Add to results
       table.insert(results, member)
     end
@@ -3151,6 +3167,9 @@ declare module "@internal/redis" {
       currentTime: string,
       batchSize: string,
       shardCount: string,
+      workerQueueKey: string,
+      workerItemsKey: string,
+      visibilityTimeoutMs: string,
       callback?: Callback<string[]>
     ): Result<string[], Context>;