feat(run-engine): add sharded TTL system for automatic run expiration

Implements a scalable TTL system for the run queue: - Add ttlExpiresAt field to InputPayload for absolute expiration time - Add sharded TTL sorted sets (ttl:shard:N) for tracking runs by expiration - Modify enqueueMessage Lua script to add runs to TTL sorted set - Modify dequeueMessagesFromQueue to check TTL and skip expired runs - Add dequeueTtlExpiredRuns Lua script for TTL consumer batch processing - Add TTL encoding helpers (ttlEncoding.ts) for member serialization - Add TtlSystemOptions type for configuring TTL consumer behavior TTL design: - Both normal dequeue and TTL consumer paths use Lua for atomicity - Normal dequeue checks TTL, discards expired, returns expired list - TTL consumer batch dequeues expired runs and acks normal queue - Sharding distributes load across multiple sorted sets WIP: TTL consumer loops still need to be implemented https://claude.ai/code/session_01AyzQp6tbj7th5QRTCYjJR5
2026-02-02 01:06:13 +00:00
parent a50a5f5695
commit a27bbfb018
4 changed files with 302 additions and 17 deletions
@@ -42,7 +42,9 @@ import {
  RunQueueKeyProducer,
  RunQueueKeyProducerEnvironment,
  RunQueueSelectionStrategy,
+  TtlSystemOptions,
 } from "./types.js";
+import { encodeTtlMember } from "./ttlEncoding.js";
 import { WorkerQueueResolver } from "./workerQueueResolver.js";

 const SemanticAttributes = {
@@ -92,6 +94,12 @@ export type RunQueueOptions = {
    processMarkedJitterInMs?: number;
    callback: ConcurrencySweeperCallback;
  };
+  /**
+   * TTL system options for automatic run expiration.
+   * When enabled, runs with a ttlExpiresAt will be automatically expired
+   * if not dequeued before that time.
+   */
+  ttlSystem?: TtlSystemOptions;
 };

 export interface ConcurrencySweeperCallback {
@@ -1473,6 +1481,16 @@ export class RunQueue {
    const messageData = JSON.stringify(message);
    const messageScore = String(message.timestamp);

+    // TTL handling
+    const ttlExpiresAt = message.ttlExpiresAt ?? 0;
+    const ttlShardCount = this.options.ttlSystem?.shardCount ?? this.shardCount;
+    const ttlShardKey = ttlExpiresAt > 0
+      ? this.keys.ttlShardKey(this.keys.ttlShardForRun(message.runId, ttlShardCount))
+      : "";
+    const ttlMember = ttlExpiresAt > 0
+      ? encodeTtlMember({ runId: message.runId, queueKey: message.queue, orgId: message.orgId })
+      : "";
+
    this.logger.debug("Calling enqueueMessage", {
      queueKey,
      messageKey,
@@ -1486,6 +1504,8 @@ export class RunQueue {
      messageData,
      messageScore,
      masterQueueKey,
+      ttlExpiresAt,
+      ttlShardKey,
      service: this.name,
    });

@@ -1498,10 +1518,13 @@ export class RunQueue {
      queueCurrentDequeuedKey,
      envCurrentDequeuedKey,
      envQueueKey,
+      ttlShardKey,
      queueName,
      messageId,
      messageData,
-      messageScore
+      messageScore,
+      String(ttlExpiresAt),
+      ttlMember
    );
  }

@@ -2275,8 +2298,11 @@ end
      `,
    });

+    // TTL member encoding: runId|queueKey|orgId
+    const TTL_DELIMITER = "\\x1e";
+
    this.redis.defineCommand("enqueueMessage", {
-      numberOfKeys: 8,
+      numberOfKeys: 9, // Added ttlShardKey as optional 9th key
      lua: `
 local masterQueueKey = KEYS[1]
 local queueKey = KEYS[2]
@@ -2286,11 +2312,14 @@ local envCurrentConcurrencyKey = KEYS[5]
 local queueCurrentDequeuedKey = KEYS[6]
 local envCurrentDequeuedKey = KEYS[7]
 local envQueueKey = KEYS[8]
+local ttlShardKey = KEYS[9]  -- Optional: TTL sorted set shard key

 local queueName = ARGV[1]
 local messageId = ARGV[2]
 local messageData = ARGV[3]
 local messageScore = ARGV[4]
+local ttlExpiresAt = tonumber(ARGV[5] or '0')  -- Optional: TTL expiration timestamp
+local ttlMember = ARGV[6]  -- Optional: TTL member (runId|queueKey|orgId)

 -- Write the message to the message key
 redis.call('SET', messageKey, messageData)
@@ -2301,6 +2330,11 @@ redis.call('ZADD', queueKey, messageScore, messageId)
 -- Add the message to the env queue
 redis.call('ZADD', envQueueKey, messageScore, messageId)

+-- Add to TTL sorted set if TTL is set
+if ttlExpiresAt > 0 and ttlShardKey ~= '' and ttlMember ~= '' then
+  redis.call('ZADD', ttlShardKey, ttlExpiresAt, ttlMember)
+end
+
 -- Rebalance the parent queues
 local earliestMessage = redis.call('ZRANGE', queueKey, 0, 0, 'WITHSCORES')

@@ -2375,30 +2409,58 @@ if #messages == 0 then
 end

 local results = {}
+local expiredRuns = {}  -- Track expired runs for TTL cleanup
 local dequeuedCount = 0

 -- Process messages in pairs (messageId, score)
 for i = 1, #messages, 2 do
    local messageId = messages[i]
    local messageScore = tonumber(messages[i + 1])
-    
+
    -- Get the message payload
    local messageKey = messageKeyPrefix .. messageId
    local messagePayload = redis.call('GET', messageKey)
-    
+
    if messagePayload then
-        -- Update concurrency
-        redis.call('ZREM', queueKey, messageId)
-        redis.call('ZREM', envQueueKey, messageId)
-        redis.call('SADD', queueCurrentConcurrencyKey, messageId)
-        redis.call('SADD', envCurrentConcurrencyKey, messageId)
-        
-        -- Add to results
-        table.insert(results, messageId)
-        table.insert(results, messageScore)
-        table.insert(results, messagePayload)
-        
-        dequeuedCount = dequeuedCount + 1
+        -- Parse payload to check TTL
+        local ok, messageData = pcall(cjson.decode, messagePayload)
+        local isExpired = false
+
+        if ok and messageData and messageData.ttlExpiresAt then
+            local ttlExpiresAt = tonumber(messageData.ttlExpiresAt)
+            if ttlExpiresAt and ttlExpiresAt > 0 and currentTime >= ttlExpiresAt then
+                isExpired = true
+            end
+        end
+
+        if isExpired then
+            -- Run is expired - remove from queues but don't add to results
+            redis.call('ZREM', queueKey, messageId)
+            redis.call('ZREM', envQueueKey, messageId)
+            redis.call('DEL', messageKey)
+            -- Track for TTL cleanup callback (messageId:queueKey:orgId)
+            table.insert(expiredRuns, messageId)
+            if ok and messageData then
+                table.insert(expiredRuns, messageData.queue or '')
+                table.insert(expiredRuns, messageData.orgId or '')
+            else
+                table.insert(expiredRuns, '')
+                table.insert(expiredRuns, '')
+            end
+        else
+            -- Update concurrency
+            redis.call('ZREM', queueKey, messageId)
+            redis.call('ZREM', envQueueKey, messageId)
+            redis.call('SADD', queueCurrentConcurrencyKey, messageId)
+            redis.call('SADD', envCurrentConcurrencyKey, messageId)
+
+            -- Add to results
+            table.insert(results, messageId)
+            table.insert(results, messageScore)
+            table.insert(results, messagePayload)
+
+            dequeuedCount = dequeuedCount + 1
+        end
    end
 end

@@ -2411,7 +2473,17 @@ else
  redis.call('ZADD', masterQueueKey, earliestMessage[2], queueName)
 end

-- Return results as a flat array: [messageId1, messageScore1, messagePayload1, messageId2, messageScore2, messagePayload2, ...]
+-- Return results as a flat array: [messageId1, messageScore1, messagePayload1, ...]
+-- Followed by expired runs count and expired run data: [expiredCount, expiredRunId1, queueKey1, orgId1, ...]
+-- The caller can use expiredCount to parse the expired runs
+if #expiredRuns > 0 then
+    table.insert(results, '__EXPIRED__')
+    table.insert(results, #expiredRuns / 3)  -- Count of expired runs
+    for _, v in ipairs(expiredRuns) do
+        table.insert(results, v)
+    end
+end
+
 return results
      `,
    });
@@ -2436,6 +2508,71 @@ return {messageId, queueLength} -- Return message details
      `,
    });

+    // TTL consumer: batch dequeue expired runs from a TTL shard
+    this.redis.defineCommand("dequeueTtlExpiredRuns", {
+      numberOfKeys: 1, // Just the TTL shard key
+      lua: `
+local ttlShardKey = KEYS[1]
+local currentTime = tonumber(ARGV[1])
+local maxCount = tonumber(ARGV[2] or '100')
+local keyPrefix = ARGV[3] or ''
+
+-- TTL member delimiter
+local DELIMITER = "${TTL_DELIMITER}"
+
+-- Get expired runs (score <= currentTime)
+local expiredMembers = redis.call('ZRANGEBYSCORE', ttlShardKey, '-inf', currentTime, 'WITHSCORES', 'LIMIT', 0, maxCount)
+
+if #expiredMembers == 0 then
+    return nil
+end
+
+local results = {}
+
+-- Process members in pairs (member, score)
+for i = 1, #expiredMembers, 2 do
+    local member = expiredMembers[i]
+    local score = expiredMembers[i + 1]
+
+    -- Decode TTL member: runId|queueKey|orgId
+    local parts = {}
+    for part in string.gmatch(member, "([^" .. DELIMITER .. "]+)") do
+        table.insert(parts, part)
+    end
+
+    if #parts >= 3 then
+        local runId = parts[1]
+        local queueKey = parts[2]
+        local orgId = parts[3]
+
+        -- Remove from queue sorted set
+        redis.call('ZREM', queueKey, runId)
+
+        -- Remove from env sorted set (derive from queueKey)
+        -- queueKey format: {org:xxx}:proj:xxx:env:xxx:queue:xxx[:ck:xxx]
+        local envQueueKey = string.match(queueKey, "({org:[^}]+}):.-:(env:[^:]+)")
+        if envQueueKey then
+            redis.call('ZREM', envQueueKey, runId)
+        end
+
+        -- Delete message key
+        local messageKey = keyPrefix .. "{org:" .. orgId .. "}:message:" .. runId
+        redis.call('DEL', messageKey)
+
+        -- Remove from TTL sorted set
+        redis.call('ZREM', ttlShardKey, member)
+
+        -- Add to results: [runId, queueKey, orgId, ...]
+        table.insert(results, runId)
+        table.insert(results, queueKey)
+        table.insert(results, orgId)
+    end
+end
+
+return results
+      `,
+    });
+
    this.redis.defineCommand("dequeueMessageFromKey", {
      numberOfKeys: 1,
      lua: `
@@ -2740,11 +2877,14 @@ declare module "@internal/redis" {
      queueCurrentDequeuedKey: string,
      envCurrentDequeuedKey: string,
      envQueueKey: string,
+      ttlShardKey: string, // Optional: pass empty string if no TTL
      //args
      queueName: string,
      messageId: string,
      messageData: string,
      messageScore: string,
+      ttlExpiresAt: string, // Optional: pass '0' if no TTL
+      ttlMember: string, // Optional: pass empty string if no TTL
      callback?: Callback<void>
    ): Result<void, Context>;

@@ -2878,6 +3018,16 @@ declare module "@internal/redis" {
      maxCount: string,
      callback?: Callback<string[]>
    ): Result<string[], Context>;
+
+    dequeueTtlExpiredRuns(
+      // keys
+      ttlShardKey: string,
+      // args
+      currentTime: string,
+      maxCount: string,
+      keyPrefix: string,
+      callback?: Callback<string[] | null>
+    ): Result<string[] | null, Context>;
  }
 }

@@ -17,6 +17,7 @@ const constants = {
  DEAD_LETTER_QUEUE_PART: "deadLetter",
  MASTER_QUEUE_PART: "masterQueue",
  WORKER_QUEUE_PART: "workerQueue",
+  TTL_PART: "ttl",
 } as const;

 export class RunQueueFullKeyProducer implements RunQueueKeyProducer {
@@ -301,6 +302,22 @@ export class RunQueueFullKeyProducer implements RunQueueKeyProducer {
    return `*:${constants.ENV_PART}:*:queue:*:${constants.CURRENT_CONCURRENCY_PART}`;
  }

+  /**
+   * Get the TTL sorted set key for a specific shard.
+   * TTL sorted sets are sharded to distribute load.
+   */
+  ttlShardKey(shard: number): string {
+    return [constants.TTL_PART, "shard", shard.toString()].join(":");
+  }
+
+  /**
+   * Determine which TTL shard a run belongs to based on its runId.
+   * Uses jump consistent hash for even distribution.
+   */
+  ttlShardForRun(runId: string, shardCount: number): number {
+    return jumpHash(runId, shardCount);
+  }
+
  descriptorFromQueue(queue: string): QueueDescriptor {
    const parts = queue.split(":");
    return {
@@ -0,0 +1,86 @@
+/**
+ * TTL sorted set member encoding/decoding.
+ *
+ * TTL sorted set stores runs with their expiration time as the score.
+ * The member contains enough info to ack the normal queue when the run expires.
+ */
+
+// ASCII Record Separator - safe delimiter that won't appear in IDs
+const DELIMITER = "\x1e";
+
+export interface TtlMember {
+  runId: string;
+  queueKey: string;
+  orgId: string;
+}
+
+/**
+ * Encode TTL member for storage in sorted set.
+ * Format: runId␞queueKey␞orgId
+ */
+export function encodeTtlMember(data: TtlMember): string {
+  return [data.runId, data.queueKey, data.orgId].join(DELIMITER);
+}
+
+/**
+ * Decode TTL member from sorted set.
+ */
+export function decodeTtlMember(member: string): TtlMember | undefined {
+  const parts = member.split(DELIMITER);
+  if (parts.length !== 3) {
+    return undefined;
+  }
+
+  const [runId, queueKey, orgId] = parts;
+  return { runId, queueKey, orgId };
+}
+
+/**
+ * Check if a string is an encoded TTL member.
+ */
+export function isEncodedTtlMember(member: string): boolean {
+  return member.includes(DELIMITER);
+}
+
+/**
+ * Extract runId from TTL member (first part).
+ */
+export function getRunIdFromTtlMember(member: string): string {
+  const delimPos = member.indexOf(DELIMITER);
+  if (delimPos > 0) {
+    return member.substring(0, delimPos);
+  }
+  return member;
+}
+
+/**
+ * Lua helpers for TTL member encoding/decoding.
+ * Include in Lua scripts that need to work with TTL members.
+ */
+export const LUA_TTL_ENCODING_HELPERS = `
+-- TTL Member Encoding Helpers
+local TTL_DELIMITER = "\\x1e"
+
+local function encodeTtlMember(runId, queueKey, orgId)
+  return runId .. TTL_DELIMITER .. queueKey .. TTL_DELIMITER .. orgId
+end
+
+local function decodeTtlMember(member)
+  local parts = {}
+  for part in string.gmatch(member, "([^" .. TTL_DELIMITER .. "]+)") do
+    table.insert(parts, part)
+  end
+  if #parts ~= 3 then
+    return nil, nil, nil
+  end
+  return parts[1], parts[2], parts[3]  -- runId, queueKey, orgId
+end
+
+local function getRunIdFromTtlMember(member)
+  local delimPos = string.find(member, TTL_DELIMITER, 1, true)
+  if delimPos then
+    return string.sub(member, 1, delimPos - 1)
+  end
+  return member
+end
+`;
@@ -13,6 +13,11 @@ export const InputPayload = z.object({
  concurrencyKey: z.string().optional(),
  timestamp: z.number(),
  attempt: z.number(),
+  /**
+   * TTL expiration timestamp in milliseconds (absolute time when run expires).
+   * If set, the run will be automatically expired if not dequeued before this time.
+   */
+  ttlExpiresAt: z.number().optional(),
 });
 export type InputPayload = z.infer<typeof InputPayload>;

@@ -120,6 +125,10 @@ export interface RunQueueKeyProducer {
  // Concurrency sweeper methods
  markedForAckKey(): string;
  currentConcurrencySetKeyScanPattern(): string;
+
+  // TTL sorted set methods
+  ttlShardKey(shard: number): string;
+  ttlShardForRun(runId: string, shardCount: number): number;
 }

 export type EnvQueues = {
@@ -133,3 +142,26 @@ export interface RunQueueSelectionStrategy {
    consumerId: string
  ): Promise<Array<EnvQueues>>;
 }
+
+/**
+ * Callback invoked when runs are expired by the TTL system.
+ * Receives the queue key and run IDs that were expired.
+ * Should update the database and emit events as needed.
+ */
+export interface TtlExpiredCallback {
+  (runs: Array<{ queueKey: string; runId: string; orgId: string }>): Promise<void>;
+}
+
+/**
+ * TTL system options for the run queue.
+ */
+export type TtlSystemOptions = {
+  /** Number of shards for TTL sorted sets (default: same as queue shards) */
+  shardCount?: number;
+  /** How often to poll each shard for expired runs (ms) */
+  pollIntervalMs?: number;
+  /** Max number of runs to dequeue per poll per shard */
+  batchSize?: number;
+  /** Callback when runs are expired */
+  callback: TtlExpiredCallback;
+};