From 3447b115b586e28045c585a05fba621f31506140 Mon Sep 17 00:00:00 2001 From: psteinroe Date: Fri, 2 Oct 2026 17:03:10 +0000 Subject: [PATCH] fix(orchestrator): keep heartbeating until workers have drained --- AGENTS.md | 1 + packages/pgconductor-js/src/orchestrator.ts | 29 ++++--- .../integration/orchestrator-cleanup.test.ts | 82 +++++++++++++++++++ 3 files changed, 99 insertions(+), 13 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 5ff401f..054e4b6 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -143,6 +143,7 @@ Tables live in the `pgconductor` schema with a `_private_` prefix (`_private_exe **Orchestrator** (`orchestrator.ts`) - Manages multiple workers, plus the internal event-dispatch worker when events are configured - Handles startup/shutdown coordination and heartbeats +- Heartbeats keep running while workers drain after shutdown and stop in cleanup, so a slow handler is not recovered as stale - Every 8th heartbeat recovers stale orchestrators (unlocks their executions); an orchestrator whose row was recovered stops itself when its next heartbeat has to re-insert the row - Provides `stopped` promise for graceful shutdown - Runs the internal event dispatch worker only when the conductor has `events` configured diff --git a/packages/pgconductor-js/src/orchestrator.ts b/packages/pgconductor-js/src/orchestrator.ts index 7734408..2dd848d 100644 --- a/packages/pgconductor-js/src/orchestrator.ts +++ b/packages/pgconductor-js/src/orchestrator.ts @@ -47,6 +47,7 @@ export class Orchestrator { private readonly logger: Logger; private heartbeats: Promise | null = null; + private heartbeatController: AbortController | null = null; private _stopDeferred: Deferred | null = null; private _startDeferred: Deferred | null = null; private _abortController: AbortController | null = null; @@ -225,7 +226,8 @@ export class Orchestrator { throw new Error("Received shutdown signal during startup"); } - this.heartbeats = this.runHeartbeats(this.abortController); + this.heartbeatController = new AbortController(); + this.heartbeats = this.runHeartbeats(this.heartbeatController.signal); const workerLifecycles: Promise[] = []; for (const worker of this.workers) { @@ -286,8 +288,8 @@ export class Orchestrator { /** * Stop the orchestrator gracefully: - * 1. Stop heartbeat - * 2. Stop all workers + * 1. Stop all workers + * 2. Stop heartbeat * 3. Clean up resources * * Waits until the orchestrator is fully stopped @@ -312,16 +314,16 @@ export class Orchestrator { } /** - * Run the heartbeat loop until the orchestrator aborts: + * Run the heartbeat loop until cleanup stops it, so the orchestrator stays live while workers drain: * - Updates last_heartbeat_at every heartbeat * - Checks for version mismatch shutdowns from database * - Stops when another orchestrator recovered us as stale * - Recovers stale orchestrators periodically (every 8th heartbeat) */ - private async runHeartbeats(abortController: AbortController): Promise { + private async runHeartbeats(signal: AbortSignal): Promise { for (let heartbeatCount = 1; ; heartbeatCount++) { - await waitFor(HEARTBEAT_INTERVAL_MS, { signal: abortController.signal }); - if (abortController.signal.aborted) return; + await waitFor(HEARTBEAT_INTERVAL_MS, { signal }); + if (signal.aborted) return; try { if (heartbeatCount % 8 === 0) { @@ -329,7 +331,7 @@ export class Orchestrator { { maxAge: `${STALE_ORCHESTRATOR_MAX_AGE_MS} milliseconds`, }, - { signal: abortController.signal }, + { signal }, ); } @@ -340,13 +342,13 @@ export class Orchestrator { version: PACKAGE_VERSION, migrationNumber: this.migrationStore.getLatestMigrationNumber(), }, - { signal: abortController.signal }, + { signal }, ); // Our row was gone, so another orchestrator recovered our executions if (signals.some((s) => s.registered)) { this.logger.warn("Orchestrator was recovered as stale, shutting down"); - abortController.abort(); + this.abortController.abort(); return; } @@ -356,11 +358,11 @@ export class Orchestrator { switch (signal.signal_type) { case "shutdown": - if (!abortController.signal.aborted) { + if (!this.isShuttingDown) { this.logger.info( `Received shutdown signal: ${signal.signal_payload?.reason || "unknown"}`, ); - abortController.abort(); + this.abortController.abort(); } break; @@ -475,9 +477,10 @@ export class Orchestrator { */ private async cleanup(): Promise { // Stop heartbeat and wait for one in flight - this.abortController.abort(); + this.heartbeatController?.abort(); await this.heartbeats; this.heartbeats = null; + this.heartbeatController = null; // Remove ourselves from orchestrators table and release locked executions const cleanupSignal = new AbortController().signal; diff --git a/packages/pgconductor-js/tests/integration/orchestrator-cleanup.test.ts b/packages/pgconductor-js/tests/integration/orchestrator-cleanup.test.ts index be463e1..99524f8 100644 --- a/packages/pgconductor-js/tests/integration/orchestrator-cleanup.test.ts +++ b/packages/pgconductor-js/tests/integration/orchestrator-cleanup.test.ts @@ -380,3 +380,85 @@ test("an orchestrator recovered as stale aborts its running handlers and stops", expect(handlerAborted.isSettled).toBe(true); expect(orch.isStopped).toBe(true); }); + +test.each(["stop()", "a shutdown signal"])( + "an orchestrator stopped by %s keeps heartbeating until its handlers finish", + async (trigger) => { + const db = await pool.child(); + databases.push(db); + + const taskDefinition = defineTask({ name: "slow-task" }); + const conductor = Conductor.create({ + sql: db.sql, + tasks: TaskSchemas.fromSchema([taskDefinition]), + context: {}, + }); + await conductor.ensureInstalled(); + + const handlerStarted = new Deferred(); + const finishHandler = new Deferred(); + const task = conductor.createTask(taskDefinition, { invocable: true }, async () => { + handlerStarted.resolve(); + await finishHandler.promise; + }); + + await db.client.setFakeTime({ date: new Date("2024-01-01T12:00:00Z") }); + + jest.useFakeTimers(); + const orch = Orchestrator.create({ conductor, tasks: [task] }); + await orch.start(); + await conductor.invoke(taskDefinition, {}); + + while (!handlerStarted.isSettled) { + jest.advanceTimersByTime(1000); + await db.sql`select 1`; + } + + if (trigger === "stop()") { + orch.stop(); + } else { + await db.sql` + insert into pgconductor._private_orchestrator_signals (orchestrator_id, type, payload) + values (${orch.info.id}::uuid, 'shutdown', '{}'::jsonb) + `; + } + for (let seconds = 0; seconds < 30; seconds++) { + jest.advanceTimersByTime(1000); + await db.sql`select 1`; + } + expect(orch.isShuttingDown).toBe(true); + + await db.client.setFakeTime({ date: new Date("2024-01-01T12:10:00Z") }); + for (let seconds = 0; seconds < 60; seconds++) { + jest.advanceTimersByTime(1000); + await db.sql`select 1`; + } + + const [row] = await db.sql<{ last_heartbeat_at: Date }[]>` + select last_heartbeat_at from pgconductor._private_orchestrators where id = ${orch.info.id}::uuid + `; + expect(row?.last_heartbeat_at).toEqual(new Date("2024-01-01T12:10:00Z")); + + await db.client.recoverStaleOrchestrators({ maxAge: "5 minutes" }); + const locked = await db.sql` + select id from pgconductor._private_executions + where task_key = 'slow-task' and locked_by = ${orch.info.id}::uuid + `; + expect(locked.length).toBe(1); + + finishHandler.resolve(); + await orch.stopped; + + const rows = await db.sql` + select id from pgconductor._private_orchestrators where id = ${orch.info.id}::uuid + `; + expect(rows.length).toBe(0); + const [execution] = await db.sql<{ completed: boolean }[]>` + select completed_at is not null as completed from pgconductor._private_executions + where task_key = 'slow-task' + `; + expect(execution?.completed).toBe(true); + + await db.client.clearFakeTime(); + }, +);