fix(Postgres Node): Re-use connection pool across executions (#12346)

Co-authored-by: कारतोफ्फेलस्क्रिप्ट™ <aditya@netroy.in>

packages/nodes-base/utils/__tests__/connection-pool-manager.test.ts

@@ -0,0 +1,179 @@
+import { ConnectionPoolManager } from '@utils/connection-pool-manager';
+
+const ttl = 5 * 60 * 1000;
+const cleanUpInterval = 60 * 1000;
+
+let cpm: ConnectionPoolManager;
+
+beforeAll(() => {
+	jest.useFakeTimers();
+	cpm = ConnectionPoolManager.getInstance();
+});
+
+beforeEach(async () => {
+	await cpm.purgeConnections();
+});
+
+afterAll(() => {
+	cpm.onShutdown();
+});
+
+test('getInstance returns a singleton', () => {
+	const instance1 = ConnectionPoolManager.getInstance();
+	const instance2 = ConnectionPoolManager.getInstance();
+
+	expect(instance1).toBe(instance2);
+});
+
+describe('getConnection', () => {
+	test('calls fallBackHandler only once and returns the first value', async () => {
+		// ARRANGE
+		const connectionType = {};
+		const fallBackHandler = jest.fn().mockResolvedValue(connectionType);
+		const cleanUpHandler = jest.fn();
+		const options = {
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '1',
+			fallBackHandler,
+			cleanUpHandler,
+		};
+
+		// ACT 1
+		const connection = await cpm.getConnection<string>(options);
+
+		// ASSERT 1
+		expect(fallBackHandler).toHaveBeenCalledTimes(1);
+		expect(connection).toBe(connectionType);
+
+		// ACT 2
+		const connection2 = await cpm.getConnection<string>(options);
+		// ASSERT 2
+		expect(fallBackHandler).toHaveBeenCalledTimes(1);
+		expect(connection2).toBe(connectionType);
+	});
+
+	test('creates different pools for different node versions', async () => {
+		// ARRANGE
+		const connectionType1 = {};
+		const fallBackHandler1 = jest.fn().mockResolvedValue(connectionType1);
+		const cleanUpHandler1 = jest.fn();
+
+		const connectionType2 = {};
+		const fallBackHandler2 = jest.fn().mockResolvedValue(connectionType2);
+		const cleanUpHandler2 = jest.fn();
+
+		// ACT 1
+		const connection1 = await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '1',
+			fallBackHandler: fallBackHandler1,
+			cleanUpHandler: cleanUpHandler1,
+		});
+		const connection2 = await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '2',
+			fallBackHandler: fallBackHandler2,
+			cleanUpHandler: cleanUpHandler2,
+		});
+
+		// ASSERT
+		expect(fallBackHandler1).toHaveBeenCalledTimes(1);
+		expect(connection1).toBe(connectionType1);
+
+		expect(fallBackHandler2).toHaveBeenCalledTimes(1);
+		expect(connection2).toBe(connectionType2);
+
+		expect(connection1).not.toBe(connection2);
+	});
+
+	test('calls cleanUpHandler after TTL expires', async () => {
+		// ARRANGE
+		const connectionType = {};
+		const fallBackHandler = jest.fn().mockResolvedValue(connectionType);
+		const cleanUpHandler = jest.fn();
+		await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '1',
+			fallBackHandler,
+			cleanUpHandler,
+		});
+
+		// ACT
+		jest.advanceTimersByTime(ttl + cleanUpInterval * 2);
+
+		// ASSERT
+		expect(cleanUpHandler).toHaveBeenCalledTimes(1);
+	});
+});
+
+describe('onShutdown', () => {
+	test('calls all clean up handlers', async () => {
+		// ARRANGE
+		const connectionType1 = {};
+		const fallBackHandler1 = jest.fn().mockResolvedValue(connectionType1);
+		const cleanUpHandler1 = jest.fn();
+		await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '1',
+			fallBackHandler: fallBackHandler1,
+			cleanUpHandler: cleanUpHandler1,
+		});
+
+		const connectionType2 = {};
+		const fallBackHandler2 = jest.fn().mockResolvedValue(connectionType2);
+		const cleanUpHandler2 = jest.fn();
+		await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '2',
+			fallBackHandler: fallBackHandler2,
+			cleanUpHandler: cleanUpHandler2,
+		});
+
+		// ACT 1
+		cpm.onShutdown();
+
+		// ASSERT
+		expect(cleanUpHandler1).toHaveBeenCalledTimes(1);
+		expect(cleanUpHandler2).toHaveBeenCalledTimes(1);
+	});
+
+	test('calls all clean up handlers when `exit` is emitted on process', async () => {
+		// ARRANGE
+		const connectionType1 = {};
+		const fallBackHandler1 = jest.fn().mockResolvedValue(connectionType1);
+		const cleanUpHandler1 = jest.fn();
+		await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '1',
+			fallBackHandler: fallBackHandler1,
+			cleanUpHandler: cleanUpHandler1,
+		});
+
+		const connectionType2 = {};
+		const fallBackHandler2 = jest.fn().mockResolvedValue(connectionType2);
+		const cleanUpHandler2 = jest.fn();
+		await cpm.getConnection<string>({
+			credentials: {},
+			nodeType: 'example',
+			nodeVersion: '2',
+			fallBackHandler: fallBackHandler2,
+			cleanUpHandler: cleanUpHandler2,
+		});
+
+		// ACT 1
+		// @ts-expect-error we're not supposed to emit `exit` so it's missing from
+		// the type definition
+		process.emit('exit');
+
+		// ASSERT
+		expect(cleanUpHandler1).toHaveBeenCalledTimes(1);
+		expect(cleanUpHandler2).toHaveBeenCalledTimes(1);
+	});
+});

packages/nodes-base/utils/connection-pool-manager.ts

@@ -0,0 +1,137 @@
+import { createHash } from 'crypto';
+
+let instance: ConnectionPoolManager;
+
+// 5 minutes
+const ttl = 5 * 60 * 1000;
+
+// 1 minute
+const cleanUpInterval = 60 * 1000;
+
+type RegistrationOptions = {
+	credentials: unknown;
+	nodeType: string;
+	nodeVersion?: string;
+};
+
+type GetConnectionOption<Pool> = RegistrationOptions & {
+	/** When a node requests for a connection pool, but none is available, this handler is called to create new instance of the pool, which then cached and re-used until it goes stale.  */
+	fallBackHandler: () => Promise<Pool>;
+
+	/** When a pool hasn't been used in a while, or when the server is shutting down, this handler is invoked to close the pool */
+	cleanUpHandler: (pool: Pool) => Promise<void>;
+};
+
+type Registration<Pool> = {
+	/** This is an instance of a Connection Pool class, that gets reused across multiple executions */
+	pool: Pool;
+
+	/** @see GetConnectionOption['closeHandler'] */
+	cleanUpHandler: (pool: Pool) => Promise<void>;
+
+	/** We keep this timestamp to check if a pool hasn't been used in a while, and if it needs to be closed */
+	lastUsed: number;
+};
+
+export class ConnectionPoolManager {
+	/**
+	 * Gets the singleton instance of the ConnectionPoolManager.
+	 * Creates a new instance if one doesn't exist.
+	 */
+	static getInstance(): ConnectionPoolManager {
+		if (!instance) {
+			instance = new ConnectionPoolManager();
+		}
+		return instance;
+	}
+
+	private map = new Map<string, Registration<unknown>>();
+
+	/**
+	 * Private constructor that initializes the connection pool manager.
+	 * Sets up cleanup handlers for process exit and stale connections.
+	 */
+	private constructor() {
+		// Close all open pools when the process exits
+		process.on('exit', () => this.onShutdown());
+
+		// Regularly close stale pools
+		setInterval(() => this.cleanupStaleConnections(), cleanUpInterval);
+	}
+
+	/**
+	 * Generates a unique key for connection pool identification.
+	 * Hashes the credentials and node information for security.
+	 */
+	private makeKey({ credentials, nodeType, nodeVersion }: RegistrationOptions): string {
+		// The credential contains decrypted secrets, that's why we hash it.
+		return createHash('sha1')
+			.update(
+				JSON.stringify({
+					credentials,
+					nodeType,
+					nodeVersion,
+				}),
+			)
+			.digest('base64');
+	}
+
+	/**
+	 * Gets or creates a connection pool for the given options.
+	 * Updates the last used timestamp for existing connections.
+	 */
+	async getConnection<T>(options: GetConnectionOption<T>): Promise<T> {
+		const key = this.makeKey(options);
+
+		let value = this.map.get(key);
+		if (!value) {
+			value = {
+				pool: await options.fallBackHandler(),
+				cleanUpHandler: options.cleanUpHandler,
+			} as Registration<unknown>;
+		}
+
+		this.map.set(key, { ...value, lastUsed: Date.now() });
+		return value.pool as T;
+	}
+
+	/**
+	 * Removes and cleans up connection pools that haven't been used within the
+	 * TTL.
+	 */
+	private cleanupStaleConnections() {
+		const now = Date.now();
+		for (const [key, { cleanUpHandler, lastUsed, pool }] of this.map.entries()) {
+			if (now - lastUsed > ttl) {
+				void cleanUpHandler(pool);
+				this.map.delete(key);
+			}
+		}
+	}
+
+	/**
+	 * Removes and cleans up all existing connection pools.
+	 */
+	async purgeConnections(): Promise<void> {
+		await Promise.all(
+			[...this.map.entries()].map(async ([key, value]) => {
+				this.map.delete(key);
+
+				return await value.cleanUpHandler(value.pool);
+			}),
+		);
+	}
+
+	/**
+	 * Cleans up all connection pools when the process is shutting down.
+	 * Does not wait for cleanup promises to resolve also does not remove the
+	 * references from the pool.
+	 *
+	 * Only call this on process shutdown.
+	 */
+	onShutdown() {
+		for (const { cleanUpHandler, pool } of this.map.values()) {
+			void cleanUpHandler(pool);
+		}
+	}
+}