Remove the per lock duration (duration same for all locks in a single RunLock instance)

Author: ericallam

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

@@ -98,7 +98,7 @@ export class RunEngine {
       logger: this.logger,
       tracer: trace.getTracer("RunLocker"),
       meter: options.meter,
-      defaultDuration: options.runLock.duration ?? 5000,
+      duration: options.runLock.duration ?? 5000,
       automaticExtensionThreshold: options.runLock.automaticExtensionThreshold ?? 1000,
       retryConfig: {
         maxRetries: 10,
@@ -1162,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(

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

@@ -66,15 +66,6 @@ export interface LockRetryConfig {
   maxTotalWaitTime?: number;
 }
 
-interface LockOptions {
-  /** Default lock duration in milliseconds (default: 5000) */
-  defaultDuration?: number;
-  /** Automatic extension threshold in milliseconds - how early to extend locks before expiration (default: 500) */
-  automaticExtensionThreshold?: number;
-  /** Retry configuration for lock acquisition */
-  retryConfig?: LockRetryConfig;
-}
-
 export class RunLocker {
   private redlock: InstanceType<typeof redlock.default>;
   private asyncLocalStorage: AsyncLocalStorage<LockContext>;
@@ -85,20 +76,20 @@ export class RunLocker {
   private activeManualContexts: Map<string, ManualLockContext> = new Map();
   private lockDurationHistogram: Histogram;
   private retryConfig: Required<LockRetryConfig>;
-  private defaultDuration: number;
+  private duration: number;
   private automaticExtensionThreshold: number;
 
   constructor(options: {
     redis: Redis;
     logger: Logger;
     tracer: Tracer;
     meter?: Meter;
-    defaultDuration?: number;
+    duration?: number;
     automaticExtensionThreshold?: number;
     retryConfig?: LockRetryConfig;
   }) {
     // Initialize configuration values
-    this.defaultDuration = options.defaultDuration ?? 5000;
+    this.duration = options.duration ?? 5000;
     this.automaticExtensionThreshold = options.automaticExtensionThreshold ?? 500;
 
     this.redlock = new Redlock([options.redis], {
@@ -157,48 +148,18 @@ export class RunLocker {
   }
 
   /** Locks resources using RedLock. It won't lock again if we're already inside a lock with the same resources. */
-  async lock<T>(
-    name: string,
-    resources: string[],
-    duration: number | undefined,
-    routine: (signal: redlock.RedlockAbortSignal) => Promise<T>
-  ): Promise<T>;
-  async lock<T>(
-    name: string,
-    resources: string[],
-    routine: (signal: redlock.RedlockAbortSignal) => Promise<T>
-  ): Promise<T>;
-  async lock<T>(
-    name: string,
-    resources: string[],
-    durationOrRoutine: number | undefined | ((signal: redlock.RedlockAbortSignal) => Promise<T>),
-    routine?: (signal: redlock.RedlockAbortSignal) => Promise<T>
-  ): Promise<T> {
+  async lock<T>(name: string, resources: string[], routine: () => Promise<T>): Promise<T> {
     const currentContext = this.asyncLocalStorage.getStore();
     const joinedResources = [...resources].sort().join(",");
 
-    // Handle overloaded parameters
-    let actualDuration: number;
-    let actualRoutine: (signal: redlock.RedlockAbortSignal) => Promise<T>;
-
-    if (typeof durationOrRoutine === "function") {
-      // Called as lock(name, resources, routine) - use default duration
-      actualDuration = this.defaultDuration;
-      actualRoutine = durationOrRoutine;
-    } else {
-      // Called as lock(name, resources, duration, routine) - use provided duration
-      actualDuration = durationOrRoutine ?? this.defaultDuration;
-      actualRoutine = routine!;
-    }
-
     return startSpan(
       this.tracer,
       "RunLocker.lock",
       async (span) => {
         if (currentContext && currentContext.resources === joinedResources) {
           span.setAttribute("nested", true);
           // We're already inside a lock with the same resources, just run the routine
-          return actualRoutine(currentContext.signal);
+          return routine();
         }
 
         span.setAttribute("nested", false);
@@ -208,14 +169,7 @@ export class RunLocker {
         const lockStartTime = performance.now();
 
         const [error, result] = await tryCatch(
-          this.#acquireAndExecute(
-            name,
-            resources,
-            actualDuration,
-            actualRoutine,
-            lockId,
-            lockStartTime
-          )
+          this.#acquireAndExecute(name, resources, this.duration, routine, lockId, lockStartTime)
         );
 
         if (error) {
@@ -229,15 +183,15 @@ export class RunLocker {
           this.logger.error("[RunLocker] Error locking resources", {
             error,
             resources,
-            duration: actualDuration,
+            duration: this.duration,
           });
           throw error;
         }
 
         return result;
       },
       {
-        attributes: { name, resources, timeout: actualDuration },
+        attributes: { name, resources, timeout: this.duration },
       }
     );
   }
@@ -247,7 +201,7 @@ export class RunLocker {
     name: string,
     resources: string[],
     duration: number,
-    routine: (signal: redlock.RedlockAbortSignal) => Promise<T>,
+    routine: () => Promise<T>,
     lockId: string,
     lockStartTime: number
   ): Promise<T> {
@@ -260,7 +214,6 @@ export class RunLocker {
       this.retryConfig;
 
     // Track timing for total wait time limit
-    const retryStartTime = performance.now();
     let totalWaitTime = 0;
 
     // Retry the lock acquisition with exponential backoff
@@ -398,7 +351,7 @@ export class RunLocker {
       let lockSuccess = true;
       try {
         const result = await this.asyncLocalStorage.run(newContext, async () => {
-          return routine(signal);
+          return routine();
         });
 
         return result;
@@ -529,47 +482,12 @@ export class RunLocker {
     condition: boolean,
     name: string,
     resources: string[],
-    duration: number | undefined,
-    routine: (signal?: redlock.RedlockAbortSignal) => Promise<T>
-  ): Promise<T>;
-  async lockIf<T>(
-    condition: boolean,
-    name: string,
-    resources: string[],
-    routine: (signal?: redlock.RedlockAbortSignal) => Promise<T>
-  ): Promise<T>;
-  async lockIf<T>(
-    condition: boolean,
-    name: string,
-    resources: string[],
-    durationOrRoutine: number | undefined | ((signal?: redlock.RedlockAbortSignal) => Promise<T>),
-    routine?: (signal?: redlock.RedlockAbortSignal) => Promise<T>
+    routine: () => Promise<T>
   ): Promise<T> {
     if (condition) {
-      // Handle overloaded parameters
-      if (typeof durationOrRoutine === "function") {
-        // Called as lockIf(condition, name, resources, routine) - use default duration
-        return this.lock(
-          name,
-          resources,
-          durationOrRoutine as (signal: redlock.RedlockAbortSignal) => Promise<T>
-        );
-      } else {
-        // Called as lockIf(condition, name, resources, duration, routine) - use provided duration
-        return this.lock(
-          name,
-          resources,
-          durationOrRoutine,
-          routine! as (signal: redlock.RedlockAbortSignal) => Promise<T>
-        );
-      }
+      return this.lock(name, resources, routine);
     } else {
-      // Handle overloaded parameters for non-lock case
-      if (typeof durationOrRoutine === "function") {
-        return durationOrRoutine();
-      } else {
-        return routine!();
-      }
+      return routine();
     }
   }
 
@@ -585,8 +503,8 @@ export class RunLocker {
     return { ...this.retryConfig };
   }
 
-  getDefaultDuration(): number {
-    return this.defaultDuration;
+  getDuration(): number {
+    return this.duration;
   }
 
   getAutomaticExtensionThreshold(): number {

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

@@ -53,7 +53,7 @@ export class CheckpointSystem {
   }): Promise<CreateCheckpointResult> {
     const prisma = tx ?? this.$.prisma;
 
-    return await this.$.runLock.lock("createCheckpoint", [runId], 5_000, async () => {
+    return await this.$.runLock.lock("createCheckpoint", [runId], async () => {
       const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
       const isValidSnapshot =
@@ -267,7 +267,7 @@ export class CheckpointSystem {
   }): Promise<ExecutionResult> {
     const prisma = tx ?? this.$.prisma;
 
-    return await this.$.runLock.lock("continueRunExecution", [runId], 5_000, async () => {
+    return await this.$.runLock.lock("continueRunExecution", [runId], async () => {
       const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
       if (snapshot.id !== snapshotId) {

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

@@ -37,7 +37,7 @@ export class DelayedRunSystem {
       this.$.tracer,
       "rescheduleDelayedRun",
       async () => {
-        return await this.$.runLock.lock("rescheduleDelayedRun", [runId], 5_000, async () => {
+        return await this.$.runLock.lock("rescheduleDelayedRun", [runId], async () => {
           const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
           //if the run isn't just created then we can't reschedule it

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

@@ -80,8 +80,7 @@ export class DequeueSystem {
           const dequeuedRun = await this.$.runLock.lock(
             "dequeueFromWorkerQueue",
             [runId],
-            5000,
-            async (signal) => {
+            async () => {
               const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
               if (!isDequeueableExecutionStatus(snapshot.executionStatus)) {
@@ -548,7 +547,7 @@ export class DequeueSystem {
       statusReason,
     });
 
-    return this.$.runLock.lock("pendingVersion", [runId], 5_000, async (signal) => {
+    return this.$.runLock.lock("pendingVersion", [runId], async () => {
       this.$.logger.debug("RunEngine.dequeueFromWorkerQueue(): Pending version lock acquired", {
         runId,
         reason,

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

@@ -56,7 +56,7 @@ export class EnqueueSystem {
   }) {
     const prisma = tx ?? this.$.prisma;
 
-    return await this.$.runLock.lockIf(!skipRunLock, "enqueueRun", [run.id], 5000, async () => {
+    return await this.$.runLock.lockIf(!skipRunLock, "enqueueRun", [run.id], async () => {
       const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot(prisma, {
         run: run,
         snapshot: {

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

@@ -263,7 +263,6 @@ export class ReleaseConcurrencySystem {
     return await this.$.runLock.lock(
       "executeReleaseConcurrencyForSnapshot",
       [snapshot.runId],
-      5_000,
       async () => {
         const latestSnapshot = await getLatestExecutionSnapshot(this.$.prisma, snapshot.runId);
 

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

@@ -85,7 +85,7 @@ export class RunAttemptSystem {
       this.$.tracer,
       "startRunAttempt",
       async (span) => {
-        return this.$.runLock.lock("startRunAttempt", [runId], 5000, async () => {
+        return this.$.runLock.lock("startRunAttempt", [runId], async () => {
           const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
 
           if (latestSnapshot.id !== snapshotId) {
@@ -441,7 +441,7 @@ export class RunAttemptSystem {
       this.$.tracer,
       "#completeRunAttemptSuccess",
       async (span) => {
-        return this.$.runLock.lock("attemptSucceeded", [runId], 5_000, async (signal) => {
+        return this.$.runLock.lock("attemptSucceeded", [runId], async () => {
           const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
 
           if (latestSnapshot.id !== snapshotId) {
@@ -594,7 +594,7 @@ export class RunAttemptSystem {
       this.$.tracer,
       "completeRunAttemptFailure",
       async (span) => {
-        return this.$.runLock.lock("attemptFailed", [runId], 5_000, async (signal) => {
+        return this.$.runLock.lock("attemptFailed", [runId], async () => {
           const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
 
           if (latestSnapshot.id !== snapshotId) {
@@ -905,7 +905,7 @@ export class RunAttemptSystem {
   }): Promise<{ wasRequeued: boolean } & ExecutionResult> {
     const prisma = tx ?? this.$.prisma;
 
-    return await this.$.runLock.lock("tryNackAndRequeue", [run.id], 5000, async (signal) => {
+    return await this.$.runLock.lock("tryNackAndRequeue", [run.id], async () => {
       //we nack the message, this allows another work to pick up the run
       const gotRequeued = await this.$.runQueue.nackMessage({
         orgId,
@@ -982,7 +982,7 @@ export class RunAttemptSystem {
     reason = reason ?? "Cancelled by user";
 
     return startSpan(this.$.tracer, "cancelRun", async (span) => {
-      return this.$.runLock.lock("cancelRun", [runId], 5_000, async (signal) => {
+      return this.$.runLock.lock("cancelRun", [runId], async () => {
         const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
 
         //already finished, do nothing

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

@@ -23,7 +23,7 @@ export class TtlSystem {
 
   async expireRun({ runId, tx }: { runId: string; tx?: PrismaClientOrTransaction }) {
     const prisma = tx ?? this.$.prisma;
-    await this.$.runLock.lock("expireRun", [runId], 5_000, async () => {
+    await this.$.runLock.lock("expireRun", [runId], async () => {
       const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
       //if we're executing then we won't expire the run

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

@@ -380,7 +380,7 @@ export class WaitpointSystem {
 
     let $waitpoints = typeof waitpoints === "string" ? [waitpoints] : waitpoints;
 
-    return await this.$.runLock.lock("blockRunWithWaitpoint", [runId], 5000, async () => {
+    return await this.$.runLock.lock("blockRunWithWaitpoint", [runId], async () => {
       let snapshot: TaskRunExecutionSnapshot = await getLatestExecutionSnapshot(prisma, runId);
 
       //block the run with the waitpoints, returning how many waitpoints are pending
@@ -549,7 +549,7 @@ export class WaitpointSystem {
     }
 
     //4. Continue the run whether it's executing or not
-    await this.$.runLock.lock("continueRunIfUnblocked", [runId], 5000, async () => {
+    await this.$.runLock.lock("continueRunIfUnblocked", [runId], async () => {
       const snapshot = await getLatestExecutionSnapshot(this.$.prisma, runId);
 
       if (isFinishedOrPendingFinished(snapshot.executionStatus)) {