diff --git a/async_postgres/pg_client/pipeline.nim b/async_postgres/pg_client/pipeline.nim index e5ddd00..d463d73 100644 --- a/async_postgres/pg_client/pipeline.nim +++ b/async_postgres/pg_client/pipeline.nim @@ -541,6 +541,12 @@ proc executeImpl(p: Pipeline): Future[seq[PipelineResult]] {.async.} = p.ops[activeOpIdx].cache in {scsHit, scsShare}: conn.pendingStmtCloses.add(p.ops[activeOpIdx].stmtName) conn.removeStmtCache(p.ops[activeOpIdx].sql) + elif activeOpIdx < p.ops.len and p.ops[activeOpIdx].cache == scsMiss and + not p.ops[activeOpIdx].cacheSuperseded: + # Cache-miss stmts are cached only on success, so a failed + # op's stmt is orphaned unless Closed. Close of an unparsed + # or already-Closed stmt is a harmless no-op. + conn.pendingStmtCloses.add(p.ops[activeOpIdx].stmtName) raise queryError # Cache misses: add to cache (skip ops superseded by a later # same-SQL op in this same pipeline — those stmts were already @@ -667,6 +673,9 @@ proc executeIsolatedImpl(p: Pipeline): Future[IsolatedPipelineResults] {.async.} # ReadyForQuery. conn.pendingStmtCloses.add(p.ops[opIdx].stmtName) conn.removeStmtCache(p.ops[opIdx].sql) + elif p.ops[opIdx].cache == scsMiss and not p.ops[opIdx].cacheSuperseded: + # Mirror executeImpl: Close the orphaned cache-miss stmt. + conn.pendingStmtCloses.add(p.ops[opIdx].stmtName) errors[opIdx] = opError elif p.ops[opIdx].cache == scsMiss and not p.ops[opIdx].cacheSuperseded: conn.addStmtCache( diff --git a/async_postgres/pg_pool.nim b/async_postgres/pg_pool.nim index 3121234..d54aed3 100644 --- a/async_postgres/pg_pool.nim +++ b/async_postgres/pg_pool.nim @@ -997,7 +997,9 @@ proc acquireImpl(pool: PgPool): Future[AcquireResult] {.async.} = rem, onOrphan = proc(fut: Future[PgConnection]) = if fut.completed(): - asyncSpawn ( + # Track the orphan close so pool.close()'s drain awaits it. + pool.pruneBackgroundTasks() + let closeFut = ( proc() {.async.} = try: let orphan = fut.read() @@ -1006,6 +1008,8 @@ proc acquireImpl(pool: PgPool): Future[AcquireResult] {.async.} = except CatchableError: discard )() + pool.pendingBackgroundTasks.add(closeFut) + asyncSpawn closeFut , ) else: diff --git a/tests/test_e2e_transaction.nim b/tests/test_e2e_transaction.nim index ebbb539..108a36d 100644 --- a/tests/test_e2e_transaction.nim +++ b/tests/test_e2e_transaction.nim @@ -2945,6 +2945,75 @@ suite "E2E: execInTransaction / queryInTransaction": waitFor t() + test "pipeline: failing scsMiss op queues Close so no server statement leaks": + # A cache-miss op that fails at Execute time (Parse succeeded) leaves its + # freshly Parsed statement on the server. The cache-add happens only on + # success, so nothing reuses it — without a queued Close every repeat of + # the failing SQL would pile up a new server statement (each run is a fresh + # miss with a fresh name). + proc t() {.async.} = + let conn = await connect(plainConfig()) + let failingSql = "SELECT $1::int / 0" + + proc countLeaked(): Future[int] {.async.} = + ( + await conn.simpleQuery( + "SELECT count(*)::int FROM pg_prepared_statements WHERE statement = '" & + failingSql & "'" + ) + )[0].rows[0].getInt(0) + + doAssert (await countLeaked()) == 0 + + for i in 0 ..< 3: + let p = newPipeline(conn) + p.addQuery(failingSql, @[toPgParam(0)]) + var raised = false + try: + discard await p.execute() + except PgQueryError: + raised = true + doAssert raised + # The queued Close rides along with the next extended-query op. + discard await conn.query("SELECT 1") + doAssert (await countLeaked()) == 0, + "failing scsMiss op leaked a server statement" + + await conn.close() + + waitFor t() + + test "pipeline: executeIsolated failing scsMiss op queues Close so no leak": + # executeIsolated counterpart: per-op SYNC still leaves the failed op's + # freshly Parsed statement on the server, and the cache-add happens only + # on success. + proc t() {.async.} = + let conn = await connect(plainConfig()) + let failingSql = "SELECT 1 / $1::int" + + proc countLeaked(): Future[int] {.async.} = + ( + await conn.simpleQuery( + "SELECT count(*)::int FROM pg_prepared_statements WHERE statement = '" & + failingSql & "'" + ) + )[0].rows[0].getInt(0) + + doAssert (await countLeaked()) == 0 + + for i in 0 ..< 3: + let p = newPipeline(conn) + p.addQuery(failingSql, @[toPgParam(0)]) + let ir = await p.executeIsolated() + doAssert ir.errors[0] != nil + # The queued Close rides along with the next extended-query op. + discard await conn.query("SELECT 1") + doAssert (await countLeaked()) == 0, "executeIsolated leaked a server statement" + + await conn.close() + + waitFor t() + test "pipeline: PgParam raw overload": proc t() {.async.} = let conn = await connect(plainConfig()) diff --git a/tests/test_pool.nim b/tests/test_pool.nim index 3dea072..a9092b4 100644 --- a/tests/test_pool.nim +++ b/tests/test_pool.nim @@ -1641,6 +1641,53 @@ suite "Acquire deadline budget": waitFor t() + when hasAsyncDispatch: + test "orphan connect close from timed-out acquire is tracked for close() drain": + proc t() {.async.} = + # A server that answers the startup message only after the acquire + # deadline: the caller-driven connect survives as an orphan (asyncdispatch + # has no cancellation), and its eventual close must be tracked in + # pendingBackgroundTasks so pool.close() drains it — an untracked spawn + # could leave the socket open past close(). + let ms = startMockServer() + proc serve() {.async.} = + let client = await ms.accept() + discard await client.readN(4) # StartupMessage length prefix + await sleepAsync(milliseconds(150)) + await client.sendBytes(buildAuthOk()) + await client.sendBytes(buildBackendKeyData(1234, 5678)) + await client.sendBytes(buildReadyForQuery()) + + let serveFut = serve() + + let pool = makePool() + pool.config.connConfig.host = "127.0.0.1" + pool.config.connConfig.port = ms.port + pool.config.connConfig.sslMode = sslDisable + pool.config.acquireTimeout = milliseconds(50) + + var msg = "" + try: + discard await pool.acquire() + except PgPoolError as e: + msg = e.msg + doAssert "timeout" in msg.toLowerAscii() + + # Wait for the orphan connect to complete and enqueue its close. + var waited = 0 + while pool.pendingBackgroundTasks.len == 0 and waited < 50: + await sleepAsync(milliseconds(10)) + inc waited + doAssert pool.pendingBackgroundTasks.len == 1 + + await pool.close() + doAssert pool.pendingBackgroundTasks.len == 0 + + await ms.closeServer() + await serveFut + + waitFor t() + test "nearly exhausted deadline returns idle conn unpinged": proc t() {.async.} = # With less than pingBudgetFloor (10ms) of budget left, acquire must