diff --git a/crates/taskito-node/src/config.rs b/crates/taskito-node/src/config.rs index 053bae5d..f05eb66f 100644 --- a/crates/taskito-node/src/config.rs +++ b/crates/taskito-node/src/config.rs @@ -140,6 +140,9 @@ pub struct WorkerOptions { pub batch_size: Option, pub task_configs: Option>, pub queue_configs: Option>, + /// Names of the injectable resources this worker serves, advertised on + /// registration so inspection surfaces them alongside heartbeat health. + pub resources: Option>, /// Opt-in decentralized mesh overlay (requires the `mesh` build feature). pub mesh: Option, } diff --git a/crates/taskito-node/src/convert/stats.rs b/crates/taskito-node/src/convert/stats.rs index 9fd7bbdc..b1959798 100644 --- a/crates/taskito-node/src/convert/stats.rs +++ b/crates/taskito-node/src/convert/stats.rs @@ -109,6 +109,11 @@ pub struct JsWorkerRow { pub pool_type: Option, pub threads: i32, pub tags: Option, + /// JSON array of resource names the worker advertised at registration. + pub resources: Option, + /// JSON object of per-resource health (`"healthy"`/`"unhealthy"`), written + /// by the worker's heartbeat. + pub resource_health: Option, } pub fn worker_to_js(worker: WorkerRow) -> JsWorkerRow { @@ -123,6 +128,8 @@ pub fn worker_to_js(worker: WorkerRow) -> JsWorkerRow { pool_type: worker.pool_type, threads: worker.threads, tags: worker.tags, + resources: worker.resources, + resource_health: worker.resource_health, } } diff --git a/crates/taskito-node/src/queue/inspect.rs b/crates/taskito-node/src/queue/inspect.rs index ed76ecd8..c7d91eaf 100644 --- a/crates/taskito-node/src/queue/inspect.rs +++ b/crates/taskito-node/src/queue/inspect.rs @@ -126,4 +126,25 @@ impl JsQueue { .await .map_err(join_to_napi_err)? } + + /// Record a heartbeat for a running worker, carrying its current resource + /// health as a JSON object (`null` clears it). Called from the JS shell + /// every 5s. Returns the ids of dead workers reaped as a side effect. + #[napi] + pub async fn worker_heartbeat( + &self, + worker_id: String, + resource_health: Option, + ) -> Result> { + let storage = self.storage.clone(); + spawn_blocking(move || { + storage + .heartbeat(&worker_id, resource_health.as_deref()) + .map_err(to_napi_err)?; + // Reaping is opportunistic — a failure must not fail the heartbeat. + Ok(storage.reap_dead_workers().unwrap_or_default()) + }) + .await + .map_err(join_to_napi_err)? + } } diff --git a/crates/taskito-node/src/worker.rs b/crates/taskito-node/src/worker.rs index 0c24fa65..ec672c15 100644 --- a/crates/taskito-node/src/worker.rs +++ b/crates/taskito-node/src/worker.rs @@ -1,8 +1,9 @@ //! Worker wiring — scheduler, dispatcher, result-drain, and worker lifecycle -//! (registration + heartbeat) so the worker shows up in the dashboard. +//! (registration + unregistration) so the worker shows up in the dashboard. +//! Heartbeats are driven from JS via [`crate::queue::JsQueue::worker_heartbeat`] +//! so they can carry resource health without a second writer flapping the column. use std::sync::Arc; -use std::time::Duration; use napi::bindgen_prelude::{spawn, spawn_blocking, Result}; use napi::threadsafe_function::{ErrorStrategy, ThreadsafeFunction, ThreadsafeFunctionCallMode}; @@ -19,29 +20,36 @@ use crate::error::invalid_arg; const DEFAULT_QUEUE: &str = "default"; const DEFAULT_CHANNEL_CAPACITY: usize = 128; -const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); /// Handle to a running worker. Hold it for the worker's lifetime; call /// [`JsWorker::stop`] to shut it down. #[napi] pub struct JsWorker { + worker_id: String, shutdown: Arc, - heartbeat_stop: Arc, + lifecycle_stop: Arc, #[cfg(feature = "mesh")] mesh_shutdown: Arc, } #[napi] impl JsWorker { - /// Stop the worker: the scheduler stops dispatching, the heartbeat loop ends - /// and unregisters, mesh gossip/steal tasks shut down, and the background + /// The id this worker registered (and claims executions) under. The JS + /// shell heartbeats with it via [`crate::queue::JsQueue::worker_heartbeat`]. + #[napi(getter)] + pub fn id(&self) -> String { + self.worker_id.clone() + } + + /// Stop the worker: the scheduler stops dispatching, the lifecycle task + /// unregisters, mesh gossip/steal tasks shut down, and the background /// tasks exit once in-flight results drain. #[napi] pub fn stop(&self) { // `notify_one` stores a permit if no waiter is parked yet, so the signal // is never lost between loop iterations. self.shutdown.notify_one(); - self.heartbeat_stop.notify_one(); + self.lifecycle_stop.notify_one(); #[cfg(feature = "mesh")] self.mesh_shutdown.notify_one(); } @@ -112,14 +120,16 @@ pub fn start_worker( let (job_tx, job_rx) = tokio::sync::mpsc::channel(capacity); let (result_tx, result_rx) = crossbeam_channel::bounded(capacity); - // Worker lifecycle: register, heartbeat every 5s, unregister on stop. - let heartbeat_stop = Arc::new(Notify::new()); + // Worker lifecycle: register (advertising resource names), unregister on + // stop. Heartbeats come from JS so they carry current resource health. + let lifecycle_stop = Arc::new(Notify::new()); spawn_worker_lifecycle( lifecycle_storage, worker_id.clone(), queues_csv, + resources_json(options.resources.as_deref()), capacity, - heartbeat_stop.clone(), + lifecycle_stop.clone(), ); // Scheduler loop: poll storage, dispatch ready jobs onto `job_tx`. With the @@ -200,18 +210,29 @@ pub fn start_worker( }); Ok(JsWorker { + worker_id, shutdown, - heartbeat_stop, + lifecycle_stop, #[cfg(feature = "mesh")] mesh_shutdown, }) } -/// Register this worker and heartbeat until `stop` is signalled, then unregister. +/// Serialize advertised resource names to the cross-SDK wire shape: a sorted +/// JSON array, or `None` when the worker registered no resources. +fn resources_json(resources: Option<&[String]>) -> Option { + let names = resources.filter(|names| !names.is_empty())?; + let mut sorted = names.to_vec(); + sorted.sort(); + serde_json::to_string(&sorted).ok() +} + +/// Register this worker, then unregister once `stop` is signalled. fn spawn_worker_lifecycle( storage: StorageBackend, worker_id: String, queues_csv: String, + resources: Option, capacity: usize, stop: Arc, ) { @@ -228,7 +249,7 @@ fn spawn_worker_lifecycle( ®_id, &queues_csv, None, - None, + resources.as_deref(), None, capacity as i32, Some(&hostname), @@ -243,20 +264,7 @@ fn spawn_worker_lifecycle( Ok(Ok(_)) => {} } - loop { - tokio::select! { - _ = stop.notified() => break, - _ = tokio::time::sleep(HEARTBEAT_INTERVAL) => { - let hb_storage = storage.clone(); - let hb_id = worker_id.clone(); - match spawn_blocking(move || hb_storage.heartbeat(&hb_id, None)).await { - Ok(Err(err)) => log::warn!("[taskito-node] worker heartbeat failed: {err}"), - Err(err) => log::warn!("[taskito-node] worker heartbeat task failed: {err}"), - Ok(Ok(_)) => {} - } - } - } - } + stop.notified().await; match spawn_blocking(move || storage.unregister_worker(&worker_id)).await { Ok(Err(err)) => log::warn!("[taskito-node] worker unregister failed: {err}"), diff --git a/sdks/node/src/queue.ts b/sdks/node/src/queue.ts index 24e11459..5601b730 100644 --- a/sdks/node/src/queue.ts +++ b/sdks/node/src/queue.ts @@ -262,6 +262,13 @@ export class Queue { scope?: ResourceScope; dispose?: (value: T) => void | Promise; pool?: PoolOptions; + /** Returns truthy while healthy; failures trigger recreation. Worker scope only. */ + healthCheck?: (value: T) => boolean | Promise; + /** Milliseconds between health checks. 0 or absent disables checking. */ + healthCheckIntervalMs?: number; + /** Failed checks tolerated (while recreation also fails) before the + * resource is marked permanently unhealthy. Default 3. */ + maxRecreationAttempts?: number; }, ): this { const scope = options?.scope ?? "worker"; @@ -270,11 +277,20 @@ export class Queue { `Resource "${name}": pool options require scope "pooled" (got "${scope}")`, ); } + if (options?.healthCheck && scope !== "worker") { + throw new ResourceError( + `Resource "${name}": health checks require scope "worker" (got "${scope}") — ` + + "task and pooled instances are already rebuilt or recycled per job", + ); + } this.resources.register(name, { factory, scope, dispose: options?.dispose, pool: options?.pool, + healthCheck: options?.healthCheck, + healthCheckIntervalMs: options?.healthCheckIntervalMs, + maxRecreationAttempts: options?.maxRecreationAttempts, }); return this; } diff --git a/sdks/node/src/resources/health.ts b/sdks/node/src/resources/health.ts new file mode 100644 index 00000000..f61e7a8b --- /dev/null +++ b/sdks/node/src/resources/health.ts @@ -0,0 +1,133 @@ +import { createLogger } from "../utils"; +import type { ResourceRuntime } from "./runtime"; +import type { ResourceDefinition } from "./types"; + +const log = createLogger("resources"); + +/** How often the daemon wakes to see whether any resource's check is due. */ +const TICK_MS = 500; + +/** Fallback for {@link ResourceDefinition.maxRecreationAttempts}. */ +const DEFAULT_MAX_RECREATION_ATTEMPTS = 3; + +/** + * Periodically checks worker-resource health and recreates failing instances. + * + * One unref'd timer wakes every {@link TICK_MS}; each resource with a positive + * `healthCheckIntervalMs` is checked independently on its own schedule. A + * failed (or throwing) check triggers recreation; when recreation also fails + * and the failure budget is spent, the resource is marked permanently + * unhealthy and never rechecked. + */ +export class HealthChecker { + private timer: ReturnType | undefined; + /** The tick currently running, if any — doubles as the overlap guard. */ + private inFlight: Promise | undefined; + /** Set by {@link HealthChecker.stop} — a stopped checker never restarts. */ + private stopped = false; + private readonly nextDue = new Map(); + private readonly failCount = new Map(); + + constructor(private readonly runtime: ResourceRuntime) {} + + /** Start the daemon. No-op unless some resource asked for health checks. */ + start(): void { + if (this.timer || this.stopped) { + return; + } + const checked = this.runtime.healthChecked(); + if (checked.size === 0) { + return; + } + const now = Date.now(); + for (const [name, def] of checked) { + this.nextDue.set(name, now + (def.healthCheckIntervalMs ?? 0)); + this.failCount.set(name, 0); + } + // The tick never rejects (see below), so the timer can't leak an + // unhandled rejection; unref keeps it from pinning the process open. + this.timer = setInterval(() => this.onTick(checked), TICK_MS); + this.timer.unref(); + } + + /** + * Stop the daemon and wait for any in-flight tick to drain, so no + * check/recreate can mutate runtime state once this resolves — safe to run + * worker teardown right after. Terminal and idempotent: a stopped checker + * never restarts. + */ + async stop(): Promise { + this.stopped = true; + if (this.timer) { + clearInterval(this.timer); + this.timer = undefined; + } + await this.inFlight; + } + + /** One timer wake-up: skip when the previous tick is still in flight. */ + private onTick(checked: ReadonlyMap): void { + if (this.inFlight || this.stopped) { + return; + } + this.inFlight = this.tick(checked).finally(() => { + this.inFlight = undefined; + }); + } + + /** Run every due check. Catches everything — timers must survive. */ + private async tick(checked: ReadonlyMap): Promise { + try { + for (const [name, def] of checked) { + if (this.stopped) { + return; // shutdown began mid-tick — leave the runtime alone + } + if (Date.now() < (this.nextDue.get(name) ?? 0) || this.runtime.isUnhealthy(name)) { + continue; + } + await this.checkAndRepair(name, def); + this.nextDue.set(name, Date.now() + (def.healthCheckIntervalMs ?? 0)); + } + } catch (error) { + log.warn(() => "health-check tick failed", error); + } + } + + /** Check one resource; on failure recreate it, exhausting into unhealthy. */ + private async checkAndRepair(name: string, def: ResourceDefinition): Promise { + if (await this.checkOne(name, def)) { + this.failCount.set(name, 0); + return; + } + const failures = (this.failCount.get(name) ?? 0) + 1; + this.failCount.set(name, failures); + const maxAttempts = def.maxRecreationAttempts ?? DEFAULT_MAX_RECREATION_ATTEMPTS; + log.warn(`resource "${name}" health check failed (${failures}/${maxAttempts})`); + if (this.stopped) { + return; // shutdown began while the check ran — don't recreate into teardown + } + if (await this.runtime.recreate(name)) { + this.failCount.set(name, 0); // a fresh instance starts with a clean slate + } else if (failures >= maxAttempts) { + this.runtime.markUnhealthy(name); + } + } + + /** Run a single health check. A throwing check counts as unhealthy. */ + private async checkOne(name: string, def: ResourceDefinition): Promise { + const { healthCheck } = def; + if (!healthCheck) { + return true; + } + const built = this.runtime.builtWorkerInstance(name); + if (!built) { + return false; // nothing built yet — recreation will build it + } + try { + return Boolean(await healthCheck(await built)); + } catch (error) { + log.warn(() => `health check for resource "${name}" threw`, error); + return false; + } + } +} diff --git a/sdks/node/src/resources/index.ts b/sdks/node/src/resources/index.ts index 6f0a0a14..7071bdc3 100644 --- a/sdks/node/src/resources/index.ts +++ b/sdks/node/src/resources/index.ts @@ -1,7 +1,8 @@ export { runWithResolver, useResource } from "./context"; +export { HealthChecker } from "./health"; export { type MockResource, mockResource } from "./mock"; export { ResourcePool } from "./pool"; -export { ResourceRuntime, type TaskScope } from "./runtime"; +export { type HealthState, ResourceRuntime, type TaskScope } from "./runtime"; export type { PoolOptions, PoolStats, diff --git a/sdks/node/src/resources/runtime.ts b/sdks/node/src/resources/runtime.ts index 051a914e..f751bb18 100644 --- a/sdks/node/src/resources/runtime.ts +++ b/sdks/node/src/resources/runtime.ts @@ -1,4 +1,6 @@ -import { ResourceNotFoundError, ResourceScopeError } from "../errors"; +import { ResourceNotFoundError, ResourceScopeError, ResourceUnavailableError } from "../errors"; +import { createLogger } from "../utils"; +import { HealthChecker } from "./health"; import { ResourcePool } from "./pool"; import type { ResourceContext, @@ -8,12 +10,19 @@ import type { ResourceScope, } from "./types"; +const log = createLogger("resources"); + /** A disposal thunk plus the resource name, for error context. */ interface Teardown { name: string; run: () => void | Promise; + /** The build this thunk disposes, so recreation can retire exactly it. */ + built?: Promise; } +/** Per-resource health, as advertised in the worker heartbeat. */ +export type HealthState = "healthy" | "unhealthy"; + /** Per-invocation task scope: caches task-scoped resources and disposes them. */ export interface TaskScope { /** Resolve a resource for this invocation (task cache first, then worker). */ @@ -42,10 +51,16 @@ export class ResourceRuntime { private readonly pools = new Map(); /** Active worker leases sharing this runtime; teardown disposes only at zero. */ private workerLeases = 0; + /** Resources marked permanently unhealthy — resolving them rejects for good. */ + private readonly unhealthy = new Set(); + /** One checker per runtime, shared by every lease — never one per worker. */ + private healthChecker: HealthChecker | undefined; /** Register (or replace) a resource definition. */ register(name: string, definition: ResourceDefinition): void { this.defs.set(name, definition as ResourceDefinition); + // A replacement definition starts with a clean bill of health. + this.unhealthy.delete(name); // Drop any built worker instance so "replace" takes effect; the old // instance is still disposed by its queued teardown. this.workerCache.delete(name); @@ -83,6 +98,11 @@ export class ResourceRuntime { /** Build a worker-scoped resource once, memoizing the promise (dedups concurrent init). */ private resolveWorker(name: string): Promise { + if (this.unhealthy.has(name)) { + return Promise.reject( + new ResourceUnavailableError(`Resource "${name}" is permanently unhealthy`), + ); + } const cached = this.workerCache.get(name); if (cached) { return cached; @@ -222,6 +242,12 @@ export class ResourceRuntime { /** Register a worker that shares this runtime's worker-scoped resources. */ acquireWorker(): void { this.workerLeases += 1; + // Start the shared checker on the first lease only — concurrent workers on + // one runtime must not race recreate() on the same resource. + if (this.workerLeases === 1) { + this.healthChecker = new HealthChecker(this); + this.healthChecker.start(); + } // Best-effort prewarm of pooled resources that asked for warm instances. // `prewarm` never rejects — build failures are logged inside the pool. for (const [name, def] of this.defs) { @@ -244,8 +270,16 @@ export class ResourceRuntime { if (this.workerLeases > 0) { return; } + // Drain the checker first: once stop() resolves, no in-flight tick can + // recreate a resource while the caches below are torn down. + const checker = this.healthChecker; + this.healthChecker = undefined; + await checker?.stop(); const pending = this.workerTeardown.splice(0); this.workerCache.clear(); + // "Permanently unhealthy" is scoped to a worker run — the next worker + // starts from scratch and may succeed where this one gave up. + this.unhealthy.clear(); // Shut pools down first — pooled instances may depend on worker resources. // `shutdown` never throws; per-instance dispose failures are logged. const pools = [...this.pools.values()]; @@ -256,6 +290,102 @@ export class ResourceRuntime { await runTeardown(pending); } + // ── Health ──────────────────────────────────────────────────────────────── + + /** Registered resource names, sorted — advertised on worker registration. */ + get names(): string[] { + return [...this.defs.keys()].sort(); + } + + /** Definitions that asked for periodic health checks (interval > 0). */ + healthChecked(): ReadonlyMap { + const out = new Map(); + for (const [name, def] of this.defs) { + if (def.healthCheck && (def.healthCheckIntervalMs ?? 0) > 0) { + out.set(name, def); + } + } + return out; + } + + /** Whether `name` was marked permanently unhealthy. */ + isUnhealthy(name: string): boolean { + return this.unhealthy.has(name); + } + + /** The built (possibly pending) worker instance, or undefined before first build. */ + builtWorkerInstance(name: string): Promise | undefined { + return this.workerCache.get(name); + } + + /** + * Mark a resource permanently unhealthy: every later resolve rejects with + * {@link ResourceUnavailableError}. Terminal — there is no auto-recovery. + */ + markUnhealthy(name: string): void { + this.unhealthy.add(name); + log.error(`resource "${name}" marked permanently unhealthy`); + } + + /** + * Dispose the current worker instance (best effort) and build a fresh one. + * Returns false when the factory fails — the caller decides whether that + * exhausts the recreation budget. + */ + async recreate(name: string): Promise { + const def = this.defs.get(name); + if (!def) { + return false; + } + await this.disposeReplacedInstance(name, def); + try { + await this.resolveWorker(name); + return true; + } catch (error) { + log.warn(() => `recreating resource "${name}" failed`, error); + return false; + } + } + + /** + * Per-resource health for the worker heartbeat, or undefined when nothing is + * registered (so the heartbeat clears the column instead of writing `{}`). + */ + healthSnapshot(): Record | undefined { + if (this.defs.size === 0) { + return undefined; + } + const out: Record = {}; + for (const name of this.defs.keys()) { + out[name] = this.unhealthy.has(name) ? "unhealthy" : "healthy"; + } + return out; + } + + /** + * Retire the currently-built worker instance ahead of recreation: drop it + * from the cache, cancel its queued teardown (recreation disposes it now), + * and run its `dispose` best-effort. + */ + private async disposeReplacedInstance(name: string, def: ResourceDefinition): Promise { + const old = this.workerCache.get(name); + if (!old) { + return; + } + this.workerCache.delete(name); + const staleIndex = this.workerTeardown.findIndex((entry) => entry.built === old); + if (staleIndex !== -1) { + this.workerTeardown.splice(staleIndex, 1); + } + try { + const value = await old; + await def.dispose?.(value); + this.counter(name).disposed += 1; + } catch (error) { + log.debug(() => `disposing resource "${name}" before recreation failed`, error); + } + } + /** * Queue a teardown thunk for a resource: run its `dispose` (if any) and record * the disposal. A build that failed is skipped — there's nothing to tear down. @@ -270,6 +400,7 @@ export class ResourceRuntime { const counter = this.counter(name); stack.push({ name, + built, run: async () => { let value: unknown; try { diff --git a/sdks/node/src/resources/types.ts b/sdks/node/src/resources/types.ts index 49d68da6..cce27ee7 100644 --- a/sdks/node/src/resources/types.ts +++ b/sdks/node/src/resources/types.ts @@ -23,6 +23,18 @@ export interface ResourceDefinition { dispose?: (value: T) => void | Promise; /** Pool tuning; only meaningful for `"pooled"` scope. */ pool?: PoolOptions; + /** + * Returns truthy while the value is healthy. Checked on the interval below; + * a failing (or throwing) check triggers recreation. Worker scope only. + */ + healthCheck?: (value: T) => boolean | Promise; + /** Milliseconds between health checks. 0 or absent disables checking. */ + healthCheckIntervalMs?: number; + /** + * Consecutive failed checks tolerated while recreation also fails before the + * resource is marked permanently unhealthy. Default 3. + */ + maxRecreationAttempts?: number; } /** Tuning for the bounded checkout pool behind a `"pooled"`-scope resource. */ diff --git a/sdks/node/src/worker.ts b/sdks/node/src/worker.ts index dacc474c..a55119ec 100644 --- a/sdks/node/src/worker.ts +++ b/sdks/node/src/worker.ts @@ -23,6 +23,9 @@ const log = createLogger("worker"); /** How often a running job polls the storage cancel flag. */ const CANCEL_POLL_INTERVAL_MS = 200; +/** How often the worker heartbeats (with resource health) to storage. */ +const HEARTBEAT_INTERVAL_MS = 5000; + /** Outcome kind -> event name + the middleware hook it triggers. */ const OUTCOMES: Record = { success: { event: "job.completed", hook: "onCompleted" }, @@ -51,6 +54,7 @@ export class Worker { private constructor( private readonly native: NativeWorker, private readonly resources: ResourceRuntime, + private readonly heartbeat: ReturnType, ) {} /** @@ -192,21 +196,40 @@ export class Worker { batchSize: run?.batchSize, taskConfigs: buildTaskConfigs(tasks), queueConfigs: buildQueueConfigs(queueLimits), + resources: resources.isEmpty ? undefined : resources.names, mesh: run?.mesh, }; const native = queue.runWorker(taskCallback, outcomeCallback, nativeOptions); // Lease the shared resource runtime only once the native worker actually // started, so its worker-scoped values survive until the last worker on this // queue stops (see ResourceRuntime). A failed start leaks no lease. + // The lease also starts the runtime's shared health checker (first lease + // only) — recreation of failing resources is per runtime, not per worker. resources.acquireWorker(); - return new Worker(native, resources); + + // Heartbeat with current resource health so inspection (and dead-worker + // reaping) see this worker as alive. Failures are logged, never thrown — + // the next beat retries. First beat goes out immediately. + const sendHeartbeat = (): void => { + const snapshot = resources.healthSnapshot(); + void queue.workerHeartbeat(native.id, snapshot && JSON.stringify(snapshot)).catch((error) => { + log.debug(() => "worker heartbeat failed", error); + }); + }; + sendHeartbeat(); + const heartbeat = setInterval(sendHeartbeat, HEARTBEAT_INTERVAL_MS); + heartbeat.unref(); + + return new Worker(native, resources, heartbeat); } /** Stop the worker; in-flight results drain before background tasks exit. */ stop(): void { + clearInterval(this.heartbeat); this.native.stop(); - // Dispose worker-scoped resources after the native worker quiesces. Best - // effort: lazy resources mean this is a no-op when none were built. + // Dispose worker-scoped resources after the native worker quiesces (the + // teardown drains the runtime's health checker before touching caches). + // Best effort: lazy resources mean this is a no-op when none were built. void this.resources.teardownWorker().catch((error) => { log.debug(() => "worker-scope resource teardown failed", error); }); diff --git a/sdks/node/test/resources/health.test.ts b/sdks/node/test/resources/health.test.ts new file mode 100644 index 00000000..ad9980c7 --- /dev/null +++ b/sdks/node/test/resources/health.test.ts @@ -0,0 +1,204 @@ +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it } from "vitest"; +import { Queue, ResourceUnavailableError, type Worker } from "../../src/index"; +import { HealthChecker } from "../../src/resources/health"; +import { ResourceRuntime } from "../../src/resources/runtime"; + +let worker: Worker | undefined; + +afterEach(() => { + worker?.stop(); + worker = undefined; +}); + +function newQueue(): Queue { + return new Queue({ dbPath: join(mkdtempSync(join(tmpdir(), "taskito-health-")), "q.db") }); +} + +async function waitFor(predicate: () => boolean, timeoutMs = 4000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (predicate()) { + return true; + } + await new Promise((resolve) => setTimeout(resolve, 20)); + } + return false; +} + +describe("HealthChecker", () => { + it("recreates a resource whose health check fails", async () => { + const rt = new ResourceRuntime(); + let builds = 0; + rt.register("svc", { + scope: "worker", + factory: () => ({ id: ++builds }), + healthCheck: () => false, + healthCheckIntervalMs: 100, + // With successful recreates the fail count must keep resetting, so even + // many failed checks past this budget never mark the resource unhealthy. + maxRecreationAttempts: 1, + }); + const first = await rt.createTaskScope().resolver("svc"); + + const checker = new HealthChecker(rt); + checker.start(); + expect(await waitFor(() => builds >= 3, 8000)).toBe(true); + await checker.stop(); + + expect(rt.isUnhealthy("svc")).toBe(false); // recreation succeeded every time + const replaced = await rt.createTaskScope().resolver("svc"); + expect(replaced).not.toBe(first); // fresh instance, new identity + expect(rt.metrics().svc?.created).toBeGreaterThanOrEqual(3); + expect(rt.metrics().svc?.disposed).toBeGreaterThanOrEqual(2); // retired builds + }); + + it("marks a resource unhealthy once recreation attempts are exhausted", async () => { + const rt = new ResourceRuntime(); + let calls = 0; + rt.register("svc", { + scope: "worker", + factory: () => { + calls += 1; + if (calls > 1) { + throw new Error("factory broken"); + } + return "svc_v1"; + }, + healthCheck: () => false, + healthCheckIntervalMs: 50, + maxRecreationAttempts: 2, + }); + expect(await rt.createTaskScope().resolver("svc")).toBe("svc_v1"); + + const checker = new HealthChecker(rt); + checker.start(); + expect(await waitFor(() => rt.isUnhealthy("svc"), 8000)).toBe(true); + await checker.stop(); + + // Terminal: resolving now (and forever) rejects. + await expect(rt.createTaskScope().resolver("svc")).rejects.toThrow(ResourceUnavailableError); + await expect(rt.createTaskScope().resolver("svc")).rejects.toThrow(/permanently unhealthy/); + }); + + it("treats a throwing health check as unhealthy and recreates", async () => { + const rt = new ResourceRuntime(); + let builds = 0; + rt.register("svc", { + scope: "worker", + factory: () => ++builds, + healthCheck: () => { + throw new Error("check exploded"); + }, + healthCheckIntervalMs: 50, + }); + await rt.createTaskScope().resolver("svc"); + + const checker = new HealthChecker(rt); + checker.start(); + expect(await waitFor(() => builds >= 2, 8000)).toBe(true); // recreate path ran + await checker.stop(); + expect(rt.isUnhealthy("svc")).toBe(false); // recreation kept succeeding + }); + + it("start() is a no-op without health-checked resources; stop() is idempotent", async () => { + const rt = new ResourceRuntime(); + rt.register("plain", { scope: "worker", factory: () => 1 }); + rt.register("disabled", { + scope: "worker", + factory: () => 2, + healthCheck: () => true, + healthCheckIntervalMs: 0, // 0 = disabled + }); + const checker = new HealthChecker(rt); + checker.start(); // no timer started — nothing asked for checking + await checker.stop(); + await checker.stop(); // safe without a timer, and repeatedly + checker.start(); // terminal: a stopped checker never restarts + await checker.stop(); + }); + + it("stop() drains an in-flight tick and blocks recreation after it", async () => { + const rt = new ResourceRuntime(); + let builds = 0; + let releaseCheck: ((healthy: boolean) => void) | undefined; + rt.register("svc", { + scope: "worker", + factory: () => ++builds, + healthCheck: () => + new Promise((resolve) => { + releaseCheck = resolve; + }), + healthCheckIntervalMs: 50, + }); + await rt.createTaskScope().resolver("svc"); + + const checker = new HealthChecker(rt); + checker.start(); + expect(await waitFor(() => releaseCheck !== undefined, 8000)).toBe(true); // check in flight + + const stopping = checker.stop(); // must wait for the pending check + let drained = false; + void stopping.then(() => { + drained = true; + }); + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(drained).toBe(false); // still blocked on the in-flight check + + releaseCheck?.(false); // fails — but shutdown already began + await stopping; + expect(builds).toBe(1); // no recreate ran into the teardown window + }); + + it("rejects health checks on task- and pooled-scoped resources", () => { + const queue = newQueue(); + expect(() => + queue.resource("perJob", () => 1, { + scope: "task", + healthCheck: () => true, + healthCheckIntervalMs: 100, + }), + ).toThrow(/health checks require scope "worker"/); + expect(() => + queue.resource("conn", () => 1, { + scope: "pooled", + healthCheck: () => true, + healthCheckIntervalMs: 100, + }), + ).toThrow(/health checks require scope "worker"/); + }); + + it("advertises resource health through the worker heartbeat", async () => { + const queue = newQueue(); + queue.resource("db", () => ({ ok: true }), { + healthCheck: (value) => value.ok, + healthCheckIntervalMs: 100, + }); + queue.task("noop", () => undefined); + + worker = queue.runWorker(); + let row: { resources?: string; resourceHealth?: string } | undefined; + const deadline = Date.now() + 12000; + while (Date.now() < deadline && row === undefined) { + const workers = await queue.listWorkers(); + row = workers.find((w) => w.resourceHealth?.includes('"healthy"')); + await new Promise((resolve) => setTimeout(resolve, 50)); + } + expect(row).toBeDefined(); + expect(JSON.parse(row?.resourceHealth ?? "{}")).toEqual({ db: "healthy" }); + expect(JSON.parse(row?.resources ?? "[]")).toEqual(["db"]); + + worker.stop(); + worker = undefined; + // Unregistration still runs on stop — the row disappears instead of flapping. + let gone = false; + const stopDeadline = Date.now() + 8000; + while (Date.now() < stopDeadline && !gone) { + gone = (await queue.listWorkers()).length === 0; + await new Promise((resolve) => setTimeout(resolve, 50)); + } + expect(gone).toBe(true); + }); +});