From 5dbb8a230539223a963c9e58fd318105526fd8cf Mon Sep 17 00:00:00 2001
From: fenos <fabri.feno@gmail.com>
Date: Wed, 30 Aug 2023 11:23:56 +0100
Subject: [PATCH] fix: handle database overload gracefully

---
 src/database/connection.ts | 87 +++++++++++++++++++++++++-------------
 src/http/error-handler.ts  |  7 ++-
 src/storage/errors.ts      | 14 ++++++
 3 files changed, 78 insertions(+), 30 deletions(-)

diff --git a/src/database/connection.ts b/src/database/connection.ts
index 57cf1d22..fd606ee9 100644
--- a/src/database/connection.ts
+++ b/src/database/connection.ts
@@ -5,6 +5,7 @@ import retry from 'async-retry'
 import { getConfig } from '../config'
 import { DbActiveConnection, DbActivePool } from '../monitoring/metrics'
 import { StorageBackendError } from '../storage'
+import KnexTimeoutError = knex.KnexTimeoutError
 
 // https://github.com/knex/knex/issues/387#issuecomment-51554522
 pg.types.setTypeParser(20, 'text', parseInt)
@@ -131,41 +132,69 @@ export class TenantConnection {
   }
 
   async transaction(instance?: Knex) {
-    const tnx = await retry(
-      async (bail) => {
-        try {
-          const pool = instance || this.pool
-          return await pool.transaction()
-        } catch (e) {
-          if (
-            e instanceof DatabaseError &&
-            e.code === '08P01' &&
-            e.message.includes('no more connections allowed')
-          ) {
-            throw e
+    try {
+      const tnx = await retry(
+        async (bail) => {
+          try {
+            const pool = instance || this.pool
+            return await pool.transaction()
+          } catch (e) {
+            if (
+              e instanceof DatabaseError &&
+              e.code === '08P01' &&
+              e.message.includes('no more connections allowed')
+            ) {
+              throw e
+            }
+
+            bail(e as Error)
+            return
           }
-
-          bail(e as Error)
-          return
+        },
+        {
+          minTimeout: 50,
+          maxTimeout: 200,
+          maxRetryTime: 3000,
+          retries: 10,
         }
-      },
-      {
-        minTimeout: 50,
-        maxTimeout: 200,
-        maxRetryTime: 3000,
-        retries: 10,
+      )
+
+      if (!tnx) {
+        throw new StorageBackendError('Could not create transaction', 500, 'transaction_failed')
       }
-    )
 
-    if (!tnx) {
-      throw new StorageBackendError('Could not create transaction', 500, 'transaction_failed')
-    }
+      if (!instance && this.options.isExternalPool) {
+        await tnx.raw(`SELECT set_config('search_path', ?, true)`, [searchPath.join(', ')])
+      }
 
-    if (!instance && this.options.isExternalPool) {
-      await tnx.raw(`SELECT set_config('search_path', ?, true)`, [searchPath.join(', ')])
-    }
+      return tnx
+    } catch (e) {
+      if (e instanceof KnexTimeoutError) {
+        throw StorageBackendError.withStatusCode(
+          'database_timeout',
+          544,
+          'The connection to the database timed out',
+          e
+        )
+      }
 
-    return tnx
+      if (
+        e instanceof DatabaseError &&
+        [
+          'remaining connection slots are reserved for non-replication superuser connections',
+          'no more connections allowed',
+        ].some((msg) => (e as DatabaseError).message.includes(msg))
+      ) {
+        throw StorageBackendError.withStatusCode(
+          'too_many_connections',
+          429,
+          'Too many connections issued to the database',
+          e
+        )
+      }
+
+      throw e
+    }
   }
 
   transactionProvider(instance?: Knex): Knex.TransactionProvider {
diff --git a/src/http/error-handler.ts b/src/http/error-handler.ts
index df43e4a9..21f85c52 100644
--- a/src/http/error-handler.ts
+++ b/src/http/error-handler.ts
@@ -20,7 +20,12 @@ export const setErrorHandler = (app: FastifyInstance) => {
 
     if (isRenderableError(error)) {
       const renderableError = error.render()
-      return reply.status(renderableError.statusCode === '500' ? 500 : 400).send(renderableError)
+      const statusCode = error.userStatusCode
+        ? error.userStatusCode
+        : renderableError.error === '500'
+        ? 500
+        : 400
+      return reply.status(statusCode).send(renderableError)
     }
 
     // Fastify errors
diff --git a/src/storage/errors.ts b/src/storage/errors.ts
index cff5a8b0..44cc3555 100644
--- a/src/storage/errors.ts
+++ b/src/storage/errors.ts
@@ -12,6 +12,7 @@ export type StorageError = {
  *  that we want to display to our users
  */
 export interface RenderableError {
+  userStatusCode?: number
   render(): StorageError
   getOriginalError(): unknown
 }
@@ -38,16 +39,29 @@ export function isS3Error(error: unknown): error is S3ServiceException {
 export class StorageBackendError extends Error implements RenderableError {
   httpStatusCode: number
   originalError: unknown
+  userStatusCode: number
 
   constructor(name: string, httpStatusCode: number, message: string, originalError?: unknown) {
     super(message)
     this.name = name
     this.httpStatusCode = httpStatusCode
+    this.userStatusCode = httpStatusCode === 500 ? 500 : 400
     this.message = message
     this.originalError = originalError
     Object.setPrototypeOf(this, StorageBackendError.prototype)
   }
 
+  static withStatusCode(
+    name: string,
+    statusCode: number,
+    message: string,
+    originalError?: unknown
+  ) {
+    const error = new StorageBackendError(name, statusCode, message, originalError)
+    error.userStatusCode = statusCode
+    return error
+  }
+
   static fromError(error?: unknown) {
     let name: string
     let httpStatusCode: number