fix: support PG reconnects (#9663)

Signed-off-by: Matt Krick <[email protected]>
ParabolInc · Apr 23, 2024 · 32574a6 · 32574a6
1 parent 01372bc
commit 32574a6
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 21 deletions.
diff --git a/packages/embedder/EmbeddingsJobQueueStream.ts b/packages/embedder/EmbeddingsJobQueueStream.ts
@@ -41,16 +41,20 @@ export class EmbeddingsJobQueueStream implements AsyncIterableIterator<DBJob> {
         .returningAll()
         .executeTakeFirst()
     }
-    const job = (await getJob(false)) || (await getJob(true))
-    if (!job) {
-      Logger.log('JobQueueStream: no jobs found')
-      // queue is empty, so sleep for a while
-      await sleep(ms('10s'))
+    try {
+      const job = (await getJob(false)) || (await getJob(true))
+      if (!job) {
+        Logger.log('JobQueueStream: no jobs found')
+        // queue is empty, so sleep for a while
+        await sleep(ms('10s'))
+        return this.next()
+      }
+      await this.orchestrator.runStep(job)
+      return {done: false, value: job}
+    } catch (e) {
+      await sleep(1000)
       return this.next()
     }
-
-    await this.orchestrator.runStep(job)
-    return {done: false, value: job}
   }
   return() {
     return Promise.resolve({done: true as const, value: undefined})

diff --git a/packages/server/postgres/getKysely.ts b/packages/server/postgres/getKysely.ts
@@ -3,21 +3,26 @@ import getPg from './getPg'
 import {DB} from './pg.d'
 
 let kysely: Kysely<DB> | undefined
+
+const makeKysely = () => {
+  const nextPg = getPg()
+  nextPg.on('poolChange' as any, makeKysely)
+  return new Kysely<DB>({
+    dialect: new PostgresDialect({
+      pool: nextPg
+    })
+    // ,log(event) {
+    //   if (event.level === 'query') {
+    //     console.log(event.query.sql)
+    //     console.log(event.query.parameters)
+    //   }
+    // }
+  })
+}
+
 const getKysely = () => {
   if (!kysely) {
-    const pg = getPg()
-    kysely = new Kysely<DB>({
-      dialect: new PostgresDialect({
-        pool: pg
-      })
-      // query logging, if you'd like it:
-      // log(event) {
-      //   if (event.level === 'query') {
-      //     console.log(event.query.sql)
-      //     console.log(event.query.parameters)
-      //   }
-      // }
-    })
+    kysely = makeKysely()
   }
   return kysely
 }

diff --git a/packages/server/postgres/getPg.ts b/packages/server/postgres/getPg.ts
@@ -1,12 +1,31 @@
 import {Pool} from 'pg'
+import sleep from '../../client/utils/sleep'
 import getPgConfig from './getPgConfig'
 
 const config = getPgConfig()
 
+const graceFullyReconnect = async () => {
+  for (let i = 0; i < 1e6; i++) {
+    const nextPool = new Pool(getPgConfig())
+    try {
+      const testClient = await nextPool.connect()
+      testClient.release()
+      nextPool.on('error', graceFullyReconnect)
+      const oldPool = pool
+      pool = nextPool
+      oldPool?.emit('changePool')
+      return
+    } catch (e) {
+      await sleep(1000)
+    }
+  }
+}
+
 let pool: Pool | undefined
 const getPg = () => {
   if (!pool) {
     pool = new Pool(config)
+    pool.on('error', graceFullyReconnect)
   }
   return pool
 }