pg-cache
Version: 
PostgreSQL connection pool LRU cache manager
151 lines (150 loc) • 4.62 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.teardownPgPools = exports.close = exports.pgCache = exports.PgPoolCacheManager = void 0;
const logger_1 = require("@launchql/logger");
const lru_cache_1 = require("lru-cache");
const log = new logger_1.Logger('pg-cache');
const ONE_HOUR_IN_MS = 1000 * 60 * 60;
const ONE_DAY = ONE_HOUR_IN_MS * 24;
const ONE_YEAR = ONE_DAY * 366;
// Kubernetes sends only SIGTERM on pod shutdown
const SYS_EVENTS = ['SIGTERM'];
class ManagedPgPool {
    pool;
    key;
    isDisposed = false;
    disposePromise = null;
    constructor(pool, key) {
        this.pool = pool;
        this.key = key;
    }
    async dispose() {
        if (this.isDisposed)
            return this.disposePromise;
        this.isDisposed = true;
        this.disposePromise = (async () => {
            try {
                if (!this.pool.ended) {
                    await this.pool.end();
                    log.success(`pg.Pool ${this.key} ended.`);
                }
                else {
                    log.info(`pg.Pool ${this.key} already ended.`);
                }
            }
            catch (err) {
                log.error(`Error ending pg.Pool ${this.key}: ${err.message}`);
                throw err;
            }
        })();
        return this.disposePromise;
    }
}
class PgPoolCacheManager {
    cleanupTasks = [];
    closed = false;
    cleanupCallbacks = new Set();
    pgCache = new lru_cache_1.LRUCache({
        max: 10,
        ttl: ONE_YEAR,
        updateAgeOnGet: true,
        dispose: (managedPool, key, reason) => {
            log.debug(`Disposing pg pool [${key}] (${reason})`);
            this.notifyCleanup(key);
            this.disposePool(managedPool);
        }
    });
    // Register a cleanup callback to be called when pools are disposed
    registerCleanupCallback(callback) {
        this.cleanupCallbacks.add(callback);
        // Return unregister function
        return () => {
            this.cleanupCallbacks.delete(callback);
        };
    }
    get(key) {
        return this.pgCache.get(key)?.pool;
    }
    has(key) {
        return this.pgCache.has(key);
    }
    set(key, pool) {
        if (this.closed)
            throw new Error(`Cannot add to cache after it has been closed (key: ${key})`);
        this.pgCache.set(key, new ManagedPgPool(pool, key));
    }
    delete(key) {
        const managedPool = this.pgCache.get(key);
        const existed = this.pgCache.delete(key);
        if (!existed && managedPool) {
            this.notifyCleanup(key);
            this.disposePool(managedPool);
        }
    }
    clear() {
        const entries = [...this.pgCache.entries()];
        this.pgCache.clear();
        for (const [key, managedPool] of entries) {
            this.notifyCleanup(key);
            this.disposePool(managedPool);
        }
    }
    async close() {
        if (this.closed)
            return;
        this.closed = true;
        this.clear();
        await this.waitForDisposals();
    }
    async waitForDisposals() {
        if (this.cleanupTasks.length === 0)
            return;
        const tasks = [...this.cleanupTasks];
        this.cleanupTasks = [];
        await Promise.allSettled(tasks);
    }
    notifyCleanup(pgPoolKey) {
        this.cleanupCallbacks.forEach(callback => {
            try {
                callback(pgPoolKey);
            }
            catch (err) {
                log.error(`Error in cleanup callback for pool ${pgPoolKey}: ${err.message}`);
            }
        });
    }
    disposePool(managedPool) {
        if (managedPool.isDisposed)
            return;
        const task = managedPool.dispose();
        this.cleanupTasks.push(task);
    }
}
exports.PgPoolCacheManager = PgPoolCacheManager;
// Create the singleton instance
exports.pgCache = new PgPoolCacheManager();
// --- Graceful Shutdown ---
const closePromise = { promise: null };
const close = async (verbose = false) => {
    if (closePromise.promise)
        return closePromise.promise;
    closePromise.promise = (async () => {
        if (verbose)
            log.info('Closing pg cache...');
        await exports.pgCache.close();
        if (verbose)
            log.success('PG cache disposed.');
    })();
    return closePromise.promise;
};
exports.close = close;
SYS_EVENTS.forEach(event => {
    process.on(event, () => {
        log.info(`Received ${event}`);
        (0, exports.close)();
    });
});
const teardownPgPools = async (verbose = false) => {
    return (0, exports.close)(verbose);
};
exports.teardownPgPools = teardownPgPools;