fix: guard pool advisory-lock reset with a sticky dirty flag

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,13 +72,15 @@ 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
 ): bool =
   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
 

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"))

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

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

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]

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