Skip to content

Commit 88c5745

Browse files
d-csTrigger.dev RepoOps
authored andcommitted
feat(run-engine): retain native block state and coordinator contracts
Add internal MemoryDB waitpoint coordination with retained results, bounded cleanup, and three-phase registration. Preserve shared waits during failed-registration cleanup and return saved completion output. Application activation remains separate. Mono-RevId: 723ff81176185d44cc63f765b80ae94550a1d6db
1 parent bd6112e commit 88c5745

34 files changed

Lines changed: 6130 additions & 426 deletions

‎internal-packages/run-engine/src/engine/errors.ts‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -188,6 +188,61 @@ export class StoreWaitpointRegistrationUnsupported extends Error {
188188
}
189189
}
190190

191+
/**
192+
* A store-format blocker reached the block seam without a caller-owned block-operation identity. The
193+
* store's compare-and-set needs a stable, caller-supplied identity that survives retries and is captured
194+
* with the logical operation — NOT re-derived from whatever snapshot is latest when the operation runs —
195+
* so a delayed operation from an older cycle carries its ORIGINAL identity and is refused stale after a
196+
* newer cycle is installed. Without it, a delayed replay would silently adopt the current cycle's identity
197+
* and be accepted. Legacy (Postgres) blockers need no identity and never reach this. In the dark default
198+
* (no store minting) this is unreachable.
199+
*/
200+
export class StoreWaitpointBlockOperationIdentityRequired extends Error {
201+
readonly runId: string;
202+
readonly waitpointIds: string[];
203+
constructor(params: { runId: string; waitpointIds: string[] }) {
204+
super(
205+
`Cannot register store-format waitpoint(s) [${params.waitpointIds.join(
206+
", "
207+
)}] against run ${params.runId}: a store block operation requires a caller-owned block identity ` +
208+
`(blockOperation.id) that survives retries; re-deriving it from the latest snapshot cannot fence a ` +
209+
`delayed operation from an older cycle`
210+
);
211+
this.name = "StoreWaitpointBlockOperationIdentityRequired";
212+
this.runId = params.runId;
213+
this.waitpointIds = params.waitpointIds;
214+
}
215+
}
216+
217+
/**
218+
* Reject an operation from an earlier run cycle before it changes MemoryDB.
219+
* Check the snapshot anchor even if cleanup has already removed the old block.
220+
* Legacy waits do not use this check.
221+
*/
222+
export class StoreWaitpointBlockCycleAdvanced extends Error {
223+
readonly runId: string;
224+
readonly waitpointIds: string[];
225+
readonly expectedCycleAnchor: string;
226+
readonly currentCycleAnchor: string;
227+
constructor(params: {
228+
runId: string;
229+
waitpointIds: string[];
230+
expectedCycleAnchor: string;
231+
currentCycleAnchor: string;
232+
}) {
233+
super(
234+
`Store block operation for run ${params.runId} carries cycle anchor ${params.expectedCycleAnchor}, ` +
235+
`but the run's active cycle has advanced to ${params.currentCycleAnchor}; refusing to re-block a ` +
236+
`run whose attempt is over`
237+
);
238+
this.name = "StoreWaitpointBlockCycleAdvanced";
239+
this.runId = params.runId;
240+
this.waitpointIds = params.waitpointIds;
241+
this.expectedCycleAnchor = params.expectedCycleAnchor;
242+
this.currentCycleAnchor = params.currentCycleAnchor;
243+
}
244+
}
245+
191246
/**
192247
* A run whose organization could hold MemoryDB block edges was asked to clear its block state inside a
193248
* rollback-capable Postgres lifecycle transaction (`tx`). Clearing MemoryDB cannot join that
@@ -210,6 +265,30 @@ export class TransactionalStoreClearUnsupported extends Error {
210265
}
211266
}
212267

268+
/**
269+
* A mixed-residency absorption accepted the MemoryDB half but its legacy half failed AND the immediate
270+
* compensating rollback of the store contribution also failed. Carries the exact fenced targets so the
271+
* caller can arm a durable retry rather than report the rejected operation as clean while leaving store
272+
* residue. `compensationCause` is the failed cleanup error.
273+
*/
274+
export class StoreAbsorptionCompensationFailed extends Error {
275+
readonly runId: string;
276+
readonly blockId: string;
277+
readonly absorptionId: string;
278+
readonly compensationCause: unknown;
279+
constructor(args: { runId: string; blockId: string; absorptionId: string; cause: unknown }) {
280+
super(
281+
`Mixed absorption for run ${args.runId} block ${args.blockId} accepted MemoryDB but failed legacy, ` +
282+
`and compensating the store contribution failed.`
283+
);
284+
this.name = "StoreAbsorptionCompensationFailed";
285+
this.runId = args.runId;
286+
this.blockId = args.blockId;
287+
this.absorptionId = args.absorptionId;
288+
this.compensationCause = args.cause;
289+
}
290+
}
291+
213292
/**
214293
* A store-format waitpoint id reached the router but no store coordinator is wired. A recognised
215294
* store id must never fall back to Postgres, so this fails loud rather than mis-routing. In the dark

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,7 @@ export class RunEngine {
293293
finishWaitpoint: async ({ payload }) => {
294294
await this.waitpointSystem.completeWaitpoint({
295295
id: payload.waitpointId,
296+
organizationId: payload.organizationId,
296297
output: payload.error
297298
? {
298299
value: payload.error,
@@ -351,8 +352,20 @@ export class RunEngine {
351352
await this.waitpointSystem.ensureWaitpointCompleted({
352353
waitpointId: payload.waitpointId,
353354
output: payload.output,
355+
organizationId: payload.organizationId,
354356
});
355357
},
358+
ensureAbsorptionCompensated: async ({ payload }) => {
359+
await this.waitpointSystem.ensureAbsorptionCompensated({
360+
runId: payload.runId,
361+
blockId: payload.blockId,
362+
absorptionId: payload.absorptionId,
363+
organizationId: payload.organizationId,
364+
});
365+
},
366+
recoverWaitpointBlock: async ({ payload }) => {
367+
await this.waitpointSystem.recoverWaitpointBlock(payload);
368+
},
356369
enqueueDelayedRun: async ({ payload }) => {
357370
await this.delayedRunSystem.enqueueDelayedRun({ runId: payload.runId });
358371
},
@@ -437,6 +450,8 @@ export class RunEngine {
437450

438451
this.checkpointSystem = new CheckpointSystem({
439452
resources,
453+
waitpointStoreEnabled: (organizationId) =>
454+
this.waitpointSystem.isStoreRoutingEnabled(organizationId),
440455
executionSnapshotSystem: this.executionSnapshotSystem,
441456
enqueueSystem: this.enqueueSystem,
442457
});
@@ -474,6 +489,7 @@ export class RunEngine {
474489
executionSnapshotSystem: this.executionSnapshotSystem,
475490
enqueueSystem: this.enqueueSystem,
476491
completionGuardDelayMs: options.completionGuardDelayMs,
492+
waitpointStore: options.waitpointStore,
477493
});
478494

479495
this.ttlSystem = new TtlSystem({
@@ -2068,6 +2084,7 @@ export class RunEngine {
20682084
workerId,
20692085
runnerId,
20702086
expectedEnvironmentId,
2087+
blockOperation,
20712088
tx,
20722089
}: {
20732090
runId: string;
@@ -2084,6 +2101,14 @@ export class RunEngine {
20842101
// id and yields nothing (404) for a foreign run, before any edge is registered. Omitted by trusted
20852102
// internal callers.
20862103
expectedEnvironmentId?: string;
2104+
// Use the block ID and cycle anchor captured by the caller.
2105+
// Required for MemoryDB waits; ignored for Postgres waits.
2106+
blockOperation?: {
2107+
id: string;
2108+
absorptionId: string;
2109+
cycleAnchor: string;
2110+
expectedPreviousBlockId?: string;
2111+
};
20872112
tx?: PrismaClientOrTransaction;
20882113
}): Promise<TaskRunExecutionSnapshot> {
20892114
return this.waitpointSystem.blockRunWithWaitpoint({
@@ -2097,6 +2122,7 @@ export class RunEngine {
20972122
workerId,
20982123
runnerId,
20992124
expectedEnvironmentId,
2125+
blockOperation,
21002126
tx,
21012127
});
21022128
}

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

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import type { SystemResources } from "./systems.js";
1616
import { snapshotRouteFromSnapshot } from "@internal/run-store";
1717
import { ServiceValidationError } from "../errors.js";
1818
import type { EnqueueSystem } from "./enqueueSystem.js";
19+
import { cycleAnchorSnapshotId } from "./waitpointSystem.js";
1920

2021
/** The description of the EXECUTING snapshot `continueRunExecution` writes. */
2122
export const CONTINUED_AFTER_SUSPENDED_DESCRIPTION = "Run was continued after being suspended";
@@ -24,6 +25,7 @@ export type CheckpointSystemOptions = {
2425
resources: SystemResources;
2526
executionSnapshotSystem: ExecutionSnapshotSystem;
2627
enqueueSystem: EnqueueSystem;
28+
waitpointStoreEnabled?: (organizationId: string) => boolean | Promise<boolean>;
2729
};
2830

2931
export class CheckpointSystem {
@@ -240,13 +242,28 @@ export class CheckpointSystem {
240242
checkpoint: taskRunCheckpoint,
241243
} satisfies CreateCheckpointResult;
242244
} else {
245+
// An unwired capability check returns false synchronously.
246+
const storeEnabled =
247+
(snapshot.executionStatus === "EXECUTING_WITH_WAITPOINTS" ||
248+
snapshot.executionStatus === "SUSPENDED") &&
249+
this.options.waitpointStoreEnabled?.(env.organizationId);
243250
//create a new execution snapshot, with the checkpoint
244251
const newSnapshot = await this.executionSnapshotSystem.createExecutionSnapshot(prisma, {
245252
run,
246253
snapshot: {
247254
executionStatus: "SUSPENDED",
248255
description: "Run was suspended after creating a checkpoint.",
249-
metadata: snapshot.metadata,
256+
metadata:
257+
storeEnabled && (await storeEnabled) === true
258+
? {
259+
...(snapshot.metadata &&
260+
typeof snapshot.metadata === "object" &&
261+
!Array.isArray(snapshot.metadata)
262+
? snapshot.metadata
263+
: {}),
264+
waitpointCycleAnchor: cycleAnchorSnapshotId(snapshot),
265+
}
266+
: snapshot.metadata,
250267
},
251268
previousSnapshotId: snapshot.id,
252269
previousSnapshot: snapshot,

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ export function createWaitpointResultRunStore(
2929
requestedIds: ids,
3030
legacyReader: async () => [],
3131
storeReader: async (requested) => {
32-
const entries = await options.store.readCompletedRecords(requested);
32+
const entries = await options.store.readCompletedRecords(requested, context.runId);
3333
for (const [id, entry] of entries) {
3434
if (entry.record.environmentId !== context.environmentId) {
3535
throw new UnresolvableCompletedWaitpointRecordsError(

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

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -331,6 +331,8 @@ describe("WaitpointSystem create/block write routing", () => {
331331
store: wired ? destination : undefined,
332332
isStoreEnabled: isEnabled,
333333
mintPolicy: { resolveMintMode: () => "redis", isCapabilityReady: () => true },
334+
withLegacyRegistrationLock: (runId, operation) =>
335+
resources.runLock.lock("routing-test", [runId], operation),
334336
});
335337
const system = new WaitpointSystem({
336338
...engine.waitpointSystem["options"],
@@ -342,8 +344,9 @@ describe("WaitpointSystem create/block write routing", () => {
342344
}
343345
: undefined,
344346
});
345-
const expectOneFlagCheck = (before: number) => {
346-
expect(checkedOrganizations.slice(before)).toEqual(wired ? [organizationId] : []);
347+
const expectFlagChecks = (before: number, enabledCount = 1) => {
348+
const count = !wired ? 0 : mode === "enabled" ? enabledCount : 1;
349+
expect(checkedOrganizations.slice(before)).toEqual(Array(count).fill(organizationId));
347350
};
348351
for (const native of [false, true]) {
349352
const id = native ? generateWaitpointId("DATETIME") : WaitpointId.generate().id;
@@ -372,23 +375,20 @@ describe("WaitpointSystem create/block write routing", () => {
372375
let beforeFlagChecks = checkedOrganizations.length;
373376
if (mode === "enabled" && native) {
374377
await expect(router.registerBlocks(params)).rejects.toThrow("three-step barrier");
375-
expectOneFlagCheck(beforeFlagChecks);
378+
expectFlagChecks(beforeFlagChecks);
376379
beforeFlagChecks = checkedOrganizations.length;
377380
await expect(router.registerBlocksLockless(params)).rejects.toThrow(
378381
"three-step barrier"
379382
);
380-
expectOneFlagCheck(beforeFlagChecks);
383+
expectFlagChecks(beforeFlagChecks);
381384
expect(legacyStore.calls).toEqual([]);
382385
} else {
383-
expect(await router.registerBlocks(params)).toEqual(
384-
mode === "enabled"
385-
? { pendingCount: 1, storeRoutingEnabled: true }
386-
: { pendingCount: 1 }
387-
);
388-
expectOneFlagCheck(beforeFlagChecks);
386+
expect(await router.registerBlocks(params)).toEqual({ pendingCount: 1 });
387+
// The cleanup fence has its own independent gate after registration's routing gate.
388+
expectFlagChecks(beforeFlagChecks, 2);
389389
beforeFlagChecks = checkedOrganizations.length;
390390
await router.registerBlocksLockless(params);
391-
expectOneFlagCheck(beforeFlagChecks);
391+
expectFlagChecks(beforeFlagChecks, 3);
392392
// Direct delegation preserves the original collection, including duplicates.
393393
if (mode !== "enabled") expect(legacyStore.blockInputs.at(-1)).toBe(ids);
394394
expect(legacyStore.calls).toEqual([
@@ -404,7 +404,8 @@ describe("WaitpointSystem create/block write routing", () => {
404404
projectId: env.projectId,
405405
organizationId,
406406
});
407-
expectOneFlagCheck(beforeFlagChecks);
407+
// System admission, absorption, and cleanup-fence gates remain separate stages.
408+
expectFlagChecks(beforeFlagChecks, 3);
408409
expect(destinationStore.calls).toEqual([]);
409410
}
410411
// Ensure a fan-out edge even when enabled registration correctly rejected this parent.
@@ -432,16 +433,16 @@ describe("WaitpointSystem create/block write routing", () => {
432433
if (!nativeDispatch) {
433434
// Complete-before-register uses the same gated payload on the immediate unblock path.
434435
beforeFlagChecks = checkedOrganizations.length;
435-
// A flag change after the check affects the next operation, not this registration's
436-
// route or derived continuation payload. There is no decision cache across operations.
436+
// A flag change after system admission is rechecked by absorption. The system still
437+
// carries its admitted org onto the continuation; the next operation evaluates afresh.
437438
disableAfterNextCheck = mode === "enabled";
438439
await system.blockRunWithWaitpoint({
439440
runId: run.id,
440441
waitpoints: ids,
441442
projectId: env.projectId,
442443
organizationId,
443444
});
444-
expectOneFlagCheck(beforeFlagChecks);
445+
expectFlagChecks(beforeFlagChecks, 2);
445446
const immediate = await resources.worker.getJob(`continueRunIfUnblocked:${run.id}`);
446447
expect(immediate?.item).toEqual(
447448
mode === "enabled" ? { runId: run.id, organizationId } : { runId: run.id }
@@ -450,7 +451,7 @@ describe("WaitpointSystem create/block write routing", () => {
450451
if (mode === "enabled") {
451452
beforeFlagChecks = checkedOrganizations.length;
452453
await system.completeWaitpoint({ id, organizationId });
453-
expectOneFlagCheck(beforeFlagChecks);
454+
expectFlagChecks(beforeFlagChecks);
454455
const disabledJob = await resources.worker.getJob(
455456
`continueRunIfUnblocked:${run.id}`
456457
);
@@ -463,7 +464,7 @@ describe("WaitpointSystem create/block write routing", () => {
463464
beforeFlagChecks = checkedOrganizations.length;
464465
const direct = await router.complete({ waitpointId: id, organizationId });
465466
expect(direct.waitpoint.status).toBe("COMPLETED");
466-
expectOneFlagCheck(beforeFlagChecks);
467+
expectFlagChecks(beforeFlagChecks);
467468
const missing = generateWaitpointId("DATETIME");
468469
legacyStore.calls.length = 0;
469470
destinationStore.calls.length = 0;
@@ -489,7 +490,7 @@ describe("WaitpointSystem create/block write routing", () => {
489490
completedAfter: new Date(),
490491
mintContext: { organizationId, parentResidency: "redis-primary" },
491492
});
492-
expectOneFlagCheck(beforeFlagChecks);
493+
expectFlagChecks(beforeFlagChecks);
493494
expect(mode === "enabled" ? legacyStore.calls : destinationStore.calls).toEqual([]);
494495
expect(mode === "enabled" ? destinationStore.calls : legacyStore.calls).toEqual([
495496
"upsertWaitpoint",

0 commit comments

Comments
 (0)