Added a span to the run lock with a name

Author: matt-aitken

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

@@ -92,7 +92,11 @@ export class RunEngine {
         },
       }
     );
-    this.runLock = new RunLocker({ redis: this.runLockRedis });
+    this.runLock = new RunLocker({
+      redis: this.runLockRedis,
+      logger: this.logger,
+      tracer: trace.getTracer("RunLocker"),
+    });
 
     const keys = new RunQueueFullKeyProducer();
 
@@ -491,7 +495,7 @@ export class RunEngine {
 
         span.setAttribute("runId", taskRun.id);
 
-        await this.runLock.lock([taskRun.id], 5000, async (signal) => {
+        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,
@@ -1162,7 +1166,7 @@ export class RunEngine {
     tx?: PrismaClientOrTransaction;
   }) {
     const prisma = tx ?? this.prisma;
-    return await this.runLock.lock([runId], 5_000, async () => {
+    return await this.runLock.lock("handleStalledSnapshot", [runId], 5_000, async () => {
       const latestSnapshot = await getLatestExecutionSnapshot(prisma, runId);
       if (latestSnapshot.id !== snapshotId) {
         this.logger.log(

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

@@ -3,6 +3,9 @@ const { default: Redlock } = require("redlock");
 import { AsyncLocalStorage } from "async_hooks";
 import { Redis } from "@internal/redis";
 import * as redlock from "redlock";
+import { tryCatch } from "@trigger.dev/core";
+import { Logger } from "@trigger.dev/core/logger";
+import { startSpan, Tracer } from "@internal/tracing";
 
 interface LockContext {
   resources: string;
@@ -12,8 +15,10 @@ interface LockContext {
 export class RunLocker {
   private redlock: InstanceType<typeof redlock.default>;
   private asyncLocalStorage: AsyncLocalStorage<LockContext>;
+  private logger: Logger;
+  private tracer: Tracer;
 
-  constructor(options: { redis: Redis }) {
+  constructor(options: { redis: Redis; logger: Logger; tracer: Tracer }) {
     this.redlock = new Redlock([options.redis], {
       driftFactor: 0.01,
       retryCount: 10,
@@ -22,30 +27,54 @@ export class RunLocker {
       automaticExtensionThreshold: 500, // time in ms
     });
     this.asyncLocalStorage = new AsyncLocalStorage<LockContext>();
+    this.logger = options.logger;
+    this.tracer = options.tracer;
   }
 
   /** 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,
     routine: (signal: redlock.RedlockAbortSignal) => Promise<T>
   ): Promise<T> {
     const currentContext = this.asyncLocalStorage.getStore();
     const joinedResources = resources.sort().join(",");
 
-    if (currentContext && currentContext.resources === joinedResources) {
-      // We're already inside a lock with the same resources, just run the routine
-      return routine(currentContext.signal);
-    }
+    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 routine(currentContext.signal);
+        }
 
-    // Different resources or not in a lock, proceed with new lock
-    return this.redlock.using(resources, duration, async (signal) => {
-      const newContext: LockContext = { resources: joinedResources, signal };
+        span.setAttribute("nested", false);
 
-      return this.asyncLocalStorage.run(newContext, async () => {
-        return routine(signal);
-      });
-    });
+        // Different resources or not in a lock, proceed with new lock
+        const [error, result] = await tryCatch(
+          this.redlock.using(resources, duration, async (signal) => {
+            const newContext: LockContext = { resources: joinedResources, signal };
+
+            return this.asyncLocalStorage.run(newContext, async () => {
+              return routine(signal);
+            });
+          })
+        );
+
+        if (error) {
+          this.logger.error("[RunLocker] Error locking resources", { error, resources, duration });
+          throw error;
+        }
+
+        return result;
+      },
+      {
+        attributes: { name, resources, timeout: duration },
+      }
+    );
   }
 
   isInsideLock(): boolean {

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([runId], 5_000, async () => {
+    return await this.$.runLock.lock("createCheckpoint", [runId], 5_000, async () => {
       const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
       const isValidSnapshot =
@@ -238,7 +238,7 @@ export class CheckpointSystem {
   }): Promise<ExecutionResult> {
     const prisma = tx ?? this.$.prisma;
 
-    return await this.$.runLock.lock([runId], 5_000, async () => {
+    return await this.$.runLock.lock("continueRunExecution", [runId], 5_000, 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([runId], 5_000, async () => {
+        return await this.$.runLock.lock("rescheduleDelayedRun", [runId], 5_000, async () => {
           const snapshot = await getLatestExecutionSnapshot(prisma, runId);
 
           //if the run isn't just created then we can't reschedule it