Merge 887d154cfe into bbee498b3e

src/lib/__snapshots__/create-config.test.ts.snap

@@ -24,6 +24,7 @@ exports[`should create default config 1`] = `
     "type": "open-source",
   },
   "buildDate": undefined,
+  "checkDbOnReady": false,
   "clientFeatureCaching": {
     "enabled": true,
     "maxAge": 3600000,

src/lib/create-config.ts

@@ -786,6 +786,10 @@ export function createConfig(options: IUnleashOptions): IUnleashConfig {
         options.prometheusImpactMetricsApi ||
         process.env.PROMETHEUS_IMPACT_METRICS_API;
 
+    const checkDbOnReady =
+        Boolean(options.checkDbOnReady) ??
+        parseEnvVarBoolean(process.env.CHECK_DB_ON_READY, false);
+
     return {
         db,
         session,
@@ -830,5 +834,6 @@ export function createConfig(options: IUnleashOptions): IUnleashConfig {
         userInactivityThresholdInDays,
         buildDate: process.env.BUILD_DATE,
         unleashFrontendToken,
+        checkDbOnReady,
     };
 }

src/lib/openapi/spec/index.ts

@@ -169,6 +169,7 @@ export * from './public-signup-token-schema.js';
 export * from './public-signup-token-update-schema.js';
 export * from './public-signup-tokens-schema.js';
 export * from './push-variants-schema.js';
+export * from './ready-check-schema.js';
 export * from './record-ui-error-schema.js';
 export * from './release-plan-milestone-schema.js';
 export * from './release-plan-milestone-strategy-schema.js';

src/lib/openapi/spec/ready-check-schema.ts

@@ -0,0 +1,22 @@
+import type { FromSchema } from 'json-schema-to-ts';
+
+export const readyCheckSchema = {
+    $id: '#/components/schemas/readyCheckSchema',
+    type: 'object',
+    description:
+        'Used by service orchestrators to decide whether this Unleash instance should be considered ready to serve requests.',
+    additionalProperties: false,
+    required: ['health'],
+    properties: {
+        health: {
+            description:
+                'The readiness state this Unleash instance is in. GOOD if the server is up and running. If the server is unhealthy you will get an unsuccessful http response.',
+            type: 'string',
+            enum: ['GOOD'],
+            example: 'GOOD',
+        },
+    },
+    components: {},
+} as const;
+
+export type ReadyCheckSchema = FromSchema<typeof readyCheckSchema>;

src/lib/routes/health-check.test.ts

@@ -28,11 +28,11 @@ afterEach(() => {
     getLogger.setMuteError(false);
 });
 
-test('should give 200 when ready', async () => {
+test('should give 200 when healthy', async () => {
     await request.get('/health').expect(200);
 });
 
-test('should give health=GOOD  when ready', async () => {
+test('should give health=GOOD when healthy', async () => {
     expect.assertions(2);
     await request
         .get('/health')

src/lib/routes/index.ts

@@ -8,6 +8,7 @@ import Controller from './controller.js';
 import { AdminApi } from './admin-api/index.js';
 import ClientApi from './client-api/index.js';
 
+import { ReadyCheckController } from './ready-check.js';
 import { HealthCheckController } from './health-check.js';
 import FrontendAPIController from '../features/frontend-api/frontend-api-controller.js';
 import EdgeController from './edge-api/index.js';
@@ -25,6 +26,10 @@ class IndexRouter extends Controller {
     ) {
         super(config);
 
+        this.use(
+            '/ready',
+            new ReadyCheckController(config, services, db).router,
+        );
         this.use('/health', new HealthCheckController(config, services).router);
         this.use(
             '/invite',

src/lib/routes/ready-check.ts

@@ -0,0 +1,114 @@
+import type { Request, Response } from 'express';
+import type { PoolClient } from 'pg';
+import type { IUnleashConfig } from '../types/option.js';
+import type { IUnleashServices } from '../services/index.js';
+import type { Db } from '../db/db.js';
+import type { Logger } from '../logger.js';
+
+import Controller from './controller.js';
+import { NONE } from '../types/permissions.js';
+import { createResponseSchema } from '../openapi/util/create-response-schema.js';
+import type { ReadyCheckSchema } from '../openapi/spec/ready-check-schema.js';
+import { emptyResponse, parseEnvVarNumber } from '../server-impl.js';
+
+export class ReadyCheckController extends Controller {
+    private logger: Logger;
+
+    private db?: Db;
+
+    constructor(
+        config: IUnleashConfig,
+        { openApiService }: Pick<IUnleashServices, 'openApiService'>,
+        db?: Db,
+    ) {
+        super(config);
+        this.logger = config.getLogger('ready-check.js');
+        this.db = db;
+
+        this.route({
+            method: 'get',
+            path: '',
+            handler: this.getReady,
+            permission: NONE,
+            middleware: [
+                openApiService.validPath({
+                    tags: ['Operational'],
+                    operationId: 'getReady',
+                    summary: 'Get instance readiness status',
+                    description:
+                        'This operation returns information about whether this Unleash instance is ready to serve requests or not. Typically used by your deployment orchestrator (e.g. Kubernetes, Docker Swarm, Mesos, et al.).',
+                    responses: {
+                        200: createResponseSchema('readyCheckSchema'),
+                        503: emptyResponse,
+                    },
+                }),
+            ],
+        });
+    }
+
+    async getReady(_: Request, res: Response<ReadyCheckSchema>): Promise<void> {
+        if (this.config.checkDbOnReady && this.db) {
+            try {
+                const timeoutMs = parseEnvVarNumber(
+                    process.env.DATABASE_STATEMENT_TIMEOUT_MS,
+                    200,
+                );
+
+                await this.runReadinessQuery(timeoutMs);
+                res.status(200).json({ health: 'GOOD' });
+                return;
+            } catch (err: any) {
+                this.logger.warn('Database readiness check failed', err);
+                res.status(503).end();
+                return;
+            }
+        }
+
+        res.status(200).json({ health: 'GOOD' });
+    }
+
+    private async runReadinessQuery(timeoutMs: number): Promise<void> {
+        const client = getKnexClient(this.db!);
+        const pendingAcquire = client.pool.acquire();
+
+        const abortDelay = setTimeout(() => pendingAcquire.abort(), timeoutMs);
+        let connection: PoolClient | undefined;
+
+        try {
+            connection = (await pendingAcquire.promise) as PoolClient;
+        } finally {
+            clearTimeout(abortDelay);
+        }
+
+        try {
+            await connection.query('BEGIN');
+            await connection.query({
+                text: `SET LOCAL statement_timeout = ${timeoutMs}`,
+            });
+            await connection.query('SELECT 1');
+            await connection.query('COMMIT');
+        } catch (error) {
+            try {
+                await connection.query('ROLLBACK');
+            } catch (rollbackError) {
+                this.logger.debug(
+                    'Failed to rollback readiness timeout transaction',
+                    rollbackError,
+                );
+            }
+            throw error;
+        } finally {
+            if (connection) {
+                await client.releaseConnection(connection);
+            }
+        }
+    }
+}
+
+function getKnexClient(db: Db): Db['client'] {
+    if (db.client?.pool && typeof db.client.releaseConnection === 'function') {
+        return db.client as Db['client'];
+    }
+
+    throw new Error('Unsupported database handle for readiness check');
+}

src/lib/types/option.ts

@@ -175,6 +175,7 @@ export interface IUnleashOptions {
     resourceLimits?: Partial<ResourceLimits>;
     userInactivityThresholdInDays?: number;
     unleashFrontendToken?: string;
+    checkDbOnReady?: boolean;
 }
 
 export interface IEmailOption {
@@ -301,4 +302,5 @@ export interface IUnleashConfig {
     userInactivityThresholdInDays: number;
     buildDate?: string;
     unleashFrontendToken?: string;
+    checkDbOnReady?: boolean;
 }

src/test/e2e/ready.e2e.test.ts

@@ -0,0 +1,102 @@
+import { setupAppWithCustomConfig } from './helpers/test-helper.js';
+import dbInit, { type ITestDb } from './helpers/database-init.js';
+
+let db: ITestDb;
+
+describe('DB is up', () => {
+    beforeAll(async () => {
+        db = await dbInit();
+    });
+
+    test('when checkDb is disabled, returns ready', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            undefined,
+            db.rawDatabase,
+        );
+        await request
+            .get('/ready')
+            .expect('Content-Type', /json/)
+            .expect(200)
+            .expect('{"health":"GOOD"}');
+    });
+
+    test('when checkDb is enabled, returns ready', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            { checkDbOnReady: true },
+            db.rawDatabase,
+        );
+        await request.get('/ready').expect(200).expect('{"health":"GOOD"}');
+    });
+
+    test('fails fast when readiness query hangs', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            { checkDbOnReady: true },
+            db.rawDatabase,
+        );
+
+        const pool = db.rawDatabase.client.pool;
+        const originalAcquire = pool.acquire.bind(pool);
+
+        pool.acquire = () => {
+            const pending = originalAcquire();
+            pending.promise = pending.promise.then((conn: any) => {
+                const originalQuery = conn.query;
+                conn.query = (queryConfig: any, ...args: any[]) => {
+                    const isSelectOne =
+                        queryConfig?.toUpperCase() === 'SELECT 1';
+
+                    if (isSelectOne) {
+                        return originalQuery.call(
+                            conn,
+                            'SELECT pg_sleep(1)',
+                            ...args,
+                        );
+                    }
+
+                    return originalQuery.call(conn, queryConfig, ...args);
+                };
+                return conn;
+            });
+            return pending;
+        };
+
+        try {
+            await request.get('/ready').expect(503);
+        } finally {
+            pool.acquire = originalAcquire;
+        }
+    });
+});
+
+describe('DB is down', () => {
+    beforeAll(async () => {
+        db = await dbInit();
+    });
+
+    test('when checkDb is disabled, returns readiness good', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            undefined,
+            db.rawDatabase,
+        );
+        await db.destroy();
+        await request
+            .get('/ready')
+            .expect('Content-Type', /json/)
+            .expect(200)
+            .expect('{"health":"GOOD"}');
+    });
+
+    test('when checkDb is enabled, fails readiness check', async () => {
+        const { request } = await setupAppWithCustomConfig(
+            db.stores,
+            { checkDbOnReady: true },
+            db.rawDatabase,
+        );
+        await db.destroy();
+        await request.get('/ready').expect(503);
+    });
+});