From 5a0edff79a7db4d4d7247b863ee4e138929706cf Mon Sep 17 00:00:00 2001
From: Ivan FB <128452529+Ivansete-status@users.noreply.github.com>
Date: Tue, 10 Sep 2024 15:07:53 +0200
Subject: [PATCH] fix: get back health check for postgres legacy (#3010)

---
 .../postgres_driver/postgres_driver.nim       |  7 ++++
 .../postgres_driver/postgres_healthcheck.nim  | 38 +++++++++++++++++++
 2 files changed, 45 insertions(+)
 create mode 100644 waku/waku_archive_legacy/driver/postgres_driver/postgres_healthcheck.nim

diff --git a/waku/waku_archive_legacy/driver/postgres_driver/postgres_driver.nim b/waku/waku_archive_legacy/driver/postgres_driver/postgres_driver.nim
index c2f83e36a3..4f3532622d 100644
--- a/waku/waku_archive_legacy/driver/postgres_driver/postgres_driver.nim
+++ b/waku/waku_archive_legacy/driver/postgres_driver/postgres_driver.nim
@@ -15,6 +15,7 @@ import
   ../../../waku_core,
   ../../common,
   ../../driver,
+  ./postgres_healthcheck,
   ../../../common/databases/db_postgres as waku_postgres
 
 type PostgresDriver* = ref object of ArchiveDriver
@@ -132,6 +133,12 @@ proc new*(
   let writeConnPool = PgAsyncPool.new(dbUrl, maxNumConnOnEachPool).valueOr:
     return err("error creating write conn pool PgAsyncPool")
 
+  if not isNil(onFatalErrorAction):
+    asyncSpawn checkConnectivity(readConnPool, onFatalErrorAction)
+
+  if not isNil(onFatalErrorAction):
+    asyncSpawn checkConnectivity(writeConnPool, onFatalErrorAction)
+
   let driver = PostgresDriver(writeConnPool: writeConnPool, readConnPool: readConnPool)
   return ok(driver)
 
diff --git a/waku/waku_archive_legacy/driver/postgres_driver/postgres_healthcheck.nim b/waku/waku_archive_legacy/driver/postgres_driver/postgres_healthcheck.nim
new file mode 100644
index 0000000000..4c9f170c9f
--- /dev/null
+++ b/waku/waku_archive_legacy/driver/postgres_driver/postgres_healthcheck.nim
@@ -0,0 +1,38 @@
+{.push raises: [].}
+
+import chronos, chronicles, results
+import ../../../common/databases/db_postgres, ../../../common/error_handling
+
+## Simple query to validate that the postgres is working and attending requests
+const HealthCheckQuery = "SELECT version();"
+const CheckConnectivityInterval = 60.seconds
+const MaxNumTrials = 20
+const TrialInterval = 1.seconds
+
+proc checkConnectivity*(
+    connPool: PgAsyncPool, onFatalErrorAction: OnFatalErrorHandler
+) {.async.} =
+  while true:
+    (await connPool.pgQuery(HealthCheckQuery)).isOkOr:
+      ## The connection failed once. Let's try reconnecting for a while.
+      ## Notice that the 'pgQuery' proc tries to establish a new connection.
+
+      block errorBlock:
+        ## Force close all the opened connections. No need to close gracefully.
+        (await connPool.resetConnPool()).isOkOr:
+          onFatalErrorAction("checkConnectivity legacy resetConnPool error: " & error)
+
+        var numTrial = 0
+        while numTrial < MaxNumTrials:
+          let res = await connPool.pgQuery(HealthCheckQuery)
+          if res.isOk():
+            ## Connection resumed. Let's go back to the normal healthcheck.
+            break errorBlock
+
+          await sleepAsync(TrialInterval)
+          numTrial.inc()
+
+        ## The connection couldn't be resumed. Let's inform the upper layers.
+        onFatalErrorAction("postgres legacy health check error: " & error)
+
+    await sleepAsync(CheckConnectivityInterval)