cherry-pick: feat(threads): graceful drain hook for in-flight work before worker shutdown (conflicts → v5.1)

components/shutdownDrain.ts

@@ -0,0 +1,117 @@
+/**
+ * Per-worker registry of graceful "drain before shutdown" hooks.
+ *
+ * On a worker restart (deploy reload, `harper dev` reload, rolling restart) the worker is told to shut
+ * down while it may still be doing work that is cheaper to finish than to interrupt — notably a
+ * replication blob *send* streaming to a peer. Tearing that down mid-stream makes the peer receiver
+ * treat the blob as diverged until it is re-requested. A registered drain lets that work reach a safe
+ * point (finish, or stall) before the worker stops accepting connections and exits.
+ *
+ * Core stays deliberately generic here — it knows nothing about blobs or replication. A component
+ * (harper-pro replication) registers a {@link ShutdownDrain}; the worker shutdown path
+ * ({@link server/threads/threadServer}) awaits {@link runShutdownDrains} before `closeServers`, and
+ * extends its termination backstops (see manageThreads `extendShutdownDeadline`) only while
+ * {@link shutdownDrainsHaveWork} reports real in-flight work — so a worker that hangs for an unrelated
+ * reason is still force-killed on the normal short timeout.
+ *
+ * State is module-local, so it is naturally per-worker (each worker is a fresh module realm).
+ */
+import harperLogger from '../utility/logging/harper_logger.ts';
+import * as env from '../utility/environment/environmentManager.ts';
+import { CONFIG_PARAMS } from '../utility/hdbTerms.ts';
+
+/** Absolute cap (ms) on how long a drain may hold the worker's shutdown open. */
+const DEFAULT_DRAIN_CEILING_MS = 600_000; // 10 minutes
+/** Largest delay a Node timer accepts; a larger value overflows and fires ~immediately. */
+export const MAX_TIMER_MS = 2_147_483_647; // 2^31 - 1
+
+/**
+ * The configured drain ceiling in ms. Config values can arrive as strings from YAML / HARPER_CONFIG,
+ * so coerce; but treat blank/null/empty as "unset" and fall back to the default (a blank config must
+ * not read as `Number('') === 0`, which would silently disable the drain). Clamp to the max timer
+ * value so a huge misconfig can't overflow setTimeout. Used both to size the drain deadline and to cap
+ * the main thread's force-terminate extension. An explicit `0` still disables draining.
+ */
+export function getShutdownDrainCeilingMs(): number {
+	const raw = env.get(CONFIG_PARAMS.REPLICATION_BLOBSENDDRAINTIMEOUT);
+	if (raw === null || raw === undefined || String(raw).trim() === '') return DEFAULT_DRAIN_CEILING_MS;
+	const ms = Number(raw);
+	return Number.isFinite(ms) && ms >= 0 ? Math.min(ms, MAX_TIMER_MS) : DEFAULT_DRAIN_CEILING_MS;
+}
+
+/**
+ * Compute a force-terminate timer delay for a worker-requested drain deadline. Clamps the requested
+ * absolute `deadlineMs` to `now + ceilingMs` (so a buggy/rogue worker message can't defer the kill
+ * unboundedly) and to a finite value (a non-finite deadline falls back to `now`, i.e. no extension),
+ * then adds `baseMs` of normal shutdown headroom. A shrink (e.g. the drain-done reset posting a
+ * now-deadline) passes through untouched. Pure so the clamp/guard arithmetic is directly testable.
+ *
+ * The final result is clamped to `MAX_TIMER_MS` too, not just `ceilingMs` on its own — `ceilingMs` is
+ * itself capped at `MAX_TIMER_MS` (see `getShutdownDrainCeilingMs`), but adding `baseMs` headroom on
+ * top can push the sum back over the limit at the extreme end of the configured ceiling, silently
+ * defeating the overflow guard the ceiling clamp was meant to provide.
+ */
+export function boundedTerminateDelay(deadlineMs: number, now: number, baseMs: number, ceilingMs: number): number {
+	const target = Number.isFinite(deadlineMs) ? deadlineMs : now;
+	const bounded = Math.min(target, now + ceilingMs);
+	return Math.min(Math.max(0, bounded - now) + baseMs, MAX_TIMER_MS);
+}
+
+export interface ShutdownDrain {
+	/** Synchronous, cheap: is there in-flight work worth draining right now? Drives deadline extension. */
+	hasWork(): boolean;
+	/**
+	 * Resolve once this hook's in-flight work has finished, stalled, or the absolute deadline has
+	 * passed. `deadlineMs` is an epoch timestamp (same clock as `Date.now()`); the hook must not run
+	 * past it.
+	 */
+	drain(deadlineMs: number): Promise<void>;
+}
+
+const drains = new Set<ShutdownDrain>();
+
+/** Register a drain hook. Returns an unregister function. */
+export function registerShutdownDrain(drain: ShutdownDrain): () => void {
+	drains.add(drain);
+	return () => drains.delete(drain);
+}
+
+/** Whether any registered hook currently has in-flight work worth draining (and extending the backstops for). */
+export function shutdownDrainsHaveWork(): boolean {
+	for (const drain of drains) {
+		try {
+			if (drain.hasWork()) return true;
+		} catch (error) {
+			harperLogger.error('Error checking shutdown drain for work', error);
+		}
+	}
+	return false;
+}
+
+/**
+ * Run every registered drain, resolving when all have settled or the absolute deadline is reached —
+ * whichever comes first. Never rejects: a drain that throws or rejects is logged and treated as
+ * settled, and a drain that ignores the deadline is abandoned (the process is exiting immediately
+ * after), so this can never wedge the shutdown sequence that follows.
+ */
+export async function runShutdownDrains(deadlineMs: number): Promise<void> {
+	if (drains.size === 0) return;
+	const settled = [...drains].map((drain) =>
+		Promise.resolve()
+			.then(() => drain.drain(deadlineMs))
+			.catch((error) => harperLogger.error('Error draining before shutdown', error))
+	);
+	const remaining = Math.max(0, deadlineMs - Date.now());
+	let timer: NodeJS.Timeout | undefined;
+	await Promise.race([
+		Promise.all(settled),
+		new Promise<void>((resolve) => {
+			timer = setTimeout(resolve, remaining);
+			timer.unref();
+		}),
+	])
+		.catch((error) => harperLogger.error('Error running shutdown drains', error))
+		.finally(() => {
+			if (timer) clearTimeout(timer);
+		});
+}

server/threads/manageThreads.js

@@ -57,9 +57,64 @@ const ADDED_PORT = 'added-port';
 const ACKNOWLEDGEMENT = 'ack';
 const REMOVE_PORT = 'remove-port';
 const FORCE_EXIT = 'force-exit';
+// Worker -> main request to push out the force-terminate backstop while the worker gracefully drains
+// in-flight work (e.g. replication blob sends) before shutdown. Carries an absolute epoch deadline.
+const EXTEND_SHUTDOWN_DEADLINE = 'extend-shutdown-deadline';
 let getThreadInfo;
+// Worker-side backstop that force-exits if the graceful shutdown sequence doesn't finish in time.
+let selfExitTimer;
+// An extended self-exit deadline requested by a drain (absolute epoch ms), honored regardless of whether
+// it is recorded before or after the SHUTDOWN handler arms the timer (the two race across listeners).
+let selfExitDrainDeadline = 0;
 _assignPackageExport('threads', connectedPorts);
 
+// Worker-side: (re)arm the self-exit backstop `delay` ms out, but never earlier than a drain deadline
+// already requested via extendShutdownDeadline (so ordering between the SHUTDOWN handler and the drain
+// extension doesn't matter).
+function armSelfExit(delay) {
+	if (selfExitDrainDeadline) {
+		// selfExitDrainDeadline is already clamped to the configured ceiling (see boundedTerminateDelay),
+		// but adding threadTerminationTimeout headroom on top can push the sum back over the max timer
+		// value at the extreme end of that ceiling, silently defeating the overflow guard.
+		const { MAX_TIMER_MS } = require('../../components/shutdownDrain.ts');
+		delay = Math.max(
+			delay,
+			Math.min(Math.max(0, selfExitDrainDeadline - Date.now()) + threadTerminationTimeout, MAX_TIMER_MS)
+		);
+	}
+	if (selfExitTimer) clearTimeout(selfExitTimer);
+	selfExitTimer = setTimeout(() => {
+		harperLogger.warn('Thread did not voluntarily terminate', threadId);
+		// Note that if this occurs, you may want to use this to debug what is currently running:
+		// require('why-is-node-running')();
+		realExit(0);
+	}, delay).unref(); // don't block the shutdown
+}
+
+// Worker-side: push both termination backstops out to `deadlineMs` (an absolute epoch timestamp) so the
+// worker can gracefully drain in-flight work before exiting. Extends the local self-exit timer and asks
+// the main thread to extend its external force-terminate timer to match. Called from the shutdown path
+// only when there is real work to drain (see threadServer / shutdownDrain), so an unrelated hang is
+// still force-killed on the normal short timeout.
+function extendShutdownDeadline(deadlineMs) {
+	selfExitDrainDeadline = Math.max(selfExitDrainDeadline, deadlineMs);
+	if (selfExitTimer) armSelfExit(0); // re-arm honoring the (now recorded) drain deadline
+	try {
+		parentPort?.postMessage({ type: EXTEND_SHUTDOWN_DEADLINE, deadlineMs });
+	} catch {}
+}
+
+// Worker-side: drop a drain extension once draining is done, restoring the normal short backstops for
+// the remaining shutdown steps (closeServers / scope disposal) so a later hang is still force-killed
+// promptly. Posting a now-deadline re-arms the main thread's terminate timer back to its normal window.
+function restoreShutdownDeadline() {
+	selfExitDrainDeadline = 0;
+	if (selfExitTimer) armSelfExit(threadTerminationTimeout);
+	try {
+		parentPort?.postMessage({ type: EXTEND_SHUTDOWN_DEADLINE, deadlineMs: Date.now() });
+	} catch {}
+}
+
 const listenersByType = new Map();
 const messagesQueuedByType = new Map();
 
@@ -79,6 +134,12 @@ module.exports = {
 	getTicketKeys,
 	setMainIsWorker,
 	setTerminateTimeout,
+<<<<<<< HEAD
+=======
+	extendShutdownDeadline,
+	restoreShutdownDeadline,
+	registerWorkerDataProvider,
+>>>>>>> 9018760e4 (Merge pull request #1621 from HarperFast/kris/blob-send-drain-core)
 	restartNumber: workerData?.restartNumber || 1,
 };
 
@@ -144,6 +205,9 @@ if (!parentPort) {
 	onMessageByType(RESOURCE_REPORT, (message, worker) => {
 		if (worker) recordResourceReport(worker, message);
 	});
+	onMessageByType(EXTEND_SHUTDOWN_DEADLINE, (message, worker) => {
+		worker?.extendTerminateDeadline?.(message.deadlineMs);
+	});
 }
 // postMessage type listeners that are registered in other ways or can be registered later
 listenersByType.set(hdbTerms.ITC_EVENT_TYPES.CHILD_STARTED, null);
@@ -412,19 +476,37 @@ async function restartWorkers(
 			if (overlapping && startReplacementThreads && !canPreStartReplacement) worker.startCopy();
 			let whenDone = new Promise((resolve) => {
 				// in case the exit inside the thread doesn't timeout, force it from the outside
-				let timeout = setTimeout(() => {
-					harperLogger.warn('Thread did not voluntarily terminate, terminating from the outside', worker.threadId);
-					if (isBun) {
-						// worker.terminate() triggers a NAPI segfault in Bun; ask the worker to self-exit instead
-						try {
-							worker.postMessage({ type: FORCE_EXIT });
-						} catch {}
-					} else {
-						worker.terminate();
-					}
-				}, threadTerminationTimeout * 2).unref();
+				const armTerminate = (delay) =>
+					setTimeout(() => {
+						harperLogger.warn('Thread did not voluntarily terminate, terminating from the outside', worker.threadId);
+						if (isBun) {
+							// worker.terminate() triggers a NAPI segfault in Bun; ask the worker to self-exit instead
+							try {
+								worker.postMessage({ type: FORCE_EXIT });
+							} catch {}
+						} else {
+							worker.terminate();
+						}
+					}, delay).unref();
+				let timeout = armTerminate(threadTerminationTimeout * 2);
+				// The worker can push this backstop out while it gracefully drains in-flight work (e.g. a
+				// replication blob send) before exiting. It only asks when it actually has such work, so a
+				// worker hung for an unrelated reason is still force-killed on the normal short timeout.
+				// This timer is armed synchronously above, in the same tick as the SHUTDOWN post, so the
+				// worker's EXTEND request can only arrive after it exists.
+				worker.extendTerminateDeadline = (deadlineMs) => {
+					clearTimeout(timeout);
+					// Clamp the worker-requested deadline to the configured ceiling (and to a finite value) so a
+					// buggy/rogue message can't defer the force-kill unboundedly; a shrink (drain-done reset)
+					// passes through untouched. See boundedTerminateDelay for the arithmetic + its unit tests.
+					const { boundedTerminateDelay, getShutdownDrainCeilingMs } = require('../../components/shutdownDrain.ts');
+					timeout = armTerminate(
+						boundedTerminateDelay(deadlineMs, Date.now(), threadTerminationTimeout * 2, getShutdownDrainCeilingMs())
+					);
+				};
 				worker.on('exit', () => {
 					clearTimeout(timeout);
+					worker.extendTerminateDeadline = undefined;
 					const index = waitingToFinish.indexOf(whenDone);
 					if (index > -1) waitingToFinish.splice(index, 1);
 					// non-overlapping types have no advance replacement, so start it once the old one is gone
@@ -794,12 +876,7 @@ if (isMainThread) {
 	onMessageByType(hdbTerms.ITC_EVENT_TYPES.SHUTDOWN, async (message) => {
 		module.exports.restartNumber = message.restartNumber;
 		parentPort.unref(); // remove this handle
-		setTimeout(() => {
-			harperLogger.warn('Thread did not voluntarily terminate', threadId);
-			// Note that if this occurs, you may want to use this to debug what is currently running:
-			// require('why-is-node-running')();
-			realExit(0);
-		}, threadTerminationTimeout).unref(); // don't block the shutdown
+		armSelfExit(threadTerminationTimeout);
 	});
 	// In Bun, worker.terminate() triggers a NAPI segfault; the main thread sends FORCE_EXIT
 	// instead, and the worker self-exits cleanly to avoid the crash.

server/threads/threadServer.js

@@ -14,7 +14,17 @@ const env = require('../../utility/environment/environmentManager.ts');
 const terms = require('../../utility/hdbTerms.ts');
 const { server } = require('../Server.ts');
 let { createServer: createSecureSocketServer } = require('node:tls');
-const { restartNumber, getWorkerIndex } = require('./manageThreads.js');
+const {
+	restartNumber,
+	getWorkerIndex,
+	extendShutdownDeadline,
+	restoreShutdownDeadline,
+} = require('./manageThreads.js');
+const {
+	runShutdownDrains,
+	shutdownDrainsHaveWork,
+	getShutdownDrainCeilingMs,
+} = require('../../components/shutdownDrain.ts');
 const { realExit } = require('./workerProcessGuard.ts');
 const { isBun } = require('../serverHelpers/Request.ts');
 const { createTLSSelector } = require('../../security/keys.ts');
@@ -173,11 +183,26 @@ function startServers() {
 						harperLogger.trace('received shutdown request', threadId);
 						// shutdown (for these threads) means stop listening for incoming requests (finish what we are working) and
 						// close connections as possible, then let the event loop complete.
+						// First, gracefully drain any in-flight work registered by components — notably a
+						// replication blob *send* streaming to a peer, which is cheaper to finish than to interrupt
+						// (interrupting leaves the peer's copy diverged until it re-requests). The drain waits only
+						// on work still making progress, bounded by an absolute deadline. When there is real work to
+						// drain we push the termination backstops out to that deadline first so the drain isn't cut
+						// short, then restore the normal short backstop once draining is done — so any later hang
+						// (closeServers / scope disposal) is still force-killed on the normal timeout, and a worker
+						// with no such work is never affected.
+						const drainDeadline = Date.now() + getShutdownDrainCeilingMs();
+						const extendedForDrain = shutdownDrainsHaveWork();
+						if (extendedForDrain) extendShutdownDeadline(drainDeadline);
 						// Wait for application scopes to finish closing before exiting — some dispose a native
 						// runtime asynchronously (e.g. @harperfast/vite's rolldown dev server), and exiting the
 						// worker while that runtime is still live crashes the process. The manageThreads backstop
 						// timers still bound this if a scope's disposal hangs.
-						closeServers()
+						runShutdownDrains(drainDeadline)
+							.then(() => {
+								if (extendedForDrain) restoreShutdownDeadline();
+							})
+							.then(() => closeServers())
 							.then(() => whenScopesClosed())
 							.then(() => {
 								realExit(0);