Skip to content

Commit 23ae77a

Browse files
ericallamericallam
committed
handle max duration timeout errors better
1 parent 30705fd commit 23ae77a

File tree

3 files changed

+137
-51
lines changed

3 files changed

+137
-51
lines changed

‎packages/cli-v3/src/entryPoints/dev-run-worker.ts

Lines changed: 0 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -377,35 +377,6 @@ const zodIpc = new ZodIpcConnection({
377377
? timeout.abortAfterTimeout(execution.run.maxDuration)
378378
: undefined;
379379

380-
signal?.addEventListener("abort", async (e) => {
381-
if (_isRunning) {
382-
_isRunning = false;
383-
_execution = undefined;
384-
385-
const usageSample = usage.stop(measurement);
386-
387-
await sender.send("TASK_RUN_COMPLETED", {
388-
execution,
389-
result: {
390-
ok: false,
391-
id: execution.run.id,
392-
error: {
393-
type: "INTERNAL_ERROR",
394-
code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED,
395-
message:
396-
signal.reason instanceof Error
397-
? signal.reason.message
398-
: String(signal.reason),
399-
},
400-
usage: {
401-
durationMs: usageSample.cpuTime,
402-
},
403-
metadata: runMetadataManager.stopAndReturnLastFlush(),
404-
},
405-
});
406-
}
407-
});
408-
409380
const { result } = await executor.execute(execution, metadata, traceContext, signal);
410381

411382
const usageSample = usage.stop(measurement);

‎packages/core/src/v3/workers/taskExecutor.ts

Lines changed: 30 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,13 @@ import { SpanKind } from "@opentelemetry/api";
22
import { VERSION } from "../../version.js";
33
import { ApiError, RateLimitError } from "../apiClient/errors.js";
44
import { ConsoleInterceptor } from "../consoleInterceptor.js";
5-
import { isInternalError, parseError, sanitizeError, TaskPayloadParsedError } from "../errors.js";
5+
import {
6+
InternalError,
7+
isInternalError,
8+
parseError,
9+
sanitizeError,
10+
TaskPayloadParsedError,
11+
} from "../errors.js";
612
import { flattenAttributes, lifecycleHooks, runMetadata, waitUntil } from "../index.js";
713
import {
814
AnyOnMiddlewareHookFunction,
@@ -341,9 +347,29 @@ export class TaskExecutor {
341347
throw new Error("Task does not have a run function");
342348
}
343349

344-
return runTimelineMetrics.measureMetric("trigger.dev/execution", "run", () =>
345-
runFn(payload, { ctx, init, signal })
346-
);
350+
// Create a promise that rejects when the signal aborts
351+
const abortPromise = signal
352+
? new Promise((_, reject) => {
353+
signal.addEventListener("abort", () => {
354+
const maxDuration = ctx.run.maxDuration;
355+
reject(
356+
new InternalError({
357+
code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED,
358+
message: `Task execution exceeded maximum duration of ${maxDuration}ms`,
359+
})
360+
);
361+
});
362+
})
363+
: undefined;
364+
365+
return runTimelineMetrics.measureMetric("trigger.dev/execution", "run", async () => {
366+
if (abortPromise) {
367+
// Race between the run function and the abort promise
368+
return await Promise.race([runFn(payload, { ctx, init, signal }), abortPromise]);
369+
}
370+
371+
return await runFn(payload, { ctx, init, signal });
372+
});
347373
}
348374

349375
async #callInitFunctions(payload: unknown, ctx: TaskRunContext, signal?: AbortSignal) {

‎packages/core/test/taskExecutor.test.ts

Lines changed: 107 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@ import {
44
RunFnParams,
55
ServerBackgroundWorker,
66
TaskMetadataWithFunctions,
7+
TaskRunErrorCodes,
78
TaskRunExecution,
89
} from "../src/v3/index.js";
910
import { TracingSDK } from "../src/v3/otel/tracingSDK.js";
@@ -52,7 +53,7 @@ describe("TaskExecutor", () => {
5253
},
5354
};
5455

55-
const result = await executeTask(task, {});
56+
const result = await executeTask(task, {}, undefined);
5657

5758
expect(result).toEqual({
5859
result: {
@@ -136,7 +137,7 @@ describe("TaskExecutor", () => {
136137
},
137138
};
138139

139-
const result = await executeTask(task, {});
140+
const result = await executeTask(task, {}, undefined);
140141

141142
// Verify hooks were called in correct order - should match registration order
142143
expect(globalSuccessOrder).toEqual(["global-2", "global-1", "task"]);
@@ -238,7 +239,7 @@ describe("TaskExecutor", () => {
238239
},
239240
};
240241

241-
const result = await executeTask(task, { test: "data" });
242+
const result = await executeTask(task, { test: "data" }, undefined);
242243

243244
// Verify hooks were called in correct order
244245
expect(globalStartOrder).toEqual(["global-1", "global-2", "task"]);
@@ -337,7 +338,7 @@ describe("TaskExecutor", () => {
337338
},
338339
};
339340

340-
const result = await executeTask(task, { test: "data" });
341+
const result = await executeTask(task, { test: "data" }, undefined);
341342

342343
// Verify hooks were called in correct order
343344
expect(globalFailureOrder).toEqual(["global-1", "global-2", "task"]);
@@ -445,7 +446,7 @@ describe("TaskExecutor", () => {
445446
},
446447
};
447448

448-
const result = await executeTask(task, { test: "data" });
449+
const result = await executeTask(task, { test: "data" }, undefined);
449450

450451
// Verify hooks were called in correct order
451452
expect(globalCompleteOrder).toEqual(["global-1", "global-2", "task"]);
@@ -533,7 +534,7 @@ describe("TaskExecutor", () => {
533534
},
534535
};
535536

536-
const result = await executeTask(task, { test: "data" });
537+
const result = await executeTask(task, { test: "data" }, undefined);
537538

538539
// Verify hooks were called in correct order
539540
expect(globalCompleteOrder).toEqual(["global", "task"]);
@@ -645,7 +646,7 @@ describe("TaskExecutor", () => {
645646
},
646647
};
647648

648-
const result = await executeTask(task, { test: "data" });
649+
const result = await executeTask(task, { test: "data" }, undefined);
649650

650651
// Verify hooks were called in correct order and stopped after second global hook
651652
expect(hookCallOrder).toEqual(["task", "global-1", "global-2"]);
@@ -707,7 +708,7 @@ describe("TaskExecutor", () => {
707708
},
708709
};
709710

710-
const result = await executeTask(task, { test: "data" });
711+
const result = await executeTask(task, { test: "data" }, undefined);
711712

712713
// Verify only task hook was called
713714
expect(hookCallOrder).toEqual(["task"]);
@@ -755,7 +756,7 @@ describe("TaskExecutor", () => {
755756
},
756757
};
757758

758-
const result = await executeTask(task, { test: "data" });
759+
const result = await executeTask(task, { test: "data" }, undefined);
759760

760761
// Verify only task hook was called
761762
expect(hookCallOrder).toEqual(["task"]);
@@ -863,7 +864,7 @@ describe("TaskExecutor", () => {
863864
},
864865
};
865866

866-
const result = await executeTask(task, { test: "data" });
867+
const result = await executeTask(task, { test: "data" }, undefined);
867868

868869
// Verify the execution order:
869870
// 1. Global middlewares (outside to inside)
@@ -935,7 +936,7 @@ describe("TaskExecutor", () => {
935936
},
936937
};
937938

938-
const result = await executeTask(task, { test: "data" });
939+
const result = await executeTask(task, { test: "data" }, undefined);
939940

940941
// Verify only the middleware-before hook ran
941942
expect(executionOrder).toEqual(["middleware-before"]);
@@ -1012,7 +1013,7 @@ describe("TaskExecutor", () => {
10121013
},
10131014
};
10141015

1015-
const result = await executeTask(task, { test: "data" });
1016+
const result = await executeTask(task, { test: "data" }, undefined);
10161017

10171018
// Verify only the global init hook ran, and failure/complete hooks were called
10181019
expect(executionOrder).toEqual(["global-init", "failure", "complete"]);
@@ -1094,7 +1095,7 @@ describe("TaskExecutor", () => {
10941095
},
10951096
};
10961097

1097-
const result = await executeTask(task, { test: "data" });
1098+
const result = await executeTask(task, { test: "data" }, undefined);
10981099

10991100
// Verify both init hooks ran, but run wasn't called, and failure/complete hooks were called
11001101
expect(executionOrder).toEqual(["global-init", "task-init", "failure", "complete"]);
@@ -1184,7 +1185,7 @@ describe("TaskExecutor", () => {
11841185
},
11851186
};
11861187

1187-
const result = await executeTask(task, { test: "data" });
1188+
const result = await executeTask(task, { test: "data" }, undefined);
11881189

11891190
// Verify init succeeded, start hook failed, and run wasn't called
11901191
expect(executionOrder).toEqual(["global-init", "global-start", "failure", "complete"]);
@@ -1295,7 +1296,7 @@ describe("TaskExecutor", () => {
12951296
},
12961297
};
12971298

1298-
const result = await executeTask(task, { test: "data" });
1299+
const result = await executeTask(task, { test: "data" }, undefined);
12991300

13001301
// Verify the execution order:
13011302
// 1. Middleware starts
@@ -1390,7 +1391,7 @@ describe("TaskExecutor", () => {
13901391
},
13911392
};
13921393

1393-
const result = await executeTask(task, { test: "data" });
1394+
const result = await executeTask(task, { test: "data" }, undefined);
13941395

13951396
// Verify cleanup hooks are called even after failure
13961397
expect(executionOrder).toEqual([
@@ -1417,9 +1418,96 @@ describe("TaskExecutor", () => {
14171418
},
14181419
});
14191420
});
1421+
1422+
test("should handle max duration abort signal and call hooks in correct order", async () => {
1423+
const executionOrder: string[] = [];
1424+
const maxDurationMs = 1000;
1425+
1426+
// Create an abort controller that we'll trigger manually
1427+
const controller = new AbortController();
1428+
1429+
// Register global init hook
1430+
lifecycleHooks.registerGlobalInitHook({
1431+
id: "test-init",
1432+
fn: async () => {
1433+
executionOrder.push("init");
1434+
return {
1435+
foo: "bar",
1436+
};
1437+
},
1438+
});
1439+
1440+
// Register failure hook
1441+
lifecycleHooks.registerGlobalFailureHook({
1442+
id: "global-failure",
1443+
fn: async ({ error }) => {
1444+
executionOrder.push("failure");
1445+
expect((error as Error).message).toBe(
1446+
`Task execution exceeded maximum duration of ${maxDurationMs}ms`
1447+
);
1448+
},
1449+
});
1450+
1451+
// Register complete hook
1452+
lifecycleHooks.registerGlobalCompleteHook({
1453+
id: "global-complete",
1454+
fn: async ({ result }) => {
1455+
executionOrder.push("complete");
1456+
expect(result.ok).toBe(false);
1457+
},
1458+
});
1459+
1460+
// Register cleanup hook
1461+
lifecycleHooks.registerGlobalCleanupHook({
1462+
id: "global-cleanup",
1463+
fn: async () => {
1464+
executionOrder.push("cleanup");
1465+
},
1466+
});
1467+
1468+
const task = {
1469+
id: "test-task",
1470+
fns: {
1471+
run: async (payload: any, params: RunFnParams<any>) => {
1472+
executionOrder.push("run-start");
1473+
1474+
// Create a promise that never resolves
1475+
await new Promise((resolve) => {
1476+
// Trigger abort after a small delay
1477+
setTimeout(() => {
1478+
controller.abort();
1479+
}, 10);
1480+
});
1481+
1482+
// This should never be reached
1483+
executionOrder.push("run-end");
1484+
},
1485+
},
1486+
};
1487+
1488+
const result = await executeTask(task, { test: "data" }, controller.signal);
1489+
1490+
// Verify hooks were called in correct order
1491+
expect(executionOrder).toEqual(["init", "run-start", "failure", "complete", "cleanup"]);
1492+
1493+
// Verify the error result
1494+
expect(result).toEqual({
1495+
result: {
1496+
ok: false,
1497+
id: "test-run-id",
1498+
error: {
1499+
type: "INTERNAL_ERROR",
1500+
code: TaskRunErrorCodes.MAX_DURATION_EXCEEDED,
1501+
message: "Task execution exceeded maximum duration of 1000ms",
1502+
stackTrace: expect.any(String),
1503+
},
1504+
skippedRetrying: false,
1505+
},
1506+
});
1507+
});
14201508
});
14211509

1422-
function executeTask(task: TaskMetadataWithFunctions, payload: any) {
1510+
function executeTask(task: TaskMetadataWithFunctions, payload: any, signal?: AbortSignal) {
14231511
const tracingSDK = new TracingSDK({
14241512
url: "http://localhost:4318",
14251513
});
@@ -1472,6 +1560,7 @@ function executeTask(task: TaskMetadataWithFunctions, payload: any) {
14721560
costInCents: 0,
14731561
baseCostInCents: 0,
14741562
priority: 0,
1563+
maxDuration: 1000,
14751564
},
14761565
machine: {
14771566
name: "micro",
@@ -1508,5 +1597,5 @@ function executeTask(task: TaskMetadataWithFunctions, payload: any) {
15081597
engine: "V2",
15091598
};
15101599

1511-
return executor.execute(execution, worker, {});
1600+
return executor.execute(execution, worker, {}, signal);
15121601
}

0 commit comments

Comments
 (0)