fix(orchestrator): keep heartbeating until workers have drained

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

packages/pgconductor-js/src/orchestrator.ts

@@ -47,6 +47,7 @@ export class Orchestrator {
 	private readonly logger: Logger;
 
 	private heartbeats: Promise<void> | null = null;
+	private heartbeatController: AbortController | null = null;
 	private _stopDeferred: Deferred<void> | null = null;
 	private _startDeferred: Deferred<void> | 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<void>[] = [];
 				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,24 +314,24 @@ 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<void> {
+	private async runHeartbeats(signal: AbortSignal): Promise<void> {
 		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) {
 					await this.db.recoverStaleOrchestrators(
 						{
 							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<void> {
 		// 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;

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<void>();
+		const finishHandler = new Deferred<void>();
+		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();
+	},
+);