feat: rebalancing connection pooling

Author: fenos

migrations/multitenant/0017-pool-mode.sql

@@ -0,0 +1 @@
+ALTER TABLE tenants ADD COLUMN IF NOT EXISTS database_pool_mode TEXT NULL;

package-lock.json

@@ -28,6 +28,7 @@
     "node": ">= 14.0.0"
   },
   "dependencies": {
+    "@aws-sdk/client-ecs": "^3.795.0",
     "@aws-sdk/client-s3": "3.654.0",
     "@aws-sdk/lib-storage": "3.654.0",
     "@aws-sdk/s3-request-presigner": "3.654.0",

package.json

@@ -86,6 +86,7 @@ type StorageConfigType = {
   databaseURL: string
   databaseSSLRootCert?: string
   databasePoolURL?: string
+  databasePoolMode?: 'single_use' | 'recycle'
   databaseMaxConnections: number
   databaseFreePoolAfterInactivity: number
   databaseConnectionTimeout: number
@@ -357,6 +358,7 @@ export function getConfig(options?: { reload?: boolean }): StorageConfigType {
     databaseSSLRootCert: getOptionalConfigFromEnv('DATABASE_SSL_ROOT_CERT'),
     databaseURL: getOptionalIfMultitenantConfigFromEnv('DATABASE_URL') || '',
     databasePoolURL: getOptionalConfigFromEnv('DATABASE_POOL_URL') || '',
+    databasePoolMode: getOptionalConfigFromEnv('DATABASE_POOL_MODE'),
     databaseMaxConnections: parseInt(
       getOptionalConfigFromEnv('DATABASE_MAX_CONNECTIONS') || '20',
       10

src/config.ts

@@ -26,6 +26,7 @@ const patchSchema = {
       anonKey: { type: 'string' },
       databaseUrl: { type: 'string' },
       databasePoolUrl: { type: 'string', nullable: true },
+      databasePoolMode: { type: 'string', nullable: true },
       maxConnections: { type: 'number' },
       jwks: { type: 'object', nullable: true },
       fileSizeLimit: { type: 'number' },
@@ -112,6 +113,7 @@ export default async function routes(fastify: FastifyInstance) {
         anon_key,
         database_url,
         database_pool_url,
+        database_pool_mode,
         max_connections,
         file_size_limit,
         jwt_secret,
@@ -130,6 +132,7 @@ export default async function routes(fastify: FastifyInstance) {
         anonKey: decrypt(anon_key),
         databaseUrl: decrypt(database_url),
         databasePoolUrl: database_pool_url ? decrypt(database_pool_url) : undefined,
+        databasePoolMode: database_pool_mode,
         maxConnections: max_connections ? Number(max_connections) : undefined,
         fileSizeLimit: Number(file_size_limit),
         jwtSecret: decrypt(jwt_secret),
@@ -164,6 +167,7 @@ export default async function routes(fastify: FastifyInstance) {
       anon_key,
       database_url,
       database_pool_url,
+      database_pool_mode,
       max_connections,
       file_size_limit,
       jwt_secret,
@@ -188,6 +192,7 @@ export default async function routes(fastify: FastifyInstance) {
           : database_pool_url
           ? decrypt(database_pool_url)
           : undefined,
+      databasePoolMode: database_pool_mode,
       maxConnections: max_connections ? Number(max_connections) : undefined,
       fileSizeLimit: Number(file_size_limit),
       jwtSecret: decrypt(jwt_secret),
@@ -217,6 +222,7 @@ export default async function routes(fastify: FastifyInstance) {
     const {
       anonKey,
       databaseUrl,
+      databasePoolMode,
       fileSizeLimit,
       jwtSecret,
       jwks,
@@ -233,6 +239,7 @@ export default async function routes(fastify: FastifyInstance) {
         anon_key: encrypt(anonKey),
         database_url: encrypt(databaseUrl),
         database_pool_url: databasePoolUrl ? encrypt(databasePoolUrl) : undefined,
+        database_pool_mode: databasePoolMode,
         max_connections: maxConnections ? Number(maxConnections) : undefined,
         file_size_limit: fileSizeLimit,
         jwt_secret: encrypt(jwtSecret),
@@ -280,6 +287,7 @@ export default async function routes(fastify: FastifyInstance) {
         serviceKey,
         features,
         databasePoolUrl,
+        databasePoolMode,
         maxConnections,
         tracingMode,
         disableEvents,
@@ -295,6 +303,7 @@ export default async function routes(fastify: FastifyInstance) {
             : databasePoolUrl === null
             ? null
             : undefined,
+          database_pool_mode: databasePoolMode,
           max_connections: maxConnections ? Number(maxConnections) : undefined,
           file_size_limit: fileSizeLimit,
           jwt_secret: jwtSecret !== undefined ? encrypt(jwtSecret) : undefined,

src/http/routes/admin/tenants.ts

@@ -0,0 +1,53 @@
+import { ClusterDiscoveryECS } from '@internal/cluster/ecs'
+
+import { EventEmitter } from 'node:events'
+import { logger } from '@internal/monitoring'
+
+const clusterEvent = new EventEmitter()
+
+export class Cluster {
+  static size: number = 0
+  protected static watcher?: NodeJS.Timeout = undefined
+
+  static on(event: string, listener: (...args: any[]) => void) {
+    clusterEvent.on(event, listener)
+  }
+
+  static async init(abortSignal: AbortSignal) {
+    if (process.env.CLUSTER_DISCOVERY === 'ECS') {
+      const cluster = new ClusterDiscoveryECS()
+      Cluster.size = await cluster.getClusterSize()
+
+      logger.info('[Cluster] Initial cluster size', {
+        type: 'cluster',
+        size: Cluster.size,
+      })
+
+      Cluster.watcher = setInterval(() => {
+        cluster
+          .getClusterSize()
+          .then((size) => {
+            if (size === Cluster.size) {
+              clusterEvent.emit('change', { size })
+            }
+            Cluster.size = size
+          })
+          .catch((e) => {
+            console.error('Error getting cluster size', e)
+          })
+      }, 20 * 1000)
+
+      abortSignal.addEventListener(
+        'abort',
+        () => {
+          if (Cluster.watcher) {
+            clearInterval(Cluster.watcher)
+            clusterEvent.removeAllListeners()
+            Cluster.watcher = undefined
+          }
+        },
+        { once: true }
+      )
+    }
+  }
+}

src/internal/cluster/discovery.ts

@@ -0,0 +1,26 @@
+import { ECSClient, ListTasksCommand } from '@aws-sdk/client-ecs'
+import axios from 'axios'
+
+export class ClusterDiscoveryECS {
+  private client: ECSClient
+
+  constructor() {
+    this.client = new ECSClient()
+  }
+
+  async getClusterSize() {
+    if (!process.env.ECS_CONTAINER_METADATA_URI) {
+      throw new Error('ECS_CONTAINER_METADATA_URI is not set')
+    }
+
+    const respMetadata = await axios.get(`${process.env.ECS_CONTAINER_METADATA_URI}/task`)
+
+    const command = new ListTasksCommand({
+      serviceName: respMetadata.data.ServiceName,
+      cluster: respMetadata.data.Cluster,
+      desiredStatus: 'RUNNING',
+    })
+    const response = await this.client.send(command)
+    return response.taskArns?.length || 0
+  }
+}

src/internal/cluster/ecs.ts

@@ -1,3 +1,4 @@
 export * from './mutex'
+export * from './wait'
 export * from './async-abort-controller'
 export * from './merge-async-itertor'

src/internal/concurrency/index.ts

@@ -0,0 +1,5 @@
+export function wait(ms: number) {
+  return new Promise((resolve) => {
+    setTimeout(resolve, ms)
+  })
+}

src/internal/concurrency/wait.ts

@@ -1,7 +1,9 @@
 import { getConfig } from '../../config'
 import { getTenantConfig } from './tenant'
-import { User, TenantConnection } from './connection'
+import { TenantConnection } from './connection'
+import { User } from './pool'
 import { ERRORS } from '@internal/errors'
+import { Cluster } from '@internal/cluster/discovery'
 
 interface ConnectionOptions {
   host: string
@@ -21,17 +23,18 @@ interface ConnectionOptions {
  * @param options
  */
 export async function getPostgresConnection(options: ConnectionOptions): Promise<TenantConnection> {
-  const dbCredentials = await getDbCredentials(options.tenantId, options.host, {
+  const dbCredentials = await getDbSettings(options.tenantId, options.host, {
     disableHostCheck: options.disableHostCheck,
   })
 
   return await TenantConnection.create({
     ...dbCredentials,
     ...options,
+    clusterSize: Cluster.size,
   })
 }
 
-async function getDbCredentials(
+async function getDbSettings(
   tenantId: string,
   host: string | undefined,
   options?: { disableHostCheck?: boolean }
@@ -42,11 +45,13 @@ async function getDbCredentials(
     databaseURL,
     databaseMaxConnections,
     requestXForwardedHostRegExp,
+    databasePoolMode,
   } = getConfig()
 
   let dbUrl = databasePoolURL || databaseURL
   let maxConnections = databaseMaxConnections
   let isExternalPool = Boolean(databasePoolURL)
+  let isSingleUse = !databasePoolMode || databasePoolMode === 'single_use'
 
   if (isMultitenant) {
     if (!tenantId) {
@@ -70,11 +75,13 @@ async function getDbCredentials(
     dbUrl = tenant.databasePoolUrl || tenant.databaseUrl
     isExternalPool = Boolean(tenant.databasePoolUrl)
     maxConnections = tenant.maxConnections ?? maxConnections
+    isSingleUse = tenant.databasePoolMode ? tenant.databasePoolMode !== 'recycled' : isSingleUse
   }
 
   return {
     dbUrl,
     isExternalPool,
     maxConnections,
+    isSingleUse,
   }
 }