From 1a9105596f87911e3b4a8218f704ada3dbd844d4 Mon Sep 17 00:00:00 2001 From: Adam Hotait Date: Fri, 7 Aug 2026 23:36:09 +0200 Subject: [PATCH] fix(nats): bound the connection drain on scope close The scope finalizer awaited nc.drain() with no deadline. drain() waits for every buffered and in-flight message to be handled before it flushes and closes, so a consumer that keeps receiving messages holds the finalizer open indefinitely and strands process shutdown until the supervisor resorts to SIGKILL. Give the drain a budget, configurable via drainTimeout and defaulting to 5 seconds, and close the connection outright once it expires. A rejected drain now closes too instead of leaving the connection open. The budget is raced at the promise level rather than with Effect.timeout: finalizers run in an uninterruptible region, where the loser of an Effect-level race cannot be interrupted and the timeout never fires. --- .changeset/bound-nats-connection-drain.md | 15 +++++ packages/nats/src/NATSConnection.ts | 43 ++++++++++---- .../nats/src/internal/connectionTeardown.ts | 52 +++++++++++++++++ packages/nats/test/connectionTeardown.test.ts | 58 +++++++++++++++++++ 4 files changed, 158 insertions(+), 10 deletions(-) create mode 100644 .changeset/bound-nats-connection-drain.md create mode 100644 packages/nats/src/internal/connectionTeardown.ts create mode 100644 packages/nats/test/connectionTeardown.test.ts diff --git a/.changeset/bound-nats-connection-drain.md b/.changeset/bound-nats-connection-drain.md new file mode 100644 index 0000000..629b687 --- /dev/null +++ b/.changeset/bound-nats-connection-drain.md @@ -0,0 +1,15 @@ +--- +"@effect-messaging/nats": minor +--- + +Bound the NATS connection teardown and add `drainTimeout` to `layerNode` / `layerWebSocket`. + +The scope finalizer awaited `nc.drain()` with no deadline. `drain()` waits for every buffered and +in-flight message to be handled before it flushes and closes, so a consumer that keeps receiving +messages holds the finalizer open indefinitely and strands process shutdown until the supervisor +resorts to SIGKILL. The finalizer now gives the drain a budget (5 seconds by default) and closes the +connection outright once it expires, and a rejected drain no longer leaves the connection open. + +Note that the budget is raced at the promise level rather than with `Effect.timeout`: finalizers run +in an uninterruptible region, where the loser of an Effect-level race cannot be interrupted and the +timeout never fires. diff --git a/packages/nats/src/NATSConnection.ts b/packages/nats/src/NATSConnection.ts index 89e0630..d590e96 100644 --- a/packages/nats/src/NATSConnection.ts +++ b/packages/nats/src/NATSConnection.ts @@ -4,11 +4,13 @@ import * as NATSCore from "@nats-io/nats-core" import * as TransportNode from "@nats-io/transport-node" import * as Context from "effect/Context" +import * as Duration from "effect/Duration" import * as Effect from "effect/Effect" import * as Layer from "effect/Layer" import * as Option from "effect/Option" import type * as Scope from "effect/Scope" import * as Stream from "effect/Stream" +import * as connectionTeardown from "./internal/connectionTeardown.js" import * as utils from "./internal/utils.js" import * as NATSError from "./NATSError.js" import * as NATSMessage from "./NATSMessage.js" @@ -75,12 +77,29 @@ export interface NATSConnection { */ export const NATSConnection = Context.GenericTag("@effect-messaging/nats/NATSConnection") +/** + * @category models + * @since 0.8.0 + */ +export interface NATSConnectionOptions { + /** + * How long the scope finalizer waits for the connection to drain before closing it outright. + * Defaults to 5 seconds. + * + * @since 0.8.0 + */ + drainTimeout?: Duration.DurationInput +} + +const DEFAULT_DRAIN_TIMEOUT: Duration.DurationInput = "5 seconds" + const wrapAsync = utils.wrapAsync(NATSError.NATSConnectionError) const wrap = utils.wrap(NATSError.NATSConnectionError) /** @internal */ const make = ( - connect: () => Promise + connect: () => Promise, + options: NATSConnectionOptions = {} ): Effect.Effect => Effect.gen(function*() { const nc = yield* wrapAsync(connect, "Failed to create NATS connection") @@ -127,7 +146,9 @@ const make = ( nc } - yield* Effect.addFinalizer(() => Effect.promise(() => nc.drain())) + yield* Effect.addFinalizer(() => + connectionTeardown.closeConnection(nc, Duration.decode(options.drainTimeout ?? DEFAULT_DRAIN_TIMEOUT)) + ) return connection }) @@ -136,16 +157,18 @@ const make = ( * @since 0.1.0 * @category Layers */ -export const layerWebSocket = (options: NATSCore.ConnectionOptions): Layer.Layer< - NATSConnection, - NATSError.NATSConnectionError -> => Layer.scoped(NATSConnection, make(() => NATSCore.wsconnect(options))) +export const layerWebSocket = ( + options: NATSCore.ConnectionOptions, + connectionOptions: NATSConnectionOptions = {} +): Layer.Layer => + Layer.scoped(NATSConnection, make(() => NATSCore.wsconnect(options), connectionOptions)) /** * @since 0.1.0 * @category Layers */ -export const layerNode = (options: TransportNode.NodeConnectionOptions): Layer.Layer< - NATSConnection, - NATSError.NATSConnectionError -> => Layer.scoped(NATSConnection, make(() => TransportNode.connect(options))) +export const layerNode = ( + options: TransportNode.NodeConnectionOptions, + connectionOptions: NATSConnectionOptions = {} +): Layer.Layer => + Layer.scoped(NATSConnection, make(() => TransportNode.connect(options), connectionOptions)) diff --git a/packages/nats/src/internal/connectionTeardown.ts b/packages/nats/src/internal/connectionTeardown.ts new file mode 100644 index 0000000..07a64d1 --- /dev/null +++ b/packages/nats/src/internal/connectionTeardown.ts @@ -0,0 +1,52 @@ +import * as Duration from "effect/Duration" +import * as Effect from "effect/Effect" + +/** The part of a NATS connection this module needs, kept narrow so it can be stubbed in tests. */ +export interface Closeable { + readonly drain: () => Promise + readonly close: () => Promise +} + +/** + * Waits for an already-started promise, falling back to `onTimeout` once `timeout` elapses. + * + * Raced at the promise level on purpose. Finalizers run in an uninterruptible region, where the + * loser of an `Effect.timeout` race cannot be interrupted, so an Effect-level timeout never fires. + * The timer is cleared either way so a completed teardown never holds the event loop open. + */ +const settleWithin = (promise: Promise, timeout: Duration.Duration, onTimeout: A): Effect.Effect => + Effect.promise(() => { + let timer: ReturnType | undefined + + return Promise.race([ + promise, + new Promise((resolve) => { + timer = setTimeout(() => resolve(onTimeout), Duration.toMillis(timeout)) + }) + ]).finally(() => clearTimeout(timer)) + }) + +/** + * Drains the connection, giving up and closing it outright once `drainTimeout` elapses. + * + * `drain()` unsubscribes, waits for every buffered and in-flight message to be handled, then + * flushes and closes. A consumer that keeps receiving messages can hold that open indefinitely, + * stranding process shutdown until the supervisor resorts to SIGKILL. + */ +export const closeConnection = (connection: Closeable, drainTimeout: Duration.Duration): Effect.Effect => + Effect.gen(function*() { + // Settled rather than caught, so a rejected drain counts as "not drained" and still gets closed, + // and so nothing is left unhandled once the budget expires and no one awaits this any more. + const drained = connection.drain().then(() => true, () => false) + + if (yield* settleWithin(drained, drainTimeout, false)) { + return + } + + yield* Effect.logWarning( + `NATS connection did not drain within ${Duration.format(drainTimeout)}, closing it instead` + ) + + const closed = connection.close().then(() => undefined, () => undefined) + yield* settleWithin(closed, drainTimeout, undefined) + }) diff --git a/packages/nats/test/connectionTeardown.test.ts b/packages/nats/test/connectionTeardown.test.ts new file mode 100644 index 0000000..441bfda --- /dev/null +++ b/packages/nats/test/connectionTeardown.test.ts @@ -0,0 +1,58 @@ +import { describe, expect, it } from "@effect/vitest" +import { Duration, Effect } from "effect" +import { closeConnection } from "../src/internal/connectionTeardown.js" + +const DRAIN_TIMEOUT = Duration.millis(50) +// Wide enough that a slow machine cannot fail it, tight enough that an unbounded wait cannot pass. +const TEARDOWN_BUDGET_MILLIS = 2_000 + +const stub = (drain: () => Promise) => { + const calls: Array<"drain" | "close"> = [] + return { + calls, + connection: { + drain: () => { + calls.push("drain") + return drain() + }, + close: () => { + calls.push("close") + return Promise.resolve() + } + } + } +} + +const never = () => new Promise(() => {}) + +const teardown = (drain: () => Promise) => + Effect.gen(function*() { + const { calls, connection } = stub(drain) + const [elapsed] = yield* Effect.timed(closeConnection(connection, DRAIN_TIMEOUT)) + return { calls, elapsedMillis: Duration.toMillis(elapsed) } + }) + +describe("closeConnection", () => { + it.live("closes the connection when the drain never settles", () => + Effect.gen(function*() { + const { calls, elapsedMillis } = yield* teardown(never) + + expect(calls).toEqual(["drain", "close"]) + expect(elapsedMillis).toBeLessThan(TEARDOWN_BUDGET_MILLIS) + })) + + it.live("closes the connection when the drain rejects", () => + Effect.gen(function*() { + const { calls } = yield* teardown(() => Promise.reject(new Error("already closing"))) + + expect(calls).toEqual(["drain", "close"]) + })) + + it.live("leaves a healthy drain to finish on its own", () => + Effect.gen(function*() { + const { calls, elapsedMillis } = yield* teardown(() => Promise.resolve()) + + expect(calls).toEqual(["drain"]) + expect(elapsedMillis).toBeLessThan(Duration.toMillis(DRAIN_TIMEOUT)) + })) +})