From 7ba37f9755494c37eadd317a3aefe3106fc91237 Mon Sep 17 00:00:00 2001 From: fox0430 Date: Fri, 3 Jul 2026 17:15:40 +0900 Subject: [PATCH] Stop pool deadline timeout from poisoning a connection already handed off to a waiter --- async_postgres/pg_pool.nim | 75 ++++++++++++++++------ tests/test_e2e_transaction.nim | 110 +++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 20 deletions(-) diff --git a/async_postgres/pg_pool.nim b/async_postgres/pg_pool.nim index a5a02cd..eb2f971 100644 --- a/async_postgres/pg_pool.nim +++ b/async_postgres/pg_pool.nim @@ -1809,10 +1809,13 @@ macro withTransactionDeadline*(pool: PgPool, args: varargs[untyped]): untyped = ## pool.withTransactionDeadline(conn, TransactionOptions(...), seconds(5)): ## await conn.exec(...) ## - ## **On deadline exceeded:** if a connection was already acquired, it is - ## invalidated via `invalidateOnTimeout` (marked `csClosed`) and the pool - ## drops it on release. ROLLBACK is *not* attempted. `PgTimeoutError` is - ## raised. If the timeout fires while still waiting for `acquire()`, the + ## **On deadline exceeded:** `PgTimeoutError` is raised. Under chronos the + ## body future is cancelled: a connection with a request in flight is + ## invalidated (server-side CancelRequest, dropped on release), while one + ## that unwinds cleanly (grace ROLLBACK succeeded) returns to the pool + ## healthy. Under asyncdispatch the still-running body keeps the connection; + ## it is invalidated via `invalidateOnTimeout` and dropped on its eventual + ## release. If the timeout fires while still waiting for `acquire()`, the ## waiter remains queued (cancelled best-effort) until the underlying ## acquire future settles; this is unavoidable under asyncdispatch. ## @@ -1867,8 +1870,12 @@ macro withTransactionDeadline*(pool: PgPool, args: varargs[untyped]): untyped = let bodyFnSym = genSym(nskProc, "poolTxBodyDeadline") let bodyFutSym = genSym(nskLet, "bodyFut") let connOptSym = genSym(nskVar, "connOpt") + let releasedSym = genSym(nskVar, "released") + let cancelledSym = genSym(nskLet, "cancelled") let resetSessionSym = bindSym"resetSession" let csReadySym = bindSym"csReady" + let csClosedSym = bindSym"csClosed" + let cancelNoWaitSym = bindSym"cancelNoWait" let tsInTxSym = bindSym"tsInTransaction" let tsInFailedSym = bindSym"tsInFailedTransaction" let timeoutErrSym = bindSym"AsyncTimeoutError" @@ -1886,6 +1893,7 @@ macro withTransactionDeadline*(pool: PgPool, args: varargs[untyped]): untyped = let `totalDurSym` = `deadline` let `deadlineMomentSym` = Moment.now() + `totalDurSym` var `connOptSym` = none(PgConnection) + var `releasedSym` = false proc `bodyFnSym`(): Future[void] {.async.} = let `connIdent` = await `poolSym`.acquire() `connOptSym` = some(`connIdent`) @@ -1915,9 +1923,17 @@ macro withTransactionDeadline*(pool: PgPool, args: varargs[untyped]): untyped = `connIdent`, `ckTxRollbackSym`, `csrCleanupFailedSym`, `cleanupErrSym` ) raise `eSym` + except CancelledError as `cancelledSym`: + # Cancelled mid-request (chronos deadline): abort it server-side and + # mark csClosed so release() discards the conn instead of reusing it. + if `connIdent`.state notin {`csReadySym`, `csClosedSym`}: + `cancelNoWaitSym`(`connIdent`) + `connIdent`.state = `csClosedSym` + raise `cancelledSym` finally: await `resetSessionSym`(`poolSym`, `connIdent`) `connIdent`.release() + `releasedSym` = true let `bodyFutSym` = `bodyFnSym`() try: @@ -1931,20 +1947,20 @@ macro withTransactionDeadline*(pool: PgPool, args: varargs[untyped]): untyped = # the-same-tick should suppress the timeout report. if `bodyFutSym`.completed(): discard - elif `connOptSym`.isSome: - # invalidateOnTimeout marks the connection csClosed and raises - # PgTimeoutError — control does not return from this call. - # Ordering note: under chronos, `wait` cancels `bodyFut`, which runs - # `bodyFn`'s `finally` (resetSession + release) before this handler - # gets control. The connection has therefore already been returned - # to the pool when we invalidate it here. That is safe: the pool's - # next `acquire` health-check drops `csClosed` connections, so the - # bad conn cannot escape back to a caller. - `connOptSym`.get.`invalidateSym`("withTransactionDeadline (pool) exceeded") - else: + elif `connOptSym`.isNone: raise newException( PgTimeoutError, "withTransactionDeadline (pool): acquire timed out" ) + elif `releasedSym`: + # chronos: the cancellation already ran bodyFn's finally, so the conn + # is back in the pool — possibly handed to a waiter — and must not be + # invalidated here. bodyFn marked it csClosed before release if a + # request was in flight. + raise newException(PgTimeoutError, "withTransactionDeadline (pool) exceeded") + else: + # asyncdispatch: bodyFn still owns the conn; invalidateOnTimeout marks + # it csClosed (and raises PgTimeoutError) so release() will discard it. + `connOptSym`.get.`invalidateSym`("withTransactionDeadline (pool) exceeded") macro withTransactionRetryDeadline*( pool: PgPool, retryOpts: RetryOptions, args: varargs[untyped] @@ -1965,8 +1981,9 @@ macro withTransactionRetryDeadline*( ## the pool's health check rather than reused. Worst-case wall-clock is ## `deadline`, not `maxAttempts * deadline`. ## - ## **On deadline exceeded:** the in-flight connection (if any) is invalidated - ## and `PgTimeoutError` is raised — never retried. **On a retryable error:** + ## **On deadline exceeded:** `PgTimeoutError` is raised — never retried; the + ## in-flight connection is handled as in `withTransactionDeadline`. + ## **On a retryable error:** ## ROLLBACK runs with `rollbackGrace` and the transaction is retried if budget ## remains. See `withTransactionDeadline` for the acquire-race / `completed()` ## rationale and the in-body `conn.exec(...)` warning. **Idempotency:** `body` @@ -2002,20 +2019,28 @@ macro withTransactionRetryDeadline*( let deadlineMomentSym = genSym(nskLet, "deadlineMoment") let bodyFnSym = genSym(nskProc, "poolTxBodyRetryDeadline") let connOptSym = genSym(nskVar, "connOpt") + let releasedSym = genSym(nskVar, "released") + let cancelledSym = genSym(nskLet, "cancelled") let resetSessionSym = bindSym"resetSession" + let csReadySym = bindSym"csReady" + let csClosedSym = bindSym"csClosed" + let cancelNoWaitSym = bindSym"cancelNoWait" let remainingSym = bindSym"remainingDeadlineDuration" let graceSym = bindSym"rollbackGrace" let invalidateSym = bindSym"invalidateOnTimeout" # The pool variant acquires a fresh connection per attempt, so its cleanup # (ROLLBACK + release) happens inside bodyFn; the outer loop adds no cleanup # and omits the conn-state retry gate (connForStateCheck = nil). + # See withTransactionDeadline for the released/isNone branch rationale. let timeoutElse = quote: - if `connOptSym`.isSome: - `connOptSym`.get.`invalidateSym`("withTransactionRetryDeadline (pool) exceeded") - else: + if `connOptSym`.isNone: raise newException( PgTimeoutError, "withTransactionRetryDeadline (pool): acquire timed out" ) + elif `releasedSym`: + raise newException(PgTimeoutError, "withTransactionRetryDeadline (pool) exceeded") + else: + `connOptSym`.get.`invalidateSym`("withTransactionRetryDeadline (pool) exceeded") let loop = buildRetryDeadlineLoop( bodyFnSym, retryOptsSym, @@ -2031,8 +2056,10 @@ macro withTransactionRetryDeadline*( let `totalDurSym` = `deadline` let `deadlineMomentSym` = Moment.now() + `totalDurSym` var `connOptSym` = none(PgConnection) + var `releasedSym` = false proc `bodyFnSym`(): Future[void] {.async.} = `connOptSym` = none(PgConnection) + `releasedSym` = false let `connIdent` = await `poolSym`.acquire() `connOptSym` = some(`connIdent`) try: @@ -2047,9 +2074,17 @@ macro withTransactionRetryDeadline*( except CatchableError as `eSym`: `bodyCleanup` raise `eSym` + except CancelledError as `cancelledSym`: + # Cancelled mid-request (chronos deadline): abort it server-side and + # mark csClosed so release() discards the conn instead of reusing it. + if `connIdent`.state notin {`csReadySym`, `csClosedSym`}: + `cancelNoWaitSym`(`connIdent`) + `connIdent`.state = `csClosedSym` + raise `cancelledSym` finally: await `resetSessionSym`(`poolSym`, `connIdent`) `connIdent`.release() + `releasedSym` = true `loop` diff --git a/tests/test_e2e_transaction.nim b/tests/test_e2e_transaction.nim index 1da307d..d1636b1 100644 --- a/tests/test_e2e_transaction.nim +++ b/tests/test_e2e_transaction.nim @@ -1294,6 +1294,60 @@ suite "E2E: Deadline-bounded Transaction": waitFor t() + when hasChronos: + # Same handoff-poisoning regression guard as the withTransactionDeadline + # version above; the retry variant's timeoutElse branches identically on + # `releasedSym`, so a clean unwind must hand a healthy conn to the waiter. + test "pool.withTransactionRetryDeadline hands off healthy conn when body unwinds cleanly on deadline": + proc t() {.async.} = + let pool = + await newPool(PoolConfig(connConfig: plainConfig(), minSize: 1, maxSize: 1)) + let blocker = await pool.acquire() + + var raised = false + proc macroTask() {.async.} = + try: + pool.withTransactionRetryDeadline( + RetryOptions(maxAttempts: 5), conn, milliseconds(200) + ): + # BEGIN completes well within the deadline, leaving the conn + # csReady. The cancellation arrives at the sleepAsync await with + # the conn idle, so bodyFn's `except CancelledError` leaves it + # untouched and finally's resetSession + release hands a healthy + # conn to the queued waiter. + await sleepAsync(seconds(1)) + except PgTimeoutError: + raised = true + + let macroFut = macroTask() + # Yield so macroTask queues its acquire before the waiter. + await sleepAsync(milliseconds(5)) + + var waiterOk = false + proc waiterTask() {.async.} = + let c = await pool.acquire() + # A poisoned (csClosed) conn handed via FIFO would fail here. + let res = await c.query("SELECT 1::int") + waiterOk = res.rows.len == 1 and res.rows[0].getInt(0) == 1'i32 + c.release() + + let waiterFut = waiterTask() + + # Release blocker: conn goes to macroTask (FIFO head), waiter queued. + blocker.release() + + await macroFut + doAssert raised, "deadline should have raised PgTimeoutError" + + # macroTask released a healthy conn; waiter should complete with a + # usable connection rather than a poisoned one. + await waiterFut + doAssert waiterOk, "waiter should have received a usable connection" + + await pool.close() + + waitFor t() + test "pool.withTransactionDeadline drops invalidated connection on deadline": proc t() {.async.} = let pool = @@ -1316,6 +1370,62 @@ suite "E2E: Deadline-bounded Transaction": waitFor t() + when hasChronos: + # Under chronos, wait() cancels bodyFn and runs its finally (release) + # before the timeout handler runs. When the body unwinds cleanly (conn + # csReady), release() FIFO-hands a healthy connection to a queued waiter, + # skipping the health check. The handler must NOT then invalidate that + # already-handed-off connection — a regression that unconditionally calls + # invalidateOnTimeout would mark it csClosed *after* the handoff, leaving + # the waiter with a poisoned connection. + test "pool.withTransactionDeadline hands off healthy conn when body unwinds cleanly on deadline": + proc t() {.async.} = + let pool = + await newPool(PoolConfig(connConfig: plainConfig(), minSize: 1, maxSize: 1)) + let blocker = await pool.acquire() + + var raised = false + proc macroTask() {.async.} = + try: + pool.withTransactionDeadline(conn, milliseconds(200)): + # BEGIN completes well within the deadline, leaving the conn + # csReady. The cancellation arrives at the sleepAsync await with + # the conn idle, so bodyFn's `except CancelledError` leaves it + # untouched and finally's resetSession + release hands a healthy + # conn to the queued waiter. + await sleepAsync(seconds(1)) + except PgTimeoutError: + raised = true + + let macroFut = macroTask() + # Yield so macroTask queues its acquire before the waiter. + await sleepAsync(milliseconds(5)) + + var waiterOk = false + proc waiterTask() {.async.} = + let c = await pool.acquire() + # A poisoned (csClosed) conn handed via FIFO would fail here. + let res = await c.query("SELECT 1::int") + waiterOk = res.rows.len == 1 and res.rows[0].getInt(0) == 1'i32 + c.release() + + let waiterFut = waiterTask() + + # Release blocker: conn goes to macroTask (FIFO head), waiter queued. + blocker.release() + + await macroFut + doAssert raised, "deadline should have raised PgTimeoutError" + + # macroTask released a healthy conn; waiter should complete with a + # usable connection rather than a poisoned one. + await waiterFut + doAssert waiterOk, "waiter should have received a usable connection" + + await pool.close() + + waitFor t() + test "pool.withTransactionDeadline raises PgTimeoutError when acquire times out": proc t() {.async.} = # maxSize=1: a single held connection forces the next acquire to queue.