From 7f4d581424e72e6948478542aa613ad1e3648465 Mon Sep 17 00:00:00 2001 From: fox0430 Date: Sun, 5 Jul 2026 17:48:34 +0900 Subject: [PATCH] Consolidate orphan-hook bounded-wait primitive into async_backend.nim --- async_postgres/async_backend.nim | 43 ++++++++++------ async_postgres/pg_connection/lifecycle.nim | 38 +++++++------- async_postgres/pg_connection/simple_query.nim | 15 +++--- async_postgres/pg_pool.nim | 50 +++++++++---------- tests/test_pool.nim | 6 ++- 5 files changed, 82 insertions(+), 70 deletions(-) diff --git a/async_postgres/async_backend.nim b/async_postgres/async_backend.nim index 20f19af..9df7c47 100644 --- a/async_postgres/async_backend.nim +++ b/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: diff --git a/async_postgres/pg_connection/lifecycle.nim b/async_postgres/pg_connection/lifecycle.nim index 0c768ff..f1f7e40 100644 --- a/async_postgres/pg_connection/lifecycle.nim +++ b/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) diff --git a/async_postgres/pg_connection/simple_query.nim b/async_postgres/pg_connection/simple_query.nim index 281a610..d4c4d18 100644 --- a/async_postgres/pg_connection/simple_query.nim +++ b/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) diff --git a/async_postgres/pg_pool.nim b/async_postgres/pg_pool.nim index f6cd299..99e040b 100644 --- a/async_postgres/pg_pool.nim +++ b/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, diff --git a/tests/test_pool.nim b/tests/test_pool.nim index 16dde45..7a4c1f4 100644 --- a/tests/test_pool.nim +++ b/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.