Fixed incorrect span timings around checkpoints by implementing a precise wall clock that resets after restores

Author: ericallam

.env.example

@@ -63,4 +63,5 @@ COORDINATOR_SECRET=coordinator-secret # generate the actual secret with `openssl
 # These are needed for the object store (for handling large payloads/outputs)
 # OBJECT_STORE_BASE_URL="https://{bucket}.{accountId}.r2.cloudflarestorage.com"
 # OBJECT_STORE_ACCESS_KEY_ID=
-# OBJECT_STORE_SECRET_ACCESS_KEY=
+# OBJECT_STORE_SECRET_ACCESS_KEY=
+# RUNTIME_WAIT_THRESHOLD_IN_MS=10000

apps/webapp/app/env.server.ts

@@ -116,6 +116,7 @@ const EnvironmentSchema = z.object({
   DEV_OTEL_LOG_SCHEDULED_DELAY_MILLIS: z.string().default("200"),
   DEV_OTEL_LOG_EXPORT_TIMEOUT_MILLIS: z.string().default("30000"),
   DEV_OTEL_LOG_MAX_QUEUE_SIZE: z.string().default("512"),
+  RUNTIME_WAIT_THRESHOLD_IN_MS: z.coerce.number().int().default(30000),
 });
 
 export type Environment = z.infer<typeof EnvironmentSchema>;

apps/webapp/app/v3/environmentVariables/environmentVariablesRepository.server.ts

@@ -473,6 +473,10 @@ export class EnvironmentVariablesRepository implements Repository {
         key: "TRIGGER_API_URL",
         value: env.APP_ORIGIN,
       },
+      {
+        key: "TRIGGER_RUNTIME_WAIT_THRESHOLD_IN_MS",
+        value: String(env.RUNTIME_WAIT_THRESHOLD_IN_MS),
+      },
     ];
   }
 

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

@@ -591,14 +591,6 @@ export class MarQS {
     concurrencyLimitKey: string;
     currentConcurrencyKey: string;
   }) {
-    logger.debug("Calling dequeueMessage", {
-      messageQueue,
-      parentQueue,
-      visibilityQueue,
-      concurrencyLimitKey,
-      currentConcurrencyKey,
-    });
-
     const result = await this.redis.dequeueMessage(
       messageQueue,
       parentQueue,

packages/cli-v3/src/workers/prod/worker-facade.ts

@@ -201,7 +201,9 @@ const zodIpc = new ZodIpcConnection({
   },
 });
 
-const prodRuntimeManager = new ProdRuntimeManager(zodIpc);
+const prodRuntimeManager = new ProdRuntimeManager(zodIpc, {
+  waitThresholdInMs: parseInt(process.env.TRIGGER_RUNTIME_WAIT_THRESHOLD_IN_MS ?? "30000", 10),
+});
 
 runtime.setGlobalRuntimeManager(prodRuntimeManager);
 

packages/core/src/v3/clock/clock.ts

@@ -2,4 +2,5 @@ export type ClockTime = [number, number];
 
 export interface Clock {
   preciseNow(): ClockTime;
+  reset(): void;
 }

packages/core/src/v3/clock/durableClock.ts

@@ -27,6 +27,10 @@ export class ClockAPI {
     return this.#getClock().preciseNow();
   }
 
+  public reset(): void {
+    this.#getClock().reset();
+  }
+
   #getClock(): Clock {
     return getGlobal(API_NAME) ?? SIMPLE_CLOCK;
   }

packages/core/src/v3/clock/index.ts

@@ -0,0 +1,46 @@
+import { PreciseDate } from "@google-cloud/precise-date";
+import { Clock, ClockTime } from "./clock";
+
+export type PreciseWallClockOptions = {
+  origin?: ClockTime;
+  now?: PreciseDate;
+};
+
+export class PreciseWallClock implements Clock {
+  private _origin: {
+    clockTime: ClockTime;
+    preciseDate: PreciseDate;
+  };
+
+  get #originClockTime() {
+    return this._origin.clockTime;
+  }
+
+  get #originPreciseDate() {
+    return this._origin.preciseDate;
+  }
+
+  constructor(options: PreciseWallClockOptions = {}) {
+    this._origin = {
+      clockTime: options.origin ?? process.hrtime(),
+      preciseDate: options.now ?? new PreciseDate(),
+    };
+  }
+
+  preciseNow(): [number, number] {
+    const elapsedHrTime = process.hrtime(this.#originClockTime);
+    const elapsedNanoseconds = BigInt(elapsedHrTime[0]) * BigInt(1e9) + BigInt(elapsedHrTime[1]);
+
+    const preciseDate = new PreciseDate(this.#originPreciseDate.getFullTime() + elapsedNanoseconds);
+    const dateStruct = preciseDate.toStruct();
+
+    return [dateStruct.seconds, dateStruct.nanos];
+  }
+
+  reset() {
+    this._origin = {
+      clockTime: process.hrtime(),
+      preciseDate: new PreciseDate(),
+    };
+  }
+}

packages/core/src/v3/clock/preciseWallClock.ts

@@ -8,4 +8,8 @@ export class SimpleClock implements Clock {
 
     return [nowStruct.seconds, nowStruct.nanos];
   }
+
+  reset() {
+    // do nothing
+  }
 }

packages/core/src/v3/clock/simpleClock.ts

@@ -36,7 +36,7 @@ export { taskContextManager, TaskContextSpanProcessor } from "./tasks/taskContex
 export type { RuntimeManager } from "./runtime/manager";
 export { DevRuntimeManager } from "./runtime/devRuntimeManager";
 export { ProdRuntimeManager } from "./runtime/prodRuntimeManager";
-export { DurableClock } from "./clock/durableClock";
+export { PreciseWallClock as DurableClock } from "./clock/preciseWallClock";
 export { TriggerTracer } from "./tracer";
 
 export type { TaskLogger } from "./logger/taskLogger";

packages/core/src/v3/index.ts

@@ -136,12 +136,16 @@ export class TracingSDK {
 
     traceProvider.addSpanProcessor(
       new TaskContextSpanProcessor(
-        process.env.OTEL_BATCH_PROCESSING_ENABLED === "1"
+        getEnvVar("OTEL_BATCH_PROCESSING_ENABLED") === "1"
           ? new BatchSpanProcessor(spanExporter, {
-              maxExportBatchSize: parseInt(process.env.OTEL_SPAN_MAX_EXPORT_BATCH_SIZE ?? "64"),
-              scheduledDelayMillis: parseInt(process.env.OTEL_SPAN_SCHEDULED_DELAY_MILLIS ?? "200"),
-              exportTimeoutMillis: parseInt(process.env.OTEL_SPAN_EXPORT_TIMEOUT_MILLIS ?? "30000"),
-              maxQueueSize: parseInt(process.env.OTEL_SPAN_MAX_QUEUE_SIZE ?? "512"),
+              maxExportBatchSize: parseInt(getEnvVar("OTEL_SPAN_MAX_EXPORT_BATCH_SIZE") ?? "64"),
+              scheduledDelayMillis: parseInt(
+                getEnvVar("OTEL_SPAN_SCHEDULED_DELAY_MILLIS") ?? "200"
+              ),
+              exportTimeoutMillis: parseInt(
+                getEnvVar("OTEL_SPAN_EXPORT_TIMEOUT_MILLIS") ?? "30000"
+              ),
+              maxQueueSize: parseInt(getEnvVar("OTEL_SPAN_MAX_QUEUE_SIZE") ?? "512"),
             })
           : new SimpleSpanProcessor(spanExporter)
       )
@@ -169,12 +173,12 @@ export class TracingSDK {
 
     loggerProvider.addLogRecordProcessor(
       new TaskContextLogProcessor(
-        process.env.OTEL_BATCH_PROCESSING_ENABLED === "1"
+        getEnvVar("OTEL_BATCH_PROCESSING_ENABLED") === "1"
           ? new BatchLogRecordProcessor(logExporter, {
-              maxExportBatchSize: parseInt(process.env.OTEL_LOG_MAX_EXPORT_BATCH_SIZE ?? "64"),
-              scheduledDelayMillis: parseInt(process.env.OTEL_LOG_SCHEDULED_DELAY_MILLIS ?? "200"),
-              exportTimeoutMillis: parseInt(process.env.OTEL_LOG_EXPORT_TIMEOUT_MILLIS ?? "30000"),
-              maxQueueSize: parseInt(process.env.OTEL_LOG_MAX_QUEUE_SIZE ?? "512"),
+              maxExportBatchSize: parseInt(getEnvVar("OTEL_LOG_MAX_EXPORT_BATCH_SIZE") ?? "64"),
+              scheduledDelayMillis: parseInt(getEnvVar("OTEL_LOG_SCHEDULED_DELAY_MILLIS") ?? "200"),
+              exportTimeoutMillis: parseInt(getEnvVar("OTEL_LOG_EXPORT_TIMEOUT_MILLIS") ?? "30000"),
+              maxQueueSize: parseInt(getEnvVar("OTEL_LOG_MAX_QUEUE_SIZE") ?? "512"),
             })
           : new SimpleLogRecordProcessor(logExporter)
       )

packages/core/src/v3/otel/tracingSDK.ts

@@ -1,3 +1,4 @@
+import { clock } from "../clock-api";
 import {
   BatchTaskRunExecutionResult,
   ProdChildToWorkerMessages,
@@ -9,6 +10,11 @@ import {
 } from "../schemas";
 import { ZodIpcConnection } from "../zodIpc";
 import { RuntimeManager } from "./manager";
+import { setTimeout } from "node:timers/promises";
+
+export type ProdRuntimeManagerOptions = {
+  waitThresholdInMs?: number;
+};
 
 export class ProdRuntimeManager implements RuntimeManager {
   _taskWaits: Map<
@@ -21,15 +27,16 @@ export class ProdRuntimeManager implements RuntimeManager {
     { resolve: (value: BatchTaskRunExecutionResult) => void; reject: (err?: any) => void }
   > = new Map();
 
-  _waitForRestore: { resolve: (value?: any) => void; reject: (err?: any) => void } | undefined;
+  _waitForRestore: { resolve: (value: "restore") => void; reject: (err?: any) => void } | undefined;
 
   _tasks: Map<string, TaskMetadataWithFilePath> = new Map();
 
   constructor(
     private ipc: ZodIpcConnection<
       typeof ProdWorkerToChildMessages,
       typeof ProdChildToWorkerMessages
-    >
+    >,
+    private options: ProdRuntimeManagerOptions = {}
   ) {}
 
   disable(): void {
@@ -47,20 +54,16 @@ export class ProdRuntimeManager implements RuntimeManager {
   }
 
   async waitForDuration(ms: number): Promise<void> {
-    let timeout: NodeJS.Timeout | undefined;
-
     const now = Date.now();
 
-    const resolveAfterDuration = new Promise((resolve) => {
-      timeout = setTimeout(resolve, ms);
-    });
+    const resolveAfterDuration = setTimeout(ms, "duration" as const);
 
-    if (ms < 10_000) {
+    if (ms <= this.waitThresholdInMs) {
       await resolveAfterDuration;
       return;
     }
 
-    const waitForRestore = new Promise((resolve, reject) => {
+    const waitForRestore = new Promise<"restore">((resolve, reject) => {
       this._waitForRestore = { resolve, reject };
     });
 
@@ -81,16 +84,17 @@ export class ProdRuntimeManager implements RuntimeManager {
 
     // The coordinator can then cancel any in-progress checkpoints
     this.ipc.send("CANCEL_CHECKPOINT", {});
-
-    clearTimeout(timeout);
   }
 
   resumeAfterRestore(): void {
     if (!this._waitForRestore) {
       return;
     }
 
-    this._waitForRestore.resolve();
+    // Resets the clock to the current time
+    clock.reset();
+
+    this._waitForRestore.resolve("restore");
     this._waitForRestore = undefined;
   }
 
@@ -155,4 +159,8 @@ export class ProdRuntimeManager implements RuntimeManager {
 
     this._taskWaits.delete(execution.run.id);
   }
+
+  private get waitThresholdInMs(): number {
+    return this.options.waitThresholdInMs ?? 30_000;
+  }
 }

packages/core/src/v3/runtime/prodRuntimeManager.ts

@@ -107,14 +107,21 @@ export const parentTask = task({
 
 export const childTask = task({
   id: "child-task",
-  run: async (payload: { message: string; forceError: boolean }, { ctx }) => {
+  run: async (
+    payload: { message: string; forceError: boolean; delayInSeconds?: number },
+    { ctx }
+  ) => {
     logger.info("Child task payload", { payload });
     logger.info("Child task payload 2", { payload });
     logger.info("Child task payload 3", { payload });
     logger.info("Child task payload 4", { payload });
     logger.info("Child task payload 5", { payload });
 
-    await wait.for({ seconds: 10 });
+    await wait.for({ seconds: payload.delayInSeconds ?? 5 });
+
+    logger.info("Child task payload 6", { payload });
+    logger.info("Child task payload 7", { payload });
+    logger.info("Child task payload 8", { payload });
 
     const response = await fetch("https://jsonhero.io/api/create.json", {
       method: "POST",