Skip to content

Backport #3575: fix(world-postgres): make step creation atomic #3746

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
github-actions wants to merge 1 commit into
base: stable
Choose a base branch
from
Open

Backport #3575: fix(world-postgres): make step creation atomic #3746

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions .changeset/atomic-step-created-event.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
---
'@workflow/world-postgres': patch
---

Commit step entities and their `step_created` events atomically.
133 changes: 97 additions & 36 deletions packages/world-postgres/src/storage.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -870,31 +870,6 @@ export function createEventsStorage(drizzle: Drizzle): Storage['events'] {
// skipped when this is already set.
let value: { createdAt: Date } | undefined;

// Handle step_created event: create step entity
if (data.eventType === 'step_created') {
const eventData = (data as any).eventData as {
stepName: string;
input: any;
};
const [stepValue] = await drizzle
.insert(Schema.steps)
.values({
runId: effectiveRunId,
stepId: data.correlationId!,
stepName: eventData.stepName,
input: eventData.input as SerializedContent,
status: 'pending',
attempt: 0,
// Propagate specVersion from the event to the step entity
specVersion: effectiveSpecVersion,
})
.onConflictDoNothing()
.returning();
if (stepValue) {
step = deserializeStepError(compact(stepValue));
}
}

// Handle step_started event: increment attempt and set the step to
// running, then write the matching event log entry in the same
// transaction. The guarded UPDATE takes the step row lock; keeping the
Expand Down Expand Up @@ -1465,17 +1440,103 @@ export function createEventsStorage(drizzle: Drizzle): Storage['events'] {

try {
if (!value) {
[value] = await drizzle
.insert(events)
.values({
runId: effectiveRunId,
eventId: getEventId(),
correlationId: data.correlationId,
eventType: data.eventType,
eventData: storedEventData,
specVersion: effectiveSpecVersion,
})
.returning({ createdAt: events.createdAt });
// Handle step_created event: create the step entity and append its
// event log entry in ONE transaction. Committed separately (as this
// used to be), a crash between the two writes leaves a step row with
// no matching `step_created` event, which no replay can reconstruct.
if (data.eventType === 'step_created') {
const eventData = (data as any).eventData as {
stepName: string;
input: any;
};
const created = await drizzle.transaction(async (tx) => {
let [stepValue] = await tx
.insert(Schema.steps)
.values({
runId: effectiveRunId,
stepId: data.correlationId!,
stepName: eventData.stepName,
input: eventData.input as SerializedContent,
status: 'pending',
attempt: 0,
// Propagate specVersion from the event to the step entity
specVersion: effectiveSpecVersion,
})
.onConflictDoNothing()
.returning();
if (!stepValue) {
const [existingEvent] = await tx
.select({ eventId: events.eventId })
.from(events)
.where(
and(
eq(events.runId, effectiveRunId),
eq(events.correlationId, data.correlationId!),
eq(events.eventType, 'step_created')
)
)
.limit(1);
if (existingEvent) {
throw new EntityConflictError(
`step_created for correlationId "${data.correlationId}" already exists in run "${effectiveRunId}"`
);
}

// A row without its matching event was left by the old
// non-transactional path. Keep the row and complete the
// missing event inside this transaction so existing orphans
// remain recoverable while new partial writes cannot escape.
[stepValue] = await tx
.select()
.from(Schema.steps)
.where(
and(
eq(Schema.steps.runId, effectiveRunId),
eq(Schema.steps.stepId, data.correlationId!)
)
)
.limit(1);
if (!stepValue) {
throw new EntityConflictError(
`step_created for correlationId "${data.correlationId}" already exists in run "${effectiveRunId}"`
);
}
}

const [eventValue] = await tx
.insert(events)
.values({
runId: effectiveRunId,
eventId: getEventId(),
correlationId: data.correlationId,
eventType: data.eventType,
eventData: storedEventData,
specVersion: effectiveSpecVersion,
})
.returning({ createdAt: events.createdAt });
if (!eventValue) {
throw new EntityConflictError(
`step_created for run "${effectiveRunId}" could not be created`
);
}
return { eventValue, stepValue };
});

step = deserializeStepError(compact(created.stepValue));
value = created.eventValue;
} else {
[value] = await drizzle
.insert(events)
.values({
runId: effectiveRunId,
eventId: getEventId(),
correlationId: data.correlationId,
eventType: data.eventType,
eventData: storedEventData,
specVersion: effectiveSpecVersion,
})
.returning({ createdAt: events.createdAt });
}
}
} catch (err) {
// Translate unique-violation on the entity-creation partial index
Expand Down
100 changes: 99 additions & 1 deletion packages/world-postgres/test/storage.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1245,7 +1245,105 @@ describe('Storage (Postgres integration)', () => {
stepName: 'test-step',
input: new Uint8Array(),
})
).rejects.toMatchObject({ name: 'EntityConflictError' });
).rejects.toMatchObject({
name: 'EntityConflictError',
message: `step_created for correlationId "step_seq_dup" already exists in run "${testRunId}"`,
});
});

it('recovers an orphaned step row before a plain step_started event', async () => {
const stepId = 'step_orphan';
await drizzle.insert(DrizzleSchema.steps).values({
runId: testRunId,
stepId,
stepName: 'test-step',
input: new Uint8Array(),
status: 'pending',
attempt: 0,
specVersion: SPEC_VERSION_CURRENT,
});

const recovered = await createStep(events, testRunId, {
stepId,
stepName: 'test-step',
input: new Uint8Array(),
});
expect(recovered.stepId).toBe(stepId);

await updateStep(events, testRunId, stepId, 'step_started');

const evts = await events.list({
runId: testRunId,
pagination: {},
});
expect(
evts.data
.filter((event) => event.correlationId === stepId)
.map((event) => event.eventType)
).toEqual(['step_created', 'step_started']);
});

it('rolls back the step entity when the matching event insert fails', async () => {
await pool.query(`
CREATE FUNCTION workflow.reject_step_created_event_for_test()
RETURNS trigger
LANGUAGE plpgsql
AS $$
BEGIN
IF NEW.type = 'step_created'
AND NEW.correlation_id = 'step_partial_write'
THEN
RAISE EXCEPTION 'forced step_created event insert failure';
END IF;
RETURN NEW;
END;
$$;

CREATE TRIGGER reject_step_created_event_for_test
BEFORE INSERT ON workflow.workflow_events
FOR EACH ROW
EXECUTE FUNCTION workflow.reject_step_created_event_for_test();
`);

try {
await expect(
createStep(events, testRunId, {
stepId: 'step_partial_write',
stepName: 'test-step',
input: new Uint8Array(),
})
).rejects.toMatchObject({
cause: {
message: expect.stringMatching(
/forced step_created event insert failure/
),
},
});

const stepRows = await drizzle
.select({ stepId: DrizzleSchema.steps.stepId })
.from(DrizzleSchema.steps)
.where(eq(DrizzleSchema.steps.stepId, 'step_partial_write'));
expect(stepRows).toEqual([]);

const evts = await events.list({
runId: testRunId,
pagination: {},
});
expect(
evts.data.filter(
(event) =>
event.eventType === 'step_created' &&
event.correlationId === 'step_partial_write'
)
).toHaveLength(0);
} finally {
await pool.query(`
DROP TRIGGER reject_step_created_event_for_test
ON workflow.workflow_events;
DROP FUNCTION workflow.reject_step_created_event_for_test();
`);
}
});

it('should reject duplicate wait_created with EntityConflictError', async () => {
Expand Down