From 3c012da721bb6a767eb7ddd2e0a62ff3e839b1c7 Mon Sep 17 00:00:00 2001 From: fox0430 Date: Fri, 14 Aug 2026 17:41:41 +0900 Subject: [PATCH] fix(pg_pool): reserve capacity for in-flight replenish connects to prevent maxSize overshoot --- async_postgres/pg_pool.nim | 52 +++++++++++----- tests/test_pool.nim | 120 +++++++++++++++++++++++++++++++++++++ 2 files changed, 156 insertions(+), 16 deletions(-) diff --git a/async_postgres/pg_pool.nim b/async_postgres/pg_pool.nim index c301c85..0720561 100644 --- a/async_postgres/pg_pool.nim +++ b/async_postgres/pg_pool.nim @@ -525,6 +525,24 @@ proc respawnForStrandedWaiter(pool: PgPool) = pool.active.inc pool.spawnConnectForWaiter() +proc settleReplenishConnect(pool: PgPool, conn: PgConnection, now: Moment) = + ## Settle a fresh replenishment connection: the pre-reserved `pool.active` + ## slot is consumed on waiter handoff, released on park or close. + conn.ownerPool = pool + pool.metrics.createCount.inc + pool.consecutiveConnectFailures = 0 + # The pool may have closed while awaiting connect — close instead of + # parking. closeNoWait (not `await`) so a pending cancellation can't + # interrupt the close. + if pool.closed: + pool.closeNoWait(conn) + pool.active.dec + elif pool.tryHandoffToWaiter(conn): + discard # reservation stays consumed: the conn is now the waiter's + else: + pool.active.dec + pool.idle.addLast(PooledConn(conn: conn, lastUsedAt: now)) + proc maintenanceLoop(pool: PgPool) {.async.} = while not pool.closed: await sleepAsync(pool.config.maintenanceInterval) @@ -579,31 +597,33 @@ proc maintenanceLoop(pool: PgPool) {.async.} = if needed > 0: if pool.closed: break + # Reserve the in-flight capacity up front: uncounted replenish connects + # would let concurrent acquires overshoot maxSize once these park. + pool.active.inc(needed) var connectFuts: seq[Future[PgConnection]] for i in 0 ..< needed: var connCfg = pool.config.connConfig if connCfg.connectTimeout == ZeroDuration: connCfg.connectTimeout = pool.config.maintenanceInterval connectFuts.add(connect(connCfg)) - await allFutures(connectFuts) - for f in connectFuts: - if not f.failed(): - let conn = f.read() - conn.ownerPool = pool - pool.metrics.createCount.inc - pool.consecutiveConnectFailures = 0 - # The pool may have been closed while we awaited connect. Parking a - # fresh conn in a closed pool leaks its socket — re-check and close. - # closeNoWait (not `await tracedClose`) because a pending cancellation - # would interrupt a fresh await before the close runs. - if pool.closed: - pool.closeNoWait(conn) - elif pool.tryHandoffToWaiter(conn): - pool.active.inc + try: + await allFutures(connectFuts) + for f in connectFuts: + if not f.failed(): + pool.settleReplenishConnect(f.read(), now) + except CancelledError as e: + # close() cancels us mid-connect: settle the conns that already + # landed, release the rest, then re-raise so cancelAndWait + # observes completion. + for f in connectFuts: + if not f.finished or f.failed(): + pool.active.dec else: - pool.idle.addLast(PooledConn(conn: conn, lastUsedAt: now)) + pool.settleReplenishConnect(f.read(), now) + raise e for f in connectFuts: if f.failed(): + pool.active.dec pool.consecutiveConnectFailures.inc if pool.consecutiveConnectFailures > 0: let delay = computeConnectBackoff( diff --git a/tests/test_pool.nim b/tests/test_pool.nim index 8a50786..f13a9dd 100644 --- a/tests/test_pool.nim +++ b/tests/test_pool.nim @@ -3612,6 +3612,126 @@ suite "Pool replenish close-race": waitFor t() +suite "Pool replenish capacity race": + test "in-flight replenish connects hold a reservation so concurrent acquires cannot overshoot maxSize": + # Regression: the maintenance loop once opened replenish connects without + # counting them in `pool.active` (the only violation of the reservation + # contract `spawnConnectForWaiter` documents). A burst of acquires + # arriving mid-replenish saw spare capacity (`active < maxSize`), dialed + # their own connections, and the replenish connections parked on top — + # pushing `idle + active` to 2x maxSize in a fixed pool (minSize == + # maxSize) and breaching the DB's per-user connection limit. The + # in-flight connects must hold capacity reservations, queuing the burst + # as waiters until they settle. + proc t() {.async.} = + let ms = startMockServer() + + let pool = makePool(minSize = 2, maxSize = 2) + pool.config.connConfig = mockConfig(ms.port) + pool.config.connConfig.connectTimeout = seconds(5) + pool.config.maintenanceInterval = milliseconds(10) + pool.maintenanceTask = maintenanceLoop(pool) + + # The loop sleeps one interval, then opens the replenish connects. Gate + # the handshake server-side so both connects are provably in flight with + # their reservations held (with the OLD code active would be 0 here). + let client1 = await ms.accept() + await drainStartupMessage(client1) + let client2 = await ms.accept() + await drainStartupMessage(client2) + doAssert pool.active == 2 # the in-flight replenish reservations + + # The burst must queue as waiters rather than dialing its own connects. + let futA = pool.acquire() + let futB = pool.acquire() + await sleepAsync(milliseconds(20)) + doAssert pool.waiterCount == 2 + doAssert pool.active == 2 + doAssert pool.metrics.createCount == 0 # nothing has settled yet + + # Complete the handshakes: the replenish conns hand off to the waiters. + await sendFullHandshake(client1) + await sendFullHandshake(client2) + + let connA = await futA + let connB = await futB + doAssert pool.active == 2 + doAssert pool.idle.len == 0 + doAssert pool.active + pool.idle.len <= pool.config.maxSize + doAssert pool.metrics.createCount == 2 # exactly the replenish connects + + pool.release(connA) + pool.release(connB) + doAssert pool.active == 0 + doAssert pool.idle.len == 2 + + await pool.close() + await closeServer(ms) + + waitFor t() + + test "replenish connect failures release their reservations": + # The replenish reservations must be released when the connects fail, or + # `active` would stay inflated and `acquire` would queue behind phantom + # capacity. + proc t() {.async.} = + let pool = makePool(minSize = 2, maxSize = 4) + pool.config.connConfig = mockConfig(1) # nothing listens on port 1 + pool.config.maintenanceInterval = milliseconds(10) + pool.maintenanceTask = maintenanceLoop(pool) + + await sleepAsync(milliseconds(60)) + + doAssert pool.active == 0 # reservations released despite 2 failed connects + doAssert pool.idle.len == 0 + doAssert pool.metrics.createCount == 0 + + pool.closed = true + await cancelAndWait(pool.maintenanceTask) + + waitFor t() + + when hasChronos: + test "close cancelling in-flight replenish releases its reservations (landed and in-flight)": + # chronos-only: one replenish connect lands before close() cancels the + # loop (settled via the except path); the other stays in flight. + proc t() {.async.} = + let ms = startMockServer() + + let pool = makePool(minSize = 2, maxSize = 2) + pool.config.connConfig = mockConfig(ms.port) + pool.config.connConfig.connectTimeout = seconds(5) + pool.config.maintenanceInterval = milliseconds(10) + pool.maintenanceTask = maintenanceLoop(pool) + + let client1 = await ms.accept() + await drainStartupMessage(client1) + let client2 = await ms.accept() + await drainStartupMessage(client2) + doAssert pool.active == 2 # in-flight replenish reservations + + await sendFullHandshake(client1) + await sleepAsync(milliseconds(100)) # let the client-side connect settle + + await pool.close() + doAssert pool.active == 0 + doAssert pool.metrics.closeCount == 1 # landed conn settled via closeNoWait + + # Read the landed conn until EOF: the pool must have closed it. + var closedByPool = false + try: + while true: + discard await client1.readN(1) + except CatchableError: + closedByPool = true + doAssert closedByPool + + await closeServer(ms) + await closeClient(client1) + await closeClient(client2) + + waitFor t() + suite "Pool acquire close-race": test "acquire discards a connection won after the pool is closed": # Regression: acquireImpl's `await connect()` in the new-conn branch could