Visitar URL original
fix(run-engine,webapp): make batch expiration resumable and schedule … · triggerdotdev/trigger.dev@49f0d3c · GitHub
Skip to content

Commit 49f0d3c

Browse files
ericallamericallam
committed
fix(run-engine,webapp): make batch expiration resumable and schedule it before blocking
Aborting the batch and completing the parent waitpoint are two writes. A crash between them left the parent blocked forever, because the retry saw a non-pending batch and returned early. An already-aborted batch now falls through to complete its waitpoint, which is a no-op if it already happened. The reaper is also scheduled before the parent is blocked, so a failed enqueue can no longer leave a blocked parent with nothing to recover it. Adds a bounded counter for expiration outcomes, and drops a logger mock from the grant tests to match the no-mocks policy.
1 parent 9e3cbd8 commit 49f0d3c

4 files changed

Lines changed: 99 additions & 30 deletions

File tree

‎apps/webapp/app/runEngine/services/createBatch.server.ts‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -126,18 +126,18 @@ export class CreateBatchService extends WithRunEngine {
126126

127127
// Block parent run if this is a batchTriggerAndWait
128128
if (body.parentRunId && body.resumeParentOnCompletion) {
129+
await this._engine.scheduleExpireBatch({
130+
batchId: batch.id,
131+
availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS),
132+
});
133+
129134
await this._engine.blockRunWithCreatedBatch({
130135
runId: RunId.fromFriendlyId(body.parentRunId),
131136
batchId: batch.id,
132137
environmentId: environment.id,
133138
projectId: environment.projectId,
134139
organizationId: environment.organizationId,
135140
});
136-
137-
await this._engine.scheduleExpireBatch({
138-
batchId: batch.id,
139-
availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS),
140-
});
141141
}
142142

143143
// Initialize batch metadata in Redis (without items)

‎apps/webapp/test/batchStreamGrants.test.ts‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,15 +3,6 @@ import { describe, expect, vi } from "vitest";
33

44
vi.setConfig({ testTimeout: 30_000 });
55

6-
vi.mock("../app/services/logger.server", () => ({
7-
logger: {
8-
info: vi.fn(),
9-
warn: vi.fn(),
10-
debug: vi.fn(),
11-
error: vi.fn(),
12-
},
13-
}));
14-
156
import { BatchStreamGrants } from "../app/runEngine/concerns/batchStreamGrants.server.js";
167

178
describe("BatchStreamGrants", () => {

‎internal-packages/run-engine/src/engine/systems/batchSystem.ts‎

Lines changed: 42 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { startSpan } from "@internal/tracing";
1+
import { startSpan, type Counter } from "@internal/tracing";
22
import type { TaskRunError } from "@trigger.dev/core/v3/schemas";
33
import { isFinalRunStatus } from "../statuses.js";
44
import type { SystemResources } from "./systems.js";
@@ -12,10 +12,15 @@ export type BatchSystemOptions = {
1212
export class BatchSystem {
1313
private readonly $: SystemResources;
1414
private readonly waitpointSystem: WaitpointSystem;
15+
private readonly expirationsCounter: Counter;
1516

1617
constructor(private readonly options: BatchSystemOptions) {
1718
this.$ = options.resources;
1819
this.waitpointSystem = options.waitpointSystem;
20+
this.expirationsCounter = this.$.meter.createCounter("batch_system.expirations", {
21+
description: "Seal-timeout reaper runs, by outcome",
22+
unit: "batches",
23+
});
1924
}
2025

2126
public async scheduleCompleteBatch({ batchId }: { batchId: string }): Promise<void> {
@@ -57,6 +62,11 @@ export class BatchSystem {
5762
* waitpoint, so without this a stream that never finishes strands the parent permanently.
5863
*
5964
* Idempotent and race-safe: a batch that sealed in the meantime is left alone.
65+
*
66+
* Resumable too. The abort and the waitpoint completion are two writes, so a crash between
67+
* them would otherwise leave the parent blocked with no way back: the retry would see a
68+
* non-PENDING batch and bail. An already-ABORTED batch therefore falls through to complete
69+
* its waitpoint again, which {@link WaitpointSystem.completeWaitpoint} treats as a no-op.
6070
*/
6171
public async expireBatch({ batchId }: { batchId: string }): Promise<void> {
6272
return startSpan(this.$.tracer, "expireBatch", async (span) => {
@@ -66,33 +76,46 @@ export class BatchSystem {
6676

6777
if (!batch) {
6878
this.$.logger.debug("expireBatch: batch doesn't exist", { batchId });
79+
this.expirationsCounter.add(1, { outcome: "missing" });
6980
return;
7081
}
7182

72-
if (batch.sealed || batch.status !== "PENDING") {
73-
this.$.logger.debug("expireBatch: batch sealed or no longer PENDING, nothing to do", {
83+
if (batch.sealed || (batch.status !== "PENDING" && batch.status !== "ABORTED")) {
84+
this.$.logger.debug("expireBatch: batch sealed or already finished, nothing to do", {
7485
batchId,
7586
status: batch.status,
7687
sealed: batch.sealed,
7788
});
89+
this.expirationsCounter.add(1, { outcome: "already_settled" });
7890
return;
7991
}
8092

81-
const aborted = await this.$.runStore.updateManyBatchTaskRun(
82-
{
83-
where: { id: batchId, sealed: false, status: "PENDING" },
84-
data: {
85-
status: "ABORTED",
86-
completedAt: new Date(),
87-
processingCompletedAt: new Date(),
93+
if (batch.status === "PENDING") {
94+
const aborted = await this.$.runStore.updateManyBatchTaskRun(
95+
{
96+
where: { id: batchId, sealed: false, status: "PENDING" },
97+
data: {
98+
status: "ABORTED",
99+
completedAt: new Date(),
100+
processingCompletedAt: new Date(),
101+
},
88102
},
89-
},
90-
this.$.prisma
91-
);
103+
this.$.prisma
104+
);
92105

93-
if (aborted.count === 0) {
94-
this.$.logger.debug("expireBatch: lost the race to a seal, no-op", { batchId });
95-
return;
106+
if (aborted.count === 0) {
107+
const current = await this.$.runStore.findBatchTaskRunById(
108+
batchId,
109+
undefined,
110+
this.$.prisma
111+
);
112+
113+
if (!current || current.sealed || current.status !== "ABORTED") {
114+
this.$.logger.debug("expireBatch: lost the race to a seal, no-op", { batchId });
115+
this.expirationsCounter.add(1, { outcome: "lost_race" });
116+
return;
117+
}
118+
}
96119
}
97120

98121
const waitpoint = await this.$.runStore.findWaitpoint(
@@ -106,6 +129,7 @@ export class BatchSystem {
106129
this.$.logger.debug("expireBatch: no waitpoint, nothing was blocked on this batch", {
107130
batchId,
108131
});
132+
this.expirationsCounter.add(1, { outcome: "aborted_no_parent" });
109133
return;
110134
}
111135

@@ -123,6 +147,8 @@ export class BatchSystem {
123147
output: { value: JSON.stringify(error), isError: true },
124148
});
125149

150+
this.expirationsCounter.add(1, { outcome: "aborted_parent_resumed" });
151+
126152
this.$.logger.warn("expireBatch: aborted an unsealed batch and resumed its parent", {
127153
batchId,
128154
waitpointId: waitpoint.id,

‎internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -858,6 +858,58 @@ describe("RunEngine 2-Phase Batch API", () => {
858858
}
859859
});
860860

861+
containerTest(
862+
"resumes the parent when a previous run aborted the batch but died before the waitpoint",
863+
async ({ prisma, redisOptions }) => {
864+
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
865+
const engine = createEngine(prisma, redisOptions);
866+
867+
try {
868+
const { batchId, parentRun } = await setupUnsealedBatchBlockingParent(
869+
engine,
870+
prisma,
871+
authenticatedEnvironment
872+
);
873+
874+
await prisma.batchTaskRun.update({
875+
where: { id: batchId },
876+
data: {
877+
status: "ABORTED",
878+
completedAt: new Date(),
879+
processingCompletedAt: new Date(),
880+
},
881+
});
882+
883+
const beforeRetry = await prisma.waitpoint.findFirst({
884+
where: { completedByBatchId: batchId },
885+
});
886+
assertNonNullable(beforeRetry);
887+
expect(beforeRetry.status).toBe("PENDING");
888+
889+
await engine.expireBatch({ batchId });
890+
891+
const waitpoint = await prisma.waitpoint.findFirst({
892+
where: { completedByBatchId: batchId },
893+
});
894+
assertNonNullable(waitpoint);
895+
expect(waitpoint.status).toBe("COMPLETED");
896+
expect(waitpoint.outputIsError).toBe(true);
897+
898+
await vi.waitFor(
899+
async () => {
900+
const waitpoints = await prisma.taskRunWaitpoint.findMany({
901+
where: { taskRunId: parentRun.id },
902+
});
903+
expect(waitpoints.length).toBe(0);
904+
},
905+
{ timeout: 15_000 }
906+
);
907+
} finally {
908+
await engine.quit();
909+
}
910+
}
911+
);
912+
861913
containerTest(
862914
"aborts a fire-and-forget batch without a waitpoint to resolve",
863915
async ({ prisma, redisOptions }) => {

0 commit comments

Comments
 (0)