diff --git a/src/lib/__snapshots__/create-config.test.ts.snap b/src/lib/__snapshots__/create-config.test.ts.snap index fe6bece2cb..b407729b53 100644 --- a/src/lib/__snapshots__/create-config.test.ts.snap +++ b/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, diff --git a/src/lib/create-config.ts b/src/lib/create-config.ts index 3512094ffd..426e86787c 100644 --- a/src/lib/create-config.ts +++ b/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, }; } diff --git a/src/lib/openapi/spec/index.ts b/src/lib/openapi/spec/index.ts index b7d69c8ee2..6788e0a44c 100644 --- a/src/lib/openapi/spec/index.ts +++ b/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'; diff --git a/src/lib/openapi/spec/ready-check-schema.ts b/src/lib/openapi/spec/ready-check-schema.ts new file mode 100644 index 0000000000..0e0e0fee5e --- /dev/null +++ b/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; diff --git a/src/lib/routes/health-check.test.ts b/src/lib/routes/health-check.test.ts index 89ab26e8ec..b32176f398 100644 --- a/src/lib/routes/health-check.test.ts +++ b/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') diff --git a/src/lib/routes/index.ts b/src/lib/routes/index.ts index 79e4d6d061..d59a8bae36 100644 --- a/src/lib/routes/index.ts +++ b/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', diff --git a/src/lib/routes/ready-check.ts b/src/lib/routes/ready-check.ts new file mode 100644 index 0000000000..381e90e1d4 --- /dev/null +++ b/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, + 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): Promise { + 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 { + 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'); +} diff --git a/src/lib/types/option.ts b/src/lib/types/option.ts index 9a5d7a8ed6..ef0df45051 100644 --- a/src/lib/types/option.ts +++ b/src/lib/types/option.ts @@ -175,6 +175,7 @@ export interface IUnleashOptions { resourceLimits?: Partial; userInactivityThresholdInDays?: number; unleashFrontendToken?: string; + checkDbOnReady?: boolean; } export interface IEmailOption { @@ -301,4 +302,5 @@ export interface IUnleashConfig { userInactivityThresholdInDays: number; buildDate?: string; unleashFrontendToken?: string; + checkDbOnReady?: boolean; } diff --git a/src/test/e2e/ready.e2e.test.ts b/src/test/e2e/ready.e2e.test.ts new file mode 100644 index 0000000000..f84ef6c388 --- /dev/null +++ b/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); + }); +});