fix(postgres): survive a database restart without leaking connections (#4149)

logos_delivery/waku/common/databases/db_postgres/dbconn.nim

@@ -13,6 +13,9 @@ type DataProc* = proc(result: ptr PGresult) {.closure, gcsafe, raises: [].}
 
 type DbConnWrapper* = ref object
   dbConn: DbConn
+  registeredFd: Opt[asyncengine.AsyncFD]
+    ## the descriptor handed to chronos when the connection was opened. It has
+    ## to be remembered because libpq no longer knows it once the backend dies.
   open: bool
   preparedStmts: HashSet[string] ## [stmtName's]
   futBecomeFree*: Future[void]
@@ -29,6 +32,10 @@ proc inclPreparedStmt*(dbConnWrapper: DbConnWrapper, preparedStmt: string) =
 proc getDbConn*(dbConnWrapper: DbConnWrapper): DbConn =
   return dbConnWrapper.dbConn
 
+proc getRegisteredFd*(dbConnWrapper: DbConnWrapper): Opt[asyncengine.AsyncFD] =
+  ## Exposed so that the tests can assert the selector entry is given back.
+  return dbConnWrapper.registeredFd
+
 proc isPgDbConnBusy*(dbConnWrapper: DbConnWrapper): bool =
   if isNil(dbConnWrapper.futBecomeFree):
     return false
@@ -37,9 +44,6 @@ proc isPgDbConnBusy*(dbConnWrapper: DbConnWrapper): bool =
 proc isPgDbConnOpen*(dbConnWrapper: DbConnWrapper): bool =
   return dbConnWrapper.open
 
-proc setPgDbConnOpen*(dbConnWrapper: DbConnWrapper, newOpenState: bool) =
-  dbConnWrapper.open = newOpenState
-
 const MaxDbErrorLen = 512
   ## libpq can answer with very long messages -- the DETAIL and CONTEXT lines
   ## carry row data -- and this string reaches both the logs and the error chain,
@@ -61,7 +65,7 @@ proc check(db: DbConn): Result[void, string] =
 
   return ok()
 
-proc openDbConn(connString: string): Result[DbConn, string] =
+proc openDbConn(connString: string): Result[DbConnWrapper, string] =
   ## Opens a new connection.
   var conn: DbConn = nil
   try:
@@ -75,28 +79,48 @@ proc openDbConn(connString: string): Result[DbConn, string] =
 
     return err("unknown reason")
 
-  ## registering the socket fd in chronos for better wait for data
+  ## registering the socket fd in chronos for better wait for data.
+  ## The wrapper is built here so that the registered descriptor and the
+  ## remembered one cannot drift apart.
   let asyncFd = cast[asyncengine.AsyncFD](pqsocket(conn))
-  asyncengine.register(asyncFd)
+  asyncengine.register2(asyncFd).isOkOr:
+    conn.close()
+    return err("failed to register the connection socket: " & $error)
 
-  return ok(conn)
+  return ok(DbConnWrapper(dbConn: conn, registeredFd: Opt.some(asyncFd), open: true))
 
 proc new*(T: type DbConnWrapper, connString: string): Result[T, string] =
-  let dbConn = openDbConn(connString).valueOr:
+  let dbConnWrapper = openDbConn(connString).valueOr:
     return err("failed to establish a new connection: " & $error)
 
-  return ok(DbConnWrapper(dbConn: dbConn, open: true))
+  return ok(dbConnWrapper)
 
-proc closeDbConn*(
-    dbConnWrapper: DbConnWrapper
-): Result[void, string] {.raises: [OSError].} =
-  let fd = dbConnWrapper.dbConn.pqsocket()
-  if fd == -1:
-    return err("error file descriptor -1 in closeDbConn")
+proc closeDbConn*(dbConnWrapper: DbConnWrapper): Result[void, string] {.raises: [].} =
+  ## Closing must always reach pqfinish, even when giving the selector entry
+  ## back fails. Asking libpq for the descriptor here is useless: it answers -1
+  ## as soon as the backend is gone.
+  if not dbConnWrapper.open:
+    return ok()
 
-  asyncengine.unregister(cast[asyncengine.AsyncFD](fd))
+  var unregisterError = ""
+  if dbConnWrapper.registeredFd.isSome():
+    let asyncFd = dbConnWrapper.registeredFd.get()
+    ## unregister2 asserts when the descriptor is unknown to the dispatcher
+    if asyncFd in asyncengine.getThreadDispatcher():
+      when defined(windows):
+        ## chronos exposes no Result-returning unregister on Windows, where
+        ## unregistering only drops the handle from the dispatcher and cannot
+        ## fail
+        asyncengine.unregister(asyncFd)
+      else:
+        asyncengine.unregister2(asyncFd).isOkOr:
+          unregisterError = "failed to unregister the connection socket: " & $error
 
   dbConnWrapper.dbConn.close()
+  dbConnWrapper.open = false
+
+  if unregisterError.len > 0:
+    return err(unregisterError)
 
   return ok()
 
@@ -173,31 +197,47 @@ proc sendQueryPrepared(
 
   return ok()
 
+proc waitForData(
+    dbConnWrapper: DbConnWrapper, asyncFd: asyncengine.AsyncFD
+): Future[Result[void, string]] {.async.} =
+  ## Waits until the socket has something to read and gives the reader back
+  ## before returning. The caller must not touch libpq while the reader is
+  ## installed: libpq closes the socket as soon as it notices the backend is
+  ## gone, and removing a reader from an already closed descriptor fails,
+  ## leaving the selector entry behind forever.
+
+  when defined(windows):
+    return err("Postgres not supported on Windows")
+  else:
+    let futDataAvailable = newFuture[void]("futDataAvailable")
+
+    proc onDataAvailable(udata: pointer) {.gcsafe, raises: [].} =
+      if not futDataAvailable.completed():
+        futDataAvailable.complete()
+
+    asyncengine.addReader2(asyncFd, onDataAvailable).isOkOr:
+      dbConnWrapper.futBecomeFree.fail(newException(ValueError, $error))
+      return err("failed to add event reader in waitForData: " & $error)
+
+    defer:
+      asyncengine.removeReader2(asyncFd).isOkOr:
+        error "failed to remove event reader in waitForData", error = $error
+
+    await futDataAvailable
+
+    return ok()
+
 proc waitQueryToFinish(
     dbConnWrapper: DbConnWrapper, rowCallback: DataProc = nil
 ): Future[Result[void, string]] {.async.} =
   ## The 'rowCallback' param is != nil when the underlying query wants to retrieve results (SELECT.)
   ## For other queries, like "INSERT", 'rowCallback' should be nil.
 
-  let futDataAvailable = newFuture[void]("futDataAvailable")
+  let asyncFd = dbConnWrapper.registeredFd.valueOr:
+    return err("the connection socket is not registered in waitQueryToFinish")
 
-  proc onDataAvailable(udata: pointer) {.gcsafe, raises: [].} =
-    if not futDataAvailable.completed():
-      futDataAvailable.complete()
-
-  let asyncFd = cast[asyncengine.AsyncFD](pqsocket(dbConnWrapper.dbConn))
-
-  when not defined(windows):
-    asyncengine.addReader2(asyncFd, onDataAvailable).isOkOr:
-      dbConnWrapper.futBecomeFree.fail(newException(ValueError, $error))
-      return err("failed to add event reader in waitQueryToFinish: " & $error)
-    defer:
-      asyncengine.removeReader2(asyncFd).isOkOr:
-        return err("failed to remove event reader in waitQueryToFinish: " & $error)
-  else:
-    return err("Postgres not supported on Windows")
-
-  await futDataAvailable
+  (await dbConnWrapper.waitForData(asyncFd)).isOkOr:
+    return err($error)
 
   ## Now retrieve the result from the database
   while true:

logos_delivery/waku/common/databases/db_postgres/pgasyncpool.nim

@@ -58,17 +58,24 @@ proc close*(pool: PgAsyncPool): Future[Result[void, string]] {.async.} =
   # blocking the async runtime
 
   debug "close PgAsyncPool"
-  await allFutures(pool.conns.mapIt(it.futBecomeFree))
+  ## a connection that never ran a query has no futBecomeFree to wait for
+  await allFutures(
+    pool.conns.filterIt(not it.futBecomeFree.isNil()).mapIt(it.futBecomeFree)
+  )
   debug "closing all connection PgAsyncPool"
 
+  var closeErrors = newSeq[string](0)
   for i in 0 ..< pool.conns.len:
-    if pool.conns[i].isPgDbConnOpen():
-      pool.conns[i].closeDbConn().isOkOr:
-        return err("error in close PgAsyncPool: " & $error)
-      pool.conns[i].setPgDbConnOpen(false)
+    ## one connection that cannot be closed must not keep the others open
+    ## and registered in the dispatcher
+    pool.conns[i].closeDbConn().isOkOr:
+      closeErrors.add($error)
 
   pool.conns.setLen(0)
 
+  if closeErrors.len > 0:
+    return err("error in close PgAsyncPool: " & closeErrors.join("; "))
+
   return ok()
 
 proc getFirstFreeConnIndex(pool: PgAsyncPool): DatabaseResult[int] =

logos_delivery/waku/waku_archive/driver/postgres_driver/postgres_healthcheck.nim

@@ -1,6 +1,6 @@
 {.push raises: [].}
 
-import chronos, results
+import chronos, chronicles, results
 import ../../../common/databases/db_postgres, ../../../common/error_handling
 
 ## Simple query to validate that the postgres is working and attending requests
@@ -20,7 +20,9 @@ proc checkConnectivity*(
       block errorBlock:
         ## Force close all the opened connections. No need to close gracefully.
         (await connPool.resetConnPool()).isOkOr:
-          onFatalErrorAction("checkConnectivity resetConnPool error: " & error)
+          ## Not fatal on its own. The trials below are the ones that tell
+          ## whether the database is really unreachable.
+          error "checkConnectivity resetConnPool error", error = error
 
         var numTrial = 0
         while numTrial < MaxNumTrials:

tests/waku_archive/test_driver_postgres.nim

@@ -1,12 +1,14 @@
 {.used.}
 
-import results, std/[sequtils, strutils], testutils/unittests, chronos
+import
+  results, std/[sequtils, strutils], testutils/unittests, chronos, db_connector/postgres
 import
   logos_delivery/waku/[
     waku_archive,
     waku_archive/driver/postgres_driver,
     waku_core,
     waku_core/message/digest,
+    common/databases/db_postgres/dbconn,
     common/databases/db_postgres/pgasyncpool,
   ],
   ../testlib/wakucore,
@@ -373,3 +375,154 @@ suite "Postgres driver - concurrent DDL outcomes":
     check not "ERROR:  could not acquire advisory lock".isConcurrentDdlOutcome()
     check not "ERROR:  no space left on device".isConcurrentDdlOutcome()
     check not "ERROR:  deadlock detected".isConcurrentDdlOutcome()
+
+suite "Postgres connection lifecycle":
+  ## A database restart used to take the node down: closing a connection whose
+  ## backend was gone failed, which left the remaining connections open and
+  ## their descriptors registered in the chronos dispatcher.
+
+  const RawConnString =
+    "user=postgres host=localhost port=5432 dbname=postgres password=test123"
+
+  proc backendPidsAlive(
+      pool: PgAsyncPool, pids: seq[string]
+  ): Future[Result[int, string]] {.async.} =
+    var alive = 0
+
+    proc onRow(res: ptr PGresult) {.closure, gcsafe, raises: [].} =
+      if pqntuples(res) > 0:
+        try:
+          alive = parseInt($pqgetvalue(res, 0, 0))
+        except ValueError:
+          discard
+
+    (
+      await pool.pgQuery(
+        "SELECT count(*) FROM pg_stat_activity WHERE pid IN (" & pids.join(",") & ")",
+        @[],
+        onRow,
+      )
+    ).isOkOr:
+      return err($error)
+
+    return ok(alive)
+
+  asyncTest "A connection whose backend died gives its selector entry back":
+    let wrapper = DbConnWrapper.new(RawConnString).expect("new connection")
+
+    let asyncFd = wrapper.getRegisteredFd().expect("registered fd")
+    check asyncFd in getThreadDispatcher()
+
+    ## the backend kills itself, so libpq closes the socket while the results
+    ## of this very query are being read
+    let killRes = await wrapper.dbConnQuery(
+      sql("SELECT pg_terminate_backend(pg_backend_pid())"), @[], nil, ""
+    )
+    check killRes.isErr()
+
+    ## precondition: libpq cannot tell the descriptor anymore
+    check pqsocket(wrapper.getDbConn()) == -1
+
+    check wrapper.closeDbConn().isOk()
+    check asyncFd notin getThreadDispatcher()
+
+  asyncTest "A pool closes every connection even when one is dead":
+    let pool = PgAsyncPool.new(storeMessageDbUrl, 3).expect("pool")
+    let observer = PgAsyncPool.new(storeMessageDbUrl, 1).expect("observer pool")
+
+    var pids = newSeq[string](3)
+
+    proc collectPid(index: int): DataProc =
+      return proc(res: ptr PGresult) {.closure, gcsafe, raises: [].} =
+        if pqntuples(res) > 0:
+          pids[index] = $pqgetvalue(res, 0, 0)
+
+    ## three overlapping queries force the pool to open three connections
+    var queries = newSeq[Future[Result[void, string]]](0)
+    for i in 0 ..< 3:
+      queries.add(
+        pool.pgQuery("SELECT pg_backend_pid(), pg_sleep(0.5)", @[], collectPid(i))
+      )
+
+    for queryFut in queries:
+      (await queryFut).expect("concurrent pid query")
+
+    check pids.deduplicate().len == 3
+
+    ## an idle pool always hands out its first connection, so this tells which
+    ## backend has to die for the close loop to fail on its first iteration
+    var firstConnPid: string
+    proc onFirstPid(res: ptr PGresult) {.closure, gcsafe, raises: [].} =
+      if pqntuples(res) > 0:
+        firstConnPid = $pqgetvalue(res, 0, 0)
+
+    (await pool.pgQuery("SELECT pg_backend_pid()", @[], onFirstPid)).expect(
+      "first conn pid"
+    )
+
+    let survivorPids = pids.filterIt(it != firstConnPid)
+    check survivorPids.len == 2
+
+    (await observer.pgQuery("SELECT pg_terminate_backend(" & firstConnPid & ")")).expect(
+      "terminate backend"
+    )
+
+    ## the pool only learns about the dead backend when it uses it again
+    check (await pool.pgQuery("SELECT 1")).isErr()
+
+    (await pool.close()).expect("pool close")
+
+    var stillAlive = -1
+    for _ in 0 ..< 50:
+      stillAlive = (await observer.backendPidsAlive(survivorPids)).expect("alive count")
+      if stillAlive == 0:
+        break
+      await sleepAsync(100.milliseconds)
+
+    check stillAlive == 0
+
+    (await observer.close()).expect("observer close")
+
+  asyncTest "A pool holding a never used connection can still be closed":
+    ## A failing prepare leaves behind a connection that never ran a query, so
+    ## it has no futBecomeFree for the close barrier to wait on.
+    let pool = PgAsyncPool.new(storeMessageDbUrl, 1).expect("pool")
+
+    check (
+      await pool.runStmt(
+        "stmtOverMissingTable",
+        "SELECT * FROM a_table_that_does_not_exist",
+        newSeq[string](0),
+        newSeq[int32](0),
+        newSeq[int32](0),
+      )
+    ).isErr()
+
+    (await pool.close()).expect("pool close")
+
+  asyncTest "The pool comes back after its connection died":
+    let pool = PgAsyncPool.new(storeMessageDbUrl, 1).expect("pool")
+
+    check (await pool.pgQuery("SELECT pg_terminate_backend(pg_backend_pid())")).isErr()
+
+    (await pool.resetConnPool()).expect("resetConnPool")
+    (await pool.pgQuery("SELECT 1")).expect("query after reset")
+
+    (await pool.close()).expect("pool close")
+
+  asyncTest "A pool that cannot be closed does not bring the node down":
+    let pool = PgAsyncPool.new(storeMessageDbUrl, 1).expect("pool")
+
+    check (await pool.pgQuery("SELECT pg_terminate_backend(pg_backend_pid())")).isErr()
+
+    var fatalErrors = newSeq[string](0)
+    proc onFatalError(errMsg: string) {.gcsafe, closure, raises: [].} =
+      fatalErrors.add(errMsg)
+
+    let healthFut = checkConnectivity(pool, onFatalError)
+    await sleepAsync(1.seconds)
+
+    check fatalErrors.len == 0
+
+    await healthFut.cancelAndWait()
+    (await pool.close()).expect("pool close")