v4: Improved run locking (#2173)

Author: ericallam

apps/webapp/app/env.server.ts

@@ -429,6 +429,15 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_DEQUEUE_BLOCKING_TIMEOUT_SECONDS: z.coerce.number().int().default(10),
   RUN_ENGINE_MASTER_QUEUE_CONSUMERS_INTERVAL_MS: z.coerce.number().int().default(500),
 
+  RUN_ENGINE_RUN_LOCK_DURATION: z.coerce.number().int().default(5000),
+  RUN_ENGINE_RUN_LOCK_AUTOMATIC_EXTENSION_THRESHOLD: z.coerce.number().int().default(1000),
+  RUN_ENGINE_RUN_LOCK_MAX_RETRIES: z.coerce.number().int().default(10),
+  RUN_ENGINE_RUN_LOCK_BASE_DELAY: z.coerce.number().int().default(100),
+  RUN_ENGINE_RUN_LOCK_MAX_DELAY: z.coerce.number().int().default(3000),
+  RUN_ENGINE_RUN_LOCK_BACKOFF_MULTIPLIER: z.coerce.number().default(1.8),
+  RUN_ENGINE_RUN_LOCK_JITTER_FACTOR: z.coerce.number().default(0.15),
+  RUN_ENGINE_RUN_LOCK_MAX_TOTAL_WAIT_TIME: z.coerce.number().int().default(15000),
+
   RUN_ENGINE_WORKER_REDIS_HOST: z
     .string()
     .optional()

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

@@ -75,6 +75,16 @@ function createRunEngine() {
         enableAutoPipelining: true,
         ...(env.RUN_ENGINE_RUN_LOCK_REDIS_TLS_DISABLED === "true" ? {} : { tls: {} }),
       },
+      duration: env.RUN_ENGINE_RUN_LOCK_DURATION,
+      automaticExtensionThreshold: env.RUN_ENGINE_RUN_LOCK_AUTOMATIC_EXTENSION_THRESHOLD,
+      retryConfig: {
+        maxAttempts: env.RUN_ENGINE_RUN_LOCK_MAX_RETRIES,
+        baseDelay: env.RUN_ENGINE_RUN_LOCK_BASE_DELAY,
+        maxDelay: env.RUN_ENGINE_RUN_LOCK_MAX_DELAY,
+        backoffMultiplier: env.RUN_ENGINE_RUN_LOCK_BACKOFF_MULTIPLIER,
+        jitterFactor: env.RUN_ENGINE_RUN_LOCK_JITTER_FACTOR,
+        maxTotalWaitTime: env.RUN_ENGINE_RUN_LOCK_MAX_TOTAL_WAIT_TIME,
+      },
     },
     tracer,
     meter,

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

@@ -98,6 +98,17 @@ export class RunEngine {
       logger: this.logger,
       tracer: trace.getTracer("RunLocker"),
       meter: options.meter,
+      duration: options.runLock.duration ?? 5000,
+      automaticExtensionThreshold: options.runLock.automaticExtensionThreshold ?? 1000,
+      retryConfig: {
+        maxAttempts: 10,
+        baseDelay: 100,
+        maxDelay: 3000,
+        backoffMultiplier: 1.8,
+        jitterFactor: 0.15,
+        maxTotalWaitTime: 15000,
+        ...options.runLock.retryConfig,
+      },
     });
 
     const keys = new RunQueueFullKeyProducer();
@@ -486,56 +497,52 @@ export class RunEngine {
 
         span.setAttribute("runId", taskRun.id);
 
-        await this.runLock.lock("trigger", [taskRun.id], 5000, async (signal) => {
-          //create associated waitpoint (this completes when the run completes)
-          const associatedWaitpoint = await this.waitpointSystem.createRunAssociatedWaitpoint(
-            prisma,
-            {
-              projectId: environment.project.id,
-              environmentId: environment.id,
-              completedByTaskRunId: taskRun.id,
-            }
-          );
-
-          //triggerAndWait or batchTriggerAndWait
-          if (resumeParentOnCompletion && parentTaskRunId) {
-            //this will block the parent run from continuing until this waitpoint is completed (and removed)
-            await this.waitpointSystem.blockRunWithWaitpoint({
-              runId: parentTaskRunId,
-              waitpoints: associatedWaitpoint.id,
-              projectId: associatedWaitpoint.projectId,
-              organizationId: environment.organization.id,
-              batch,
-              workerId,
-              runnerId,
-              tx: prisma,
-              releaseConcurrency,
-            });
+        //create associated waitpoint (this completes when the run completes)
+        const associatedWaitpoint = await this.waitpointSystem.createRunAssociatedWaitpoint(
+          prisma,
+          {
+            projectId: environment.project.id,
+            environmentId: environment.id,
+            completedByTaskRunId: taskRun.id,
           }
+        );
 
-          //Make sure lock extension succeeded
-          signal.throwIfAborted();
-
-          if (taskRun.delayUntil) {
-            // Schedule the run to be enqueued at the delayUntil time
-            await this.delayedRunSystem.scheduleDelayedRunEnqueuing({
-              runId: taskRun.id,
-              delayUntil: taskRun.delayUntil,
-            });
-          } else {
-            await this.enqueueSystem.enqueueRun({
-              run: taskRun,
-              env: environment,
-              workerId,
-              runnerId,
-              tx: prisma,
-            });
+        //triggerAndWait or batchTriggerAndWait
+        if (resumeParentOnCompletion && parentTaskRunId) {
+          //this will block the parent run from continuing until this waitpoint is completed (and removed)
+          await this.waitpointSystem.blockRunWithWaitpoint({
+            runId: parentTaskRunId,
+            waitpoints: associatedWaitpoint.id,
+            projectId: associatedWaitpoint.projectId,
+            organizationId: environment.organization.id,
+            batch,
+            workerId,
+            runnerId,
+            tx: prisma,
+            releaseConcurrency,
+          });
+        }
 
-            if (taskRun.ttl) {
-              await this.ttlSystem.scheduleExpireRun({ runId: taskRun.id, ttl: taskRun.ttl });
-            }
+        if (taskRun.delayUntil) {
+          // Schedule the run to be enqueued at the delayUntil time
+          await this.delayedRunSystem.scheduleDelayedRunEnqueuing({
+            runId: taskRun.id,
+            delayUntil: taskRun.delayUntil,
+          });
+        } else {
+          if (taskRun.ttl) {
+            await this.ttlSystem.scheduleExpireRun({ runId: taskRun.id, ttl: taskRun.ttl });
           }
-        });
+
+          await this.enqueueSystem.enqueueRun({
+            run: taskRun,
+            env: environment,
+            workerId,
+            runnerId,
+            tx: prisma,
+            skipRunLock: true,
+          });
+        }
 
         this.eventBus.emit("runCreated", {
           time: new Date(),
@@ -1155,7 +1162,7 @@ export class RunEngine {
     tx?: PrismaClientOrTransaction;
   }) {
     const prisma = tx ?? this.prisma;
-    return await this.runLock.lock("handleStalledSnapshot", [runId], 5_000, async () => {
+    return await this.runLock.lock("handleStalledSnapshot", [runId], async () => {
       const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
       if (latestSnapshot.id !== snapshotId) {
         this.logger.log(