Skip to content

Commit 2dc02f0

Browse files
kraftpkraftp
authored
Improve Cancellation Semantics (#1264)
Cancellation should be a terminal state. Eliminate a race condition where a workflow cancelled right before it completed would transition to "SUCCESS" or "ERROR". Now, workflows can only leave CANCELLED by being explicitly resumed.
1 parent 02585b8 commit 2dc02f0

5 files changed

Lines changed: 112 additions & 8 deletions

File tree

‎.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ packages/**/package-lock.json
99
CLAUDE.md
1010
tests/temp-dist/
1111
tests/temp-tests/
12+
.claude/
1213

1314
# Logs
1415
logs

‎src/system_database.ts‎

Lines changed: 37 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -799,9 +799,7 @@ export class SystemDatabase {
799799
async recordWorkflowOutput(workflowID: string, status: WorkflowStatusInternal): Promise<void> {
800800
const client = await this.pool.connect();
801801
try {
802-
await this.updateWorkflowStatus(client, workflowID, StatusString.SUCCESS, {
803-
update: { output: status.output, resetDeduplicationID: true, setCompletedAt: true },
804-
});
802+
await this.#recordWorkflowOutcome(client, workflowID, StatusString.SUCCESS, { output: status.output });
805803
} finally {
806804
client.release();
807805
}
@@ -811,14 +809,42 @@ export class SystemDatabase {
811809
async recordWorkflowError(workflowID: string, status: WorkflowStatusInternal): Promise<void> {
812810
const client = await this.pool.connect();
813811
try {
814-
await this.updateWorkflowStatus(client, workflowID, StatusString.ERROR, {
815-
update: { error: status.error, resetDeduplicationID: true, setCompletedAt: true },
816-
});
812+
await this.#recordWorkflowOutcome(client, workflowID, StatusString.ERROR, { error: status.error });
817813
} finally {
818814
client.release();
819815
}
820816
}
821817

818+
// Record a workflow's terminal outcome (SUCCESS or ERROR), but never overwrite
819+
// the terminal CANCELLED status: a workflow can be cancelled during its final
820+
// step, and if so it must not be able to subsequently complete. If the
821+
// workflow is cancelled, abort the function so it does not complete. This
822+
// mirrors the cancellation check done before each step.
823+
async #recordWorkflowOutcome(
824+
client: PoolClient,
825+
workflowID: string,
826+
status: (typeof StatusString)[keyof typeof StatusString],
827+
outcome: { output?: string | null; error?: string | null },
828+
): Promise<void> {
829+
let cancelled = false;
830+
try {
831+
await client.query('BEGIN');
832+
await this.updateWorkflowStatus(client, workflowID, status, {
833+
update: { ...outcome, resetDeduplicationID: true, setCompletedAt: true },
834+
where: { notStatus: StatusString.CANCELLED },
835+
throwOnFailure: false,
836+
});
837+
cancelled = (await this.getWorkflowStatusValue(client, workflowID)) === StatusString.CANCELLED;
838+
await client.query('COMMIT');
839+
} catch (e) {
840+
await client.query('ROLLBACK');
841+
throw e;
842+
}
843+
if (cancelled) {
844+
throw new DBOSWorkflowCancelledError(workflowID);
845+
}
846+
}
847+
822848
async getPendingWorkflows(executorID: string, appVersion: string): Promise<GetPendingWorkflowsOutput[]> {
823849
const getWorkflows = await this.pool.query<workflow_status>(
824850
`SELECT workflow_uuid, queue_name
@@ -3675,6 +3701,7 @@ export class SystemDatabase {
36753701
};
36763702
where?: {
36773703
status?: (typeof StatusString)[keyof typeof StatusString];
3704+
notStatus?: (typeof StatusString)[keyof typeof StatusString];
36783705
};
36793706
throwOnFailure?: boolean;
36803707
} = {},
@@ -3739,6 +3766,10 @@ export class SystemDatabase {
37393766
const param = args.push(where.status);
37403767
whereClause += ` AND status=$${param}`;
37413768
}
3769+
if (where.notStatus) {
3770+
const param = args.push(where.notStatus);
3771+
whereClause += ` AND status!=$${param}`;
3772+
}
37423773

37433774
const result = await client.query<workflow_status>(
37443775
`UPDATE "${this.schemaName}".workflow_status ${setClause} ${whereClause}`,

‎tests/dbos.test.ts‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -348,6 +348,11 @@ describe('dbos-tests', () => {
348348
});
349349
const status = await DBOS.getWorkflowStatus(workflowID);
350350
expect(status?.status).toBe(StatusString.CANCELLED);
351+
// The child is cancelled by its own inherited deadline, independently of the parent.
352+
// Wait for it to settle before reading its status to avoid a race with that cancellation.
353+
await expect(DBOS.retrieveWorkflow(childID).getResult()).rejects.toThrow(
354+
new DBOSAwaitedWorkflowCancelledError(childID),
355+
);
351356
const childStatus = await DBOS.getWorkflowStatus(childID);
352357
expect(childStatus?.status).toBe(StatusString.CANCELLED);
353358
expect(status?.deadlineEpochMS).toBe(childStatus?.deadlineEpochMS);
@@ -363,6 +368,11 @@ describe('dbos-tests', () => {
363368
await expect(handle.getResult()).rejects.toThrow(new DBOSWorkflowCancelledError(workflowID));
364369
const status = await DBOS.getWorkflowStatus(workflowID);
365370
expect(status?.status).toBe(StatusString.CANCELLED);
371+
// The child is cancelled by its own inherited deadline, independently of the parent.
372+
// Wait for it to settle before reading its status to avoid a race with that cancellation.
373+
await expect(DBOS.retrieveWorkflow(childID).getResult()).rejects.toThrow(
374+
new DBOSAwaitedWorkflowCancelledError(childID),
375+
);
366376
const childStatus = await DBOS.getWorkflowStatus(childID);
367377
expect(childStatus?.status).toBe(StatusString.CANCELLED);
368378
expect(status?.deadlineEpochMS).toBe(childStatus?.deadlineEpochMS);
@@ -380,6 +390,11 @@ describe('dbos-tests', () => {
380390
});
381391
const status = await DBOS.getWorkflowStatus(workflowID);
382392
expect(status?.status).toBe(StatusString.CANCELLED);
393+
// The child is cancelled by its own inherited deadline, independently of the parent.
394+
// Wait for it to settle before reading its status to avoid a race with that cancellation.
395+
await expect(DBOS.retrieveWorkflow(childID).getResult()).rejects.toThrow(
396+
new DBOSAwaitedWorkflowCancelledError(childID),
397+
);
383398
const childStatus = await DBOS.getWorkflowStatus(childID);
384399
expect(childStatus?.status).toBe(StatusString.CANCELLED);
385400
expect(status?.deadlineEpochMS).toBe(childStatus?.deadlineEpochMS);
@@ -395,6 +410,11 @@ describe('dbos-tests', () => {
395410
await expect(handle.getResult()).rejects.toThrow(new DBOSWorkflowCancelledError(workflowID));
396411
const status = await DBOS.getWorkflowStatus(workflowID);
397412
expect(status?.status).toBe(StatusString.CANCELLED);
413+
// The child is cancelled by its own inherited deadline, independently of the parent.
414+
// Wait for it to settle before reading its status to avoid a race with that cancellation.
415+
await expect(DBOS.retrieveWorkflow(childID).getResult()).rejects.toThrow(
416+
new DBOSAwaitedWorkflowCancelledError(childID),
417+
);
398418
const childStatus = await DBOS.getWorkflowStatus(childID);
399419
expect(childStatus?.status).toBe(StatusString.CANCELLED);
400420
expect(status?.deadlineEpochMS).toBe(childStatus?.deadlineEpochMS);

‎tests/wfqueue.test.ts‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -621,9 +621,14 @@ describe('queued-wf-tests-simple', () => {
621621
});
622622
await expect(regularHandle.getResult()).resolves.toBeUndefined();
623623

624-
// Complete the blocked workflow
624+
// Unblock the cancelled workflow's body. Even though it now runs to
625+
// completion, CANCELLED is terminal: it must not overwrite CANCELLED with
626+
// SUCCESS, and awaiting its result must raise.
625627
TestCancelQueues.blockingEvent.set();
626628
await expect(blockedHandle.getResult()).rejects.toThrow(DBOSAwaitedWorkflowCancelledError);
629+
await expect(blockedHandle.getStatus()).resolves.toMatchObject({
630+
status: StatusString.CANCELLED,
631+
});
627632

628633
// Verify all queue entries eventually get cleaned up
629634
expect(await queueEntriesAreCleanedUp()).toBe(true);

‎tests/workflow_management.test.ts‎

Lines changed: 48 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -396,6 +396,32 @@ describe('workflow-management-tests', () => {
396396
expect(result.rows[0].status).toBe(StatusString.SUCCESS);
397397
});
398398

399+
test('test-cancel-after-final-step', async () => {
400+
// A workflow cancelled after its final step completes (but before it
401+
// finishes) must not be able to complete successfully. CANCELLED is terminal.
402+
TestEndpoints.stepsCompleted = 0;
403+
const input = 5;
404+
const wfid = randomUUID();
405+
406+
const cancelledHandle = await DBOS.startWorkflow(TestEndpoints, { workflowID: wfid }).cancelAfterFinalStepWorkflow(
407+
input,
408+
);
409+
await TestEndpoints.mainThreadEvent.wait();
410+
await DBOS.cancelWorkflow(wfid);
411+
TestEndpoints.workflowEvent.set();
412+
413+
// The workflow must not complete successfully.
414+
await expect(cancelledHandle.getResult()).rejects.toThrow(DBOSWorkflowCancelledError);
415+
expect(TestEndpoints.stepsCompleted).toBe(1);
416+
await expect(DBOS.getWorkflowStatus(wfid)).resolves.toMatchObject({ status: StatusString.CANCELLED });
417+
418+
// Resuming it should let it complete successfully.
419+
const handle = await DBOS.resumeWorkflow<number>(wfid);
420+
await expect(handle.getResult()).resolves.toBe(input);
421+
await expect(DBOS.getWorkflowStatus(wfid)).resolves.toMatchObject({ status: StatusString.SUCCESS });
422+
expect(TestEndpoints.stepsCompleted).toBe(1); // cancelStep was already recorded, not re-run
423+
});
424+
399425
test('getworkflows-with-completed-at', async () => {
400426
// Successful workflow gets completedAt set.
401427
const beforeSuccess = new Date().toISOString();
@@ -515,6 +541,26 @@ describe('workflow-management-tests', () => {
515541
static async stepOne() {
516542
return Promise.resolve();
517543
}
544+
545+
static stepsCompleted = 0;
546+
static workflowEvent = new Event();
547+
static mainThreadEvent = new Event();
548+
549+
@DBOS.step()
550+
static async cancelStep() {
551+
TestEndpoints.stepsCompleted += 1;
552+
return Promise.resolve();
553+
}
554+
555+
@DBOS.workflow()
556+
static async cancelAfterFinalStepWorkflow(x: number) {
557+
// The only step runs and records its output...
558+
await TestEndpoints.cancelStep();
559+
// ...then the workflow is cancelled before it returns.
560+
TestEndpoints.mainThreadEvent.set();
561+
await TestEndpoints.workflowEvent.wait();
562+
return x;
563+
}
518564
}
519565
});
520566

@@ -2042,7 +2088,8 @@ describe('test-fork', () => {
20422088
await ResumeForkQueueWorkflow.step1Started.wait();
20432089
await DBOS.cancelWorkflow(wfid);
20442090
ResumeForkQueueWorkflow.step1Gate.set();
2045-
await expect(handle.getResult()).rejects.toThrow();
2091+
await expect(handle.getResult()).rejects.toThrow(DBOSAwaitedWorkflowCancelledError);
2092+
await expect(DBOS.getWorkflowStatus(wfid)).resolves.toMatchObject({ status: StatusString.CANCELLED });
20462093
expect(ResumeForkQueueWorkflow.stepOneCount).toBe(1);
20472094
expect(ResumeForkQueueWorkflow.stepTwoCount).toBe(0);
20482095

0 commit comments

Comments
 (0)