fix(pg_pool): reserve capacity for in-flight replenish connects to pr…

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(

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