Implement batch support to redis worker and use it in the TTL worker to batch expire runs

Author: ericallam

apps/webapp/app/env.server.ts

@@ -604,6 +604,9 @@ const EnvironmentSchema = z
     RUN_ENGINE_TTL_SYSTEM_SHARD_COUNT: z.coerce.number().int().optional(),
     RUN_ENGINE_TTL_SYSTEM_POLL_INTERVAL_MS: z.coerce.number().int().default(1_000),
     RUN_ENGINE_TTL_SYSTEM_BATCH_SIZE: z.coerce.number().int().default(100),
+    RUN_ENGINE_TTL_WORKER_CONCURRENCY: z.coerce.number().int().default(1),
+    RUN_ENGINE_TTL_WORKER_BATCH_MAX_SIZE: z.coerce.number().int().default(50),
+    RUN_ENGINE_TTL_WORKER_BATCH_MAX_WAIT_MS: z.coerce.number().int().default(5_000),
 
     /** Optional maximum TTL for all runs (e.g. "14d"). If set, runs without an explicit TTL
      *  will use this as their TTL, and runs with a TTL larger than this will be clamped. */

apps/webapp/app/v3/runEngine.server.ts

@@ -85,6 +85,9 @@ function createRunEngine() {
         shardCount: env.RUN_ENGINE_TTL_SYSTEM_SHARD_COUNT,
         pollIntervalMs: env.RUN_ENGINE_TTL_SYSTEM_POLL_INTERVAL_MS,
         batchSize: env.RUN_ENGINE_TTL_SYSTEM_BATCH_SIZE,
+        workerConcurrency: env.RUN_ENGINE_TTL_WORKER_CONCURRENCY,
+        batchMaxSize: env.RUN_ENGINE_TTL_WORKER_BATCH_MAX_SIZE,
+        batchMaxWaitMs: env.RUN_ENGINE_TTL_WORKER_BATCH_MAX_WAIT_MS,
       },
     },
     runLock: {

apps/webapp/app/v3/services/batchTriggerTask.server.ts

@@ -75,15 +75,15 @@ import {
   RunEngineOptions,
   TriggerParams,
 } from "./types.js";
-import { ttlWorkerCatalog } from "./ttlWorkerCatalog.js";
+import { createTtlWorkerCatalog } 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 ttlWorker: Worker<ReturnType<typeof createTtlWorkerCatalog>>;
   private logger: Logger;
   private tracer: Tracer;
   private meter: Meter;
@@ -341,21 +341,27 @@ export class RunEngine {
       waitpointSystem: this.waitpointSystem,
     });
 
+    const ttlWorkerCatalog = createTtlWorkerCatalog({
+      visibilityTimeoutMs: options.queue?.ttlSystem?.visibilityTimeoutMs,
+      batchMaxSize: options.queue?.ttlSystem?.batchMaxSize,
+      batchMaxWaitMs: options.queue?.ttlSystem?.batchMaxWaitMs,
+    });
+
     this.ttlWorker = new Worker({
       name: "ttl-expiration",
       redisOptions: {
         ...options.queue.redis,
         keyPrefix: `${options.queue.redis.keyPrefix}runqueue:ttl-worker:`,
       },
       catalog: ttlWorkerCatalog,
-      concurrency: { limit: 20 },
+      concurrency: { limit: options.queue?.ttlSystem?.workerConcurrency ?? 1 },
       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]);
+        expireTtlRun: async (items) => {
+          await this.ttlSystem.expireRunsBatch(items.map((i) => i.payload.runId));
         },
       },
     });

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

@@ -171,7 +171,7 @@ export class TtlSystem {
         const skipped: { runId: string; reason: string }[] = [];
 
         // Fetch all runs in a single query (no snapshot data needed)
-        const runs = await this.$.prisma.taskRun.findMany({
+        const runs = await this.$.readOnlyPrisma.taskRun.findMany({
           where: { id: { in: runIds } },
           select: {
             id: true,
@@ -182,13 +182,9 @@ export class TtlSystem {
             taskEventStore: true,
             createdAt: true,
             associatedWaitpoint: { select: { id: true } },
-            runtimeEnvironment: {
-              select: {
-                id: true,
-                organizationId: true,
-                projectId: true,
-              },
-            },
+            organizationId: true,
+            projectId: true,
+            runtimeEnvironmentId: true,
           },
         });
 
@@ -259,6 +255,11 @@ export class TtlSystem {
                 });
               }
 
+              // This should really never happen
+              if (!run.organizationId) {
+                return;
+              }
+
               // Emit event
               this.$.eventBus.emit("runExpired", {
                 run: {
@@ -273,9 +274,9 @@ export class TtlSystem {
                   status: "EXPIRED" as TaskRunStatus,
                 },
                 time: now,
-                organization: { id: run.runtimeEnvironment.organizationId },
-                project: { id: run.runtimeEnvironment.projectId },
-                environment: { id: run.runtimeEnvironment.id },
+                organization: { id: run.organizationId },
+                project: { id: run.projectId },
+                environment: { id: run.runtimeEnvironmentId },
               });
 
               expired.push(run.id);

internal-packages/run-engine/src/engine/systems/ttlSystem.ts

@@ -733,8 +733,8 @@ describe("RunEngine ttl", () => {
         expect(executionData2.run.status).toBe("PENDING");
 
         // Now wait for the TTL consumer to poll and expire the run
-        // (pollIntervalMs is 5000, so we wait 7s to allow time for the poll + processing)
-        await setTimeout(7_000);
+        // (pollIntervalMs is 5000 for TTL scan + up to 5000ms batch maxWaitMs + processing)
+        await setTimeout(13_000);
 
         // The TTL consumer should have found and expired the run
         expect(expiredEvents.length).toBe(1);
@@ -848,7 +848,8 @@ describe("RunEngine ttl", () => {
           authenticatedEnvironment.id,
           10
         );
-        await setTimeout(7_000);
+        // Wait for TTL scan (5000ms) + batch maxWaitMs (5000ms) + processing buffer
+        await setTimeout(13_000);
 
         const expiredRun = await prisma.taskRun.findUnique({
           where: { id: run.id },

internal-packages/run-engine/src/engine/tests/ttl.test.ts

@@ -1,12 +1,26 @@
 import { z } from "zod";
 
-export const ttlWorkerCatalog = {
-  expireTtlRun: {
-    schema: z.object({
-      runId: z.string(),
-      orgId: z.string(),
-      queueKey: z.string(),
-    }),
-    visibilityTimeoutMs: 30_000,
-  },
+export type TtlWorkerCatalogOptions = {
+  visibilityTimeoutMs?: number;
+  batchMaxSize?: number;
+  batchMaxWaitMs?: number;
 };
+
+export function createTtlWorkerCatalog(options?: TtlWorkerCatalogOptions) {
+  return {
+    expireTtlRun: {
+      schema: z.object({
+        runId: z.string(),
+        orgId: z.string(),
+        queueKey: z.string(),
+      }),
+      visibilityTimeoutMs: options?.visibilityTimeoutMs ?? 120_000,
+      batch: {
+        maxSize: options?.batchMaxSize ?? 50,
+        maxWaitMs: options?.batchMaxWaitMs ?? 5_000,
+      },
+    },
+  };
+}
+
+export const ttlWorkerCatalog = createTtlWorkerCatalog();

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

@@ -73,8 +73,14 @@ export type RunEngineOptions = {
       batchSize?: number;
       /** Whether TTL consumers are disabled (default: false) */
       disabled?: boolean;
-      /** Visibility timeout for TTL worker jobs (ms, default: 30000) */
+      /** Visibility timeout for TTL worker jobs (ms, default: 120000) */
       visibilityTimeoutMs?: number;
+      /** Concurrency limit for the TTL redis-worker (default: 1) */
+      workerConcurrency?: number;
+      /** Max items to accumulate before flushing a batch (default: 500) */
+      batchMaxSize?: number;
+      /** Max time (ms) to wait for more items before flushing a batch (default: 5000) */
+      batchMaxWaitMs?: number;
     };
   };
   runLock: {