Consolidate orphan-hook bounded-wait primitive into async_backend.nim

async_postgres/async_backend.nim

@@ -15,6 +15,18 @@ when hasChronos:
   import chronos
   export chronos
 
+  proc wait*[T](
+      fut: Future[T], timeout: Duration, onOrphan: proc(fut: Future[T]) {.gcsafe.}
+  ): Future[T] {.async.} =
+    ## Wait for a future with a timeout. Raises AsyncTimeoutError on timeout.
+    ##
+    ## ``onOrphan`` is never called under chronos — the inner future is properly
+    ## cancelled on timeout, so no orphan can exist. Accepted for API parity with
+    ## the asyncdispatch backend, where the hook handles futures that keep running
+    ## after ``wait()`` fires.
+    let _ = onOrphan
+    return await chronos.wait(fut, timeout)
+
   proc sleepMsAsync*(ms: int): Future[void] =
     ## Sleep for `ms` milliseconds. Wrapper around chronos Duration-based API.
     sleepAsync(milliseconds(ms))
@@ -151,31 +163,30 @@ elif hasAsyncDispatch:
   proc `>`*(a, b: Moment): bool =
     a.ticks > b.ticks
 
-  proc wait*[T](fut: Future[T], timeout: Duration): Future[T] {.async.} =
+  proc wait*[T](
+      fut: Future[T], timeout: Duration, onOrphan: proc(fut: Future[T]) {.gcsafe.} = nil
+  ): Future[T] {.async.} =
     ## Wait for a future with a timeout. Raises AsyncTimeoutError on timeout.
     ## API-compatible with chronos Future.wait().
     ##
-    ## .. warning::
-    ##   asyncdispatch has no cancellation. When a timeout fires, the inner
-    ##   future keeps running in the background until its I/O completes. For
-    ##   this reason **the underlying connection MUST NOT be reused after an
-    ##   AsyncTimeoutError** — a late write from the suspended future would
-    ##   corrupt the protocol stream of whoever reuses it next. The pg_client
-    ##   layer forces `csClosed` on timeout under asyncdispatch; see
-    ##   `invalidateOnTimeout` in pg_client.nim. chronos is not affected because
-    ##   its futures are actually cancelled.
+    ## ``onOrphan`` cleanup hook
+    ## asyncdispatch has **no cancellation**. When a timeout fires, the inner
+    ## future keeps running in the background until its I/O completes. This
+    ## ``onOrphan`` callback is registered on the inner future and called once
+    ## that orphan eventually completes. Use it to close a live connection or
+    ## release other resources the orphan holds. Without it the orphan produces
+    ## an unhandled ``FutureCompleted`` warning at best and a leaked socket /
+    ## server slot at worst.
     ##
-    ## Note: we add a no-op callback to suppress unhandled exception warnings
-    ## when the suspended inner future eventually fails.
+    ## Under chronos the ``onOrphan`` argument is accepted but never called —
+    ## futures are properly cancelled on timeout and no orphan remains.
     let ms = toMilliseconds(timeout)
     let completed = await withTimeout(fut, ms)
     if not completed:
-      # Don't call fut.fail() — the inner future is still running and would
-      # trigger "Future completed more than once" when it eventually finishes.
-      # Instead, add a callback to suppress unhandled exception warnings.
       fut.addCallback(
         proc() =
-          discard
+          if not onOrphan.isNil:
+            onOrphan(fut)
       )
       raise newException(AsyncTimeoutError, "Timeout")
     when T is void:

async_postgres/pg_connection/lifecycle.nim

@@ -517,26 +517,26 @@ proc attemptHostTimed(
     # asyncdispatch's wait() cannot cancel the attempt: on timeout it keeps
     # running in the background. If it later produces a live connection nobody
     # is waiting for it, so close the orphan instead of leaking a socket and a
-    # server slot. (chronos's wait() cancels the attempt, and connectToHost /
-    # matchesOrClose tear down their transports on the way out.)
+    # server slot. onOrphan on wait() registers the cleanup so the caller
+    # doesn't need a separate addCallback. (chronos's wait() cancels the
+    # attempt, and connectToHost / matchesOrClose tear down their transports
+    # on the way out.)
     let attempt = attemptHost(config, entry, attrs)
-    try:
-      return await attempt.wait(config.connectTimeout)
-    except AsyncTimeoutError as e:
-      proc closeOrphan() {.async.} =
-        try:
-          let orphan = attempt.read()
-          if orphan != nil:
-            await orphan.close()
-        except CatchableError:
-          discard
-
-      attempt.addCallback(
-        proc() =
-          if attempt.completed():
-            asyncSpawn closeOrphan()
-      )
-      raise e
+    return await attempt.wait(
+      config.connectTimeout,
+      onOrphan = proc(fut: Future[PgConnection]) =
+        if fut.completed():
+          asyncSpawn (
+            proc() {.async.} =
+              try:
+                let orphan = fut.read()
+                if orphan != nil:
+                  await orphan.close()
+              except CatchableError:
+                discard
+          )()
+      ,
+    )
   else:
     return await attemptHost(config, entry, attrs).wait(config.connectTimeout)
 

async_postgres/pg_connection/simple_query.nim

@@ -228,14 +228,13 @@ proc invalidateOnTimeout*(conn: PgConnection, reason: string) =
   ## `cancelNoWait`, marks the connection `csClosed` so it cannot be reused,
   ## and raises `PgTimeoutError` with `reason`.
   ##
-  ## Under asyncdispatch this is the **only** safe recovery path: the inner
-  ## future keeps running in the background after `wait()` fires, and may
-  ## still write to the socket. Reusing the connection would interleave its
-  ## stale write with a new request and corrupt the protocol stream. chronos
-  ## cancels the inner future properly, but we still invalidate unconditionally
-  ## — the server may have processed the request partially and the cached
-  ## session state (prepared statements, portals, transaction status) is no
-  ## longer reliable.
+  ## Under asyncdispatch the inner future keeps running in the background after
+  ## ``wait()`` times out (see ``wait()``'s ``onOrphan`` hook in
+  ## ``async_backend``). Reusing the connection would interleave its stale
+  ## write with a new request and corrupt the protocol stream, so we mark
+  ## ``csClosed`` unconditionally — the server may have processed the request
+  ## partially and the cached session state (prepared statements, portals,
+  ## transaction status) is no longer reliable.
   conn.cancelNoWait()
   conn.state = csClosed
   raise newException(PgTimeoutError, reason)

async_postgres/pg_pool.nim

@@ -472,29 +472,28 @@ proc spawnConnectForWaiter(pool: PgPool) =
   pool.pendingBackgroundTasks.add(fut)
   asyncSpawn fut
 
-proc closeLateConnect(pool: PgPool, connectFut: Future[PgConnection]) =
-  ## Guard against asyncdispatch's non-cancelling `wait()`: when a bounded
-  ## `connect()` times out, the inner future keeps running and may still yield
-  ## a live connection that nobody is awaiting. Register a callback to close
-  ## such an orphan rather than leaking its socket and a server-side backend
-  ## slot. No-op on chronos, whose `wait()` actually cancels the connect, so
-  ## the future never reaches `completed()`.
+proc closeLateConnect(pool: PgPool): proc(fut: Future[PgConnection]) {.gcsafe.} =
+  ## Return an ``onOrphan`` callback for ``wait()`` that closes a live
+  ## connection the orphan future produces. Under asyncdispatch the inner
+  ## future keeps running after ``wait()`` times out and may yield a
+  ## connection nobody is awaiting; this callback closes it to prevent
+  ## leaking a socket and a server-side backend slot.
+  ## Under chronos ``wait()`` cancels the inner future so ``onOrphan`` is
+  ## never called — no-op.
   when hasAsyncDispatch:
-    proc closeOrphan() {.async.} =
-      try:
-        let orphan = connectFut.read()
-        if orphan != nil:
-          await pool.tracedClose(orphan)
-      except CatchableError:
-        discard
-
-    connectFut.addCallback(
-      proc() =
-        if connectFut.completed():
-          asyncSpawn closeOrphan()
-    )
+    result = proc(fut: Future[PgConnection]) {.gcsafe.} =
+      if fut.completed():
+        asyncSpawn (
+          proc() {.async.} =
+            try:
+              let orphan = fut.read()
+              if orphan != nil:
+                await pool.tracedClose(orphan)
+            except CatchableError:
+              discard
+        )()
   else:
-    discard
+    result = nil
 
 proc maintenanceLoop(pool: PgPool) {.async.} =
   while not pool.closed:
@@ -555,7 +554,8 @@ proc maintenanceLoop(pool: PgPool) {.async.} =
         break
       let connectFut = connect(pool.config.connConfig)
       try:
-        let conn = await connectFut.wait(replenishTimeout)
+        let conn =
+          await connectFut.wait(replenishTimeout, onOrphan = pool.closeLateConnect())
         conn.ownerPool = pool
         pool.metrics.createCount.inc
         pool.consecutiveConnectFailures = 0
@@ -577,10 +577,8 @@ proc maintenanceLoop(pool: PgPool) {.async.} =
         else:
           pool.idle.addLast(PooledConn(conn: conn, lastUsedAt: now))
       except CatchableError:
-        # asyncdispatch's wait() cannot cancel connectFut: after the timeout it
-        # keeps running and may still produce a live connection nobody awaits.
-        # Close that orphan so it leaks neither a socket nor a server slot.
-        pool.closeLateConnect(connectFut)
+        # onOrphan (closeLateConnect) handles the orphan connection on timeout;
+        # here we just update failure tracking and backoff.
         pool.consecutiveConnectFailures.inc
         let delay = computeConnectBackoff(
           pool.config.connectBackoffInitial, pool.config.connectBackoffMax,

tests/test_pool.nim

@@ -2883,7 +2883,11 @@ suite "Pool replenish close-race":
         # raised, so the future is abandoned and the orphan-closer is armed,
         # exactly as the maintenance loop's except branch does.
         let connectFut = connect(mockConfig(ms.port))
-        pool.closeLateConnect(connectFut)
+        let onOrphan = pool.closeLateConnect()
+        connectFut.addCallback(
+          proc() =
+            onOrphan(connectFut)
+        )
 
         # The handshake finishes in the background; the completion callback then
         # closes the orphan.