diff --git a/async_postgres/pg_advisory_lock.nim b/async_postgres/pg_advisory_lock.nim index 146b0cd..bb0063b 100644 --- a/async_postgres/pg_advisory_lock.nim +++ b/async_postgres/pg_advisory_lock.nim @@ -21,12 +21,12 @@ ## them with the same key unintentionally. Transaction-level locks are not ## stackable and are always released at transaction end. ## -## **Pool integration:** Session-level lock acquires through this typed API -## bump a per-connection counter, and the pool releases or discards the -## connection on return so that locks never leak to subsequent borrowers. -## Raw-SQL acquires (e.g. ``conn.exec("SELECT pg_advisory_lock(1)")``) bypass -## this tracking — callers must release them explicitly or invoke -## ``advisoryUnlockAll`` before returning the connection to the pool. +## **Pool integration:** Typed acquires set a sticky ``sessionLockDirty`` +## flag; the pool runs ``pg_advisory_unlock_all`` on return based on the +## flag (not the counter), so a typed ``advisoryUnlock`` of a raw-acquired +## key cannot forge a lock-free state. Raw-SQL acquires still bypass +## tracking — callers must release them explicitly or invoke +## ``advisoryUnlockAll`` before returning the connection. ## ## Example ## ======= @@ -72,6 +72,7 @@ template acquireSessionLock( ) = discard await conn.queryValue(sql, params, timeout = t) inc conn.heldSessionLocks + conn.sessionLockDirty = true template trySessionLock( conn: PgConnection, sql: string, params: seq[PgParam], t: Duration @@ -79,6 +80,7 @@ template trySessionLock( let acquired = await conn.queryValue(bool, sql, params, timeout = t) if acquired: inc conn.heldSessionLocks + conn.sessionLockDirty = true acquired template unlockSessionLock( @@ -169,6 +171,7 @@ proc advisoryUnlockAll*( ## Release all session-level advisory locks held by the current session. discard await conn.exec("SELECT pg_advisory_unlock_all()", timeout = timeout) conn.heldSessionLocks = 0 + conn.sessionLockDirty = false # Transaction-level exclusive locks diff --git a/async_postgres/pg_connection/lifecycle.nim b/async_postgres/pg_connection/lifecycle.nim index f1f7e40..dc430da 100644 --- a/async_postgres/pg_connection/lifecycle.nim +++ b/async_postgres/pg_connection/lifecycle.nim @@ -460,6 +460,7 @@ proc close*(conn: PgConnection): Future[void] {.async.} = discard conn.state = csClosed conn.heldSessionLocks = 0 + conn.sessionLockDirty = false # Fail any pending notification waiter if conn.notifyWaiter != nil and not conn.notifyWaiter.finished: conn.notifyWaiter.fail(newException(PgError, "Connection closed")) diff --git a/async_postgres/pg_connection/types.nim b/async_postgres/pg_connection/types.nim index 2097006..487063d 100644 --- a/async_postgres/pg_connection/types.nim +++ b/async_postgres/pg_connection/types.nim @@ -266,9 +266,19 @@ type ## to the gap until the next operation. heldSessionLocks*: int ## Count of session-level `pg_advisory_lock` acquires through the typed - ## API. The pool releases or discards connections with a non-zero count - ## so that locks never leak to subsequent borrowers. Raw-SQL acquires - ## (`conn.exec("SELECT pg_advisory_lock(...)")`) bypass this counter. + ## API, minus tracked releases. Reported by the `onLeakedSessionLocks` + ## tracer hook. Raw-SQL acquires + ## (`conn.exec("SELECT pg_advisory_lock(...)")`) bypass this counter, + ## so it is best-effort under mixed typed/raw usage — the pool's + ## reset/discard decision uses `sessionLockDirty` instead so a raw + ## acquire released through the typed API cannot forge a `== 0` count + ## and leak a tracked lock into the next borrower. + sessionLockDirty*: bool + ## Sticky flag: set on any tracked session-level acquire, cleared only + ## by `advisoryUnlockAll` or connection reset. Drives the pool's + ## reset/discard decision so `pg_advisory_unlock_all` runs whenever a + ## tracked acquire ever happened, even if the tracked counter was + ## decremented back to zero by a typed unlock of a raw-acquired key. tracer*: PgTracer ## Inherited from ConnConfig on connect ownerPool*: PgPoolOwner ## Owning pool back-reference. Set when this connection is managed by diff --git a/async_postgres/pg_pool.nim b/async_postgres/pg_pool.nim index 98efa80..3cb7302 100644 --- a/async_postgres/pg_pool.nim +++ b/async_postgres/pg_pool.nim @@ -308,17 +308,18 @@ proc resetSession*(pool: PgPool, conn: PgConnection) {.async.} = ## tracer hook — are swallowed. if conn.state != csReady or conn.txStatus != tsIdle: return - if pool.config.resetQuery.len == 0 and conn.heldSessionLocks == 0: + if pool.config.resetQuery.len == 0 and not conn.sessionLockDirty: return try: - if conn.heldSessionLocks > 0: + if conn.sessionLockDirty: let t = pool.config.tracer - if t != nil and t.onLeakedSessionLocks != nil: + if t != nil and t.onLeakedSessionLocks != nil and conn.heldSessionLocks > 0: t.onLeakedSessionLocks( TraceLeakedSessionLocksData(conn: conn, count: conn.heldSessionLocks) ) discard await conn.simpleExec("SELECT pg_advisory_unlock_all()") conn.heldSessionLocks = 0 + conn.sessionLockDirty = false if pool.config.resetQuery.len > 0: discard await conn.simpleExec(pool.config.resetQuery) conn.clearStmtCache() @@ -623,7 +624,7 @@ proc releaseCore( ## `releaseImpl`. Returns flags describing the disposition of `conn` so ## the caller can report them to the tracer. if pool.closed or conn.state != csReady or conn.txStatus != tsIdle or - conn.heldSessionLocks > 0: + conn.sessionLockDirty: if pool.active > 0: pool.active.dec pool.closeNoWait(conn) @@ -665,10 +666,10 @@ proc releaseImpl(pool: PgPool, conn: PgConnection) = ## Transaction-in-progress (`txStatus != tsIdle`) is treated as failure ## to reset the session, so the connection is closed rather than leaking ## transaction state to the next borrower. - ## Session-level advisory locks (`heldSessionLocks > 0`) likewise force the + ## Session-level advisory locks (`sessionLockDirty`) likewise force the ## connection to be discarded: callers who route through `resetSession` - ## clear them ahead of time, so anything reaching here with locks still - ## held has bypassed that path and must not return to the idle queue. + ## clear them ahead of time, so anything reaching here still dirty has + ## bypassed that path and must not return to the idle queue. ## ## Double-release guard: a connection that is not currently checked out ## (`borrowed == false`) has already been returned to the pool — or never @@ -694,7 +695,8 @@ proc releaseImpl(pool: PgPool, conn: PgConnection) = tracer.onPoolDoubleRelease(TracePoolDoubleReleaseData(conn: conn)) return conn.borrowed = false - if conn.heldSessionLocks > 0 and tracer != nil and tracer.onLeakedSessionLocks != nil: + if conn.sessionLockDirty and conn.heldSessionLocks > 0 and tracer != nil and + tracer.onLeakedSessionLocks != nil: tracer.onLeakedSessionLocks( TraceLeakedSessionLocksData(conn: conn, count: conn.heldSessionLocks) ) diff --git a/tests/test_advisory_lock.nim b/tests/test_advisory_lock.nim index af53904..e6892e2 100644 --- a/tests/test_advisory_lock.nim +++ b/tests/test_advisory_lock.nim @@ -547,7 +547,10 @@ suite "Advisory Lock: pool integration": waitFor t() - test "explicit unlock keeps connection in the pool": + test "explicit unlock alone does not spare direct release from discarding": + # advisoryUnlock decrements the counter but not the sticky dirty flag, + # so a direct release() (bypassing resetSession) still discards the + # connection. Use advisoryUnlockAll or withConnection to keep it. proc t() {.async.} = let cfg = initPoolConfig(plainConfig(), minSize = 0, maxSize = 1) let pool = await newPool(cfg) @@ -558,10 +561,11 @@ suite "Advisory Lock: pool integration": await conn.advisoryLock(71002'i64) discard await conn.advisoryUnlock(71002'i64) doAssert conn.heldSessionLocks == 0 + doAssert conn.sessionLockDirty let closeBefore = pool.metrics.closeCount conn.release() - doAssert pool.idle.len == 1 - doAssert pool.metrics.closeCount == closeBefore + doAssert pool.idle.len == 0 + doAssert pool.metrics.closeCount - closeBefore == 1 waitFor t() @@ -626,6 +630,63 @@ suite "Advisory Lock: pool integration": waitFor t() + test "typed unlock of raw-acquired key does not leak the tracked lock": + # Mixed-usage regression: raw pg_advisory_lock followed by typed + # advisoryUnlock decrements the counter but must not clear the sticky + # dirty flag, so the still-held tracked lock forces a pool-side cleanup + # rather than being handed to the next borrower. + proc t() {.async.} = + let cfg = initPoolConfig(plainConfig(), minSize = 0, maxSize = 1) + let pool = await newPool(cfg) + defer: + await pool.close() + + let conn = await pool.acquire() + await conn.advisoryLock(71007'i64) # tracked K1 + discard await conn.exec("SELECT pg_advisory_lock(71008)") # raw K2 + let released = await conn.advisoryUnlock(71008'i64) # typed unlock of raw K2 + doAssert released + doAssert conn.heldSessionLocks == 0 # counter stolen, but… + doAssert conn.sessionLockDirty # …dirty flag survives. + + let closeBefore = pool.metrics.closeCount + conn.release() + doAssert pool.idle.len == 0 + doAssert pool.metrics.closeCount - closeBefore == 1 + + # K1 is released server-side by the close, so a fresh session can take it. + let probe = await connect(plainConfig()) + defer: + await probe.close() + doAssert await probe.advisoryTryLock(71007'i64) + discard await probe.advisoryUnlock(71007'i64) + + waitFor t() + + test "typed unlock of raw-acquired key: resetSession path releases via unlock_all": + proc t() {.async.} = + let cfg = initPoolConfig(plainConfig(), minSize = 0, maxSize = 1) + let pool = await newPool(cfg) + defer: + await pool.close() + + pool.withConnection(conn): + await conn.advisoryLock(71009'i64) # tracked K1 + discard await conn.exec("SELECT pg_advisory_lock(71010)") # raw K2 + discard await conn.advisoryUnlock(71010'i64) # steals counter to 0 + # resetSession keyed off dirty flag → unlock_all cleared K1 too. + doAssert pool.idle.len == 1 + doAssert not pool.idle[0].conn.sessionLockDirty + doAssert pool.idle[0].conn.heldSessionLocks == 0 + + let probe = await connect(plainConfig()) + defer: + await probe.close() + doAssert await probe.advisoryTryLock(71009'i64) + discard await probe.advisoryUnlock(71009'i64) + + waitFor t() + suite "Advisory Lock: onLeakedSessionLocks tracer hook": type LeakLog = ref object counts: seq[int] diff --git a/tests/test_pool.nim b/tests/test_pool.nim index b897536..d52bf43 100644 --- a/tests/test_pool.nim +++ b/tests/test_pool.nim @@ -503,6 +503,7 @@ suite "Pool release": pool.active = 1 let conn = mockConn() conn.heldSessionLocks = 1 + conn.sessionLockDirty = true pool.release(conn) check pool.active == 0 check pool.idle.len == 0