fix(nats): bound the connection drain on scope close

.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.

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<NATSConnection>("@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<NATSCore.NatsConnection>
+  connect: () => Promise<NATSCore.NatsConnection>,
+  options: NATSConnectionOptions = {}
 ): Effect.Effect<NATSConnection, NATSError.NATSConnectionError, Scope.Scope> =>
   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<NATSConnection, NATSError.NATSConnectionError> =>
+  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<NATSConnection, NATSError.NATSConnectionError> =>
+  Layer.scoped(NATSConnection, make(() => TransportNode.connect(options), connectionOptions))

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<void>
+  readonly close: () => Promise<void>
+}
+
+/**
+ * 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 = <A>(promise: Promise<A>, timeout: Duration.Duration, onTimeout: A): Effect.Effect<A> =>
+  Effect.promise(() => {
+    let timer: ReturnType<typeof setTimeout> | undefined
+
+    return Promise.race([
+      promise,
+      new Promise<A>((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<void> =>
+  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)
+  })

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<void>) => {
+  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<void>(() => {})
+
+const teardown = (drain: () => Promise<void>) =>
+  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))
+    }))
+})