Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ export const {db, CoreBaseModel, helpers} = initDB({
- `topologyMode`: Connection topology, either `primary-replica` (the default) or `proxy`
- `knexOptions`: Non-required additional options that will be passed to Knex before initialization
- `onKnexCreated`: Optional callback called synchronously for each Knex instance created by the dispatcher, before database health checks start. Use it to attach instrumentation, event listeners, or plugins that do not require an active database connection. If it throws, initialization fails. The callback receives the Knex instance and returns nothing.
- `onHealthcheck`: Optional synchronous callback called after each database health check with a snapshot of every host's availability, latency, and role. In `primary-replica` mode, unavailable hosts have the `unknown` role; proxy connections do not include a role. It is called even when `suppressStatusLogs` is enabled. Callback errors do not interrupt database routing or subsequent health checks.

When all connection strings point to equivalent proxy or router instances, such as SPQR routers, use `proxy` mode. Healthy endpoints are then eligible for both primary and replica queries, and the endpoint with the lowest latest health-check latency is selected:

Expand Down
18 changes: 15 additions & 3 deletions lib/core.ts
Original file line number Diff line number Diff line change
Expand Up @@ -4,9 +4,18 @@ import {type Constructor, Model} from 'objection';

import {defaultDispatcherOptions, defaultExLogger, defaultKnexOptions} from './constants';
import {PGDispatcher} from './dispatcher';
import type {BaseModel, ExLogger, TopologyMode} from './types';

export type {TopologyMode} from './types';
import type {BaseModel, ExLogger, PGHealthcheckHandler, TopologyMode} from './types';

export type {
PGConnectionRole,
PGConnectionStatus,
PGHealthcheckHandler,
PGHealthcheckStatus,
PGPrimaryReplicaConnectionStatus,
PGPrimaryReplicaHealthcheckStatus,
PGProxyHealthcheckStatus,
TopologyMode,
} from './types';

export interface CoreDBDispatcherOptions {
healthcheckInterval?: number;
Expand All @@ -25,6 +34,7 @@ export interface CoreDBConstructorArgs {
logger?: ExLogger;
modelParams?: GetModelParams;
onKnexCreated?: (knex: Knex) => void;
onHealthcheck?: PGHealthcheckHandler;
}

export function getModel(params: GetModelParams = {}): typeof BaseModel {
Expand Down Expand Up @@ -89,6 +99,7 @@ export function initDB({
logger = defaultExLogger,
modelParams,
onKnexCreated,
onHealthcheck,
}: CoreDBConstructorArgs) {
if (!connectionString) {
throw new Error('Empty connection string');
Expand All @@ -102,6 +113,7 @@ export function initDB({
knexOptions: mergedKnexOptions,
logger,
onKnexCreated,
onHealthcheck,
});

const terminate = () => {
Expand Down
83 changes: 77 additions & 6 deletions lib/dispatcher.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,14 @@ import {URL} from 'url';
import knexBuilder from 'knex';
import type {Knex} from 'knex';

import type {Dict, ExLogger, PDOptions} from './types';
import type {
Dict,
ExLogger,
PDOptions,
PGConnectionRole,
PGHealthcheckHandler,
PGHealthcheckStatus,
} from './types';

import Timeout = NodeJS.Timer;

Expand All @@ -26,6 +33,7 @@ export interface PDConstructorArgs {
logger: ExLogger;

onKnexCreated?: (knex: Knex) => void;
onHealthcheck?: PGHealthcheckHandler;
}

interface PDConnection {
Expand All @@ -51,15 +59,18 @@ export class PGDispatcher {
private connections: PDConnection[];
private options: PDOptions;
private logger: {info: InfoLogger; error: ErrorLogger};
private onHealthcheck?: PGHealthcheckHandler;
private hcTimer?: Timeout | null;
private isInit = false;
private isTerminating = false;

constructor({
connections = [],
options,
knexOptions = {},
logger,
onKnexCreated,
onHealthcheck,
}: PDConstructorArgs) {
if (!connections.length) {
throw new Error('Empty connections list is not allowed');
Expand All @@ -83,6 +94,7 @@ export class PGDispatcher {
});
this.options = options;
this.knexOptions = knexOptions;
this.onHealthcheck = onHealthcheck;

this.logger = {
info: ({message, data}) => {
Expand Down Expand Up @@ -120,6 +132,8 @@ export class PGDispatcher {
}

terminate() {
this.isTerminating = true;

if (this.hcTimer) {
clearInterval(this.hcTimer);
}
Expand Down Expand Up @@ -200,6 +214,10 @@ export class PGDispatcher {
private async initHealthcheck() {
await this.knexReady();

if (this.isTerminating) {
return;
}

const performHealthcheck = () => {
const checkups = this.connections.map((connection) =>
this.checkDatabase(connection).catch((error) => {
Expand All @@ -209,18 +227,21 @@ export class PGDispatcher {
}),
);
Promise.all(checkups).then(() => {
// Connections hold shared current state; this is not an isolated per-cycle result.
const status = this.getHealthcheckStatus();

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Build the snapshot from one health-check cycle

Health-check cycles can overlap when healthcheckInterval is shorter than a check. Each cycle mutates the shared this.connections, while getHealthcheckStatus() reads that shared state only after its own Promise.all completes. A slower older cycle can therefore publish a mixture of its own results and values written by a newer cycle. Please serialize cycles or build the snapshot from per-cycle local results.

AI generated

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The callback intentionally mirrors the existing Database current status semantics and snapshots shared current connection state rather than cycle-local results. I added a short comment to make this explicit. Serializing healthcheck cycles would change existing dispatcher behavior and is better handled separately.

this.logger.info({
message: 'Database current status',
data: {
...(this.isProxyMode ? {topologyMode: this.options.topologyMode} : {}),
connections: this.connections.map((c) => ({
host: c.host,
...(this.isProxyMode ? {} : {primary: c.primary}),
healthy: c.healthy,
latency: c.latency,
connections: this.connections.map((connection) => ({
host: connection.host,
...(this.isProxyMode ? {} : {primary: connection.primary}),
healthy: connection.healthy,
latency: connection.latency,
})),
},
});
this.notifyHealthcheck(status);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] Do not notify after terminate() has completed

terminate() clears the interval but does not cancel or await an already-running health check. Its Promise.all can finish after teardown and call onHealthcheck after terminate() has resolved, when the consumer may already have closed its exporter or registry. Please track the lifecycle and either await in-flight work or suppress notifications once termination starts.

AI generated

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in 7e244d1. Termination now marks the dispatcher before clearing its timer, prevents initialization from starting a healthcheck after termination, and suppresses callback notification from an in-flight check.

});
};

Expand Down Expand Up @@ -299,6 +320,56 @@ export class PGDispatcher {
await Promise.all(this.connections.map((c) => c.knex));
}

private getHealthcheckStatus(): PGHealthcheckStatus {
if (this.isProxyMode) {
return {
topologyMode: 'proxy',
connections: this.connections.map((connection) => ({
host: connection.host,
healthy: connection.healthy,
latency: connection.latency,
})),
};
}

return {
topologyMode: 'primary-replica',
connections: this.connections.map((connection) => ({
host: connection.host,
role: this.getConnectionRole(connection),
healthy: connection.healthy,
latency: connection.latency,
})),
};
}

private getConnectionRole(connection: PDConnection): PGConnectionRole {
if (!connection.healthy) {
return 'unknown';
}

return connection.primary ? 'primary' : 'replica';
}

private notifyHealthcheck(status: PGHealthcheckStatus) {
if (!this.onHealthcheck || this.isTerminating) {
return;
}

try {
this.onHealthcheck(status);
} catch (error) {
this.reportHealthcheckCallbackError(error);
}
}

private reportHealthcheckCallbackError(error: unknown) {
this.logger.error({
message: 'Database healthcheck callback failed',
error: error as Error,
});
}

private get healthyConnections() {
return this.connections.filter((c) => c.healthy);
}
Expand Down
26 changes: 26 additions & 0 deletions lib/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,32 @@ import type {PGDispatcher} from './dispatcher';

export type TopologyMode = 'primary-replica' | 'proxy';

export interface PGConnectionStatus {
readonly host: string;
readonly healthy: boolean;
readonly latency: number;
}

export type PGConnectionRole = 'primary' | 'replica' | 'unknown';

export interface PGPrimaryReplicaConnectionStatus extends PGConnectionStatus {
readonly role: PGConnectionRole;
}

export interface PGPrimaryReplicaHealthcheckStatus {
readonly topologyMode: 'primary-replica';
readonly connections: readonly PGPrimaryReplicaConnectionStatus[];
}

export interface PGProxyHealthcheckStatus {
readonly topologyMode: 'proxy';
readonly connections: readonly PGConnectionStatus[];
}

export type PGHealthcheckStatus = PGPrimaryReplicaHealthcheckStatus | PGProxyHealthcheckStatus;

export type PGHealthcheckHandler = (status: PGHealthcheckStatus) => void;

export interface PDOptions {
healthcheckInterval: number;
healthcheckTimeout: number;
Expand Down
110 changes: 109 additions & 1 deletion tests/dispatcher.test.js
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ function failedCheckup() {
return Promise.reject(new Error('Proxy unavailable'));
}

function createDispatcher(clients, options = {}) {
function createDispatcher(clients, options = {}, onHealthcheck) {
clients.forEach((client) => knexBuilder.mockImplementationOnce(() => client));

const logger = {
Expand All @@ -48,6 +48,7 @@ function createDispatcher(clients, options = {}) {
healthcheckTimeout: 100,
...options,
},
onHealthcheck,
});

activeDispatchers.push(dispatcher);
Expand Down Expand Up @@ -145,3 +146,110 @@ describe('PGDispatcher topology modes', () => {
}
});
});

describe('PGDispatcher healthcheck callback', () => {
test('reports a snapshot for every primary and replica connection', async () => {
const onHealthcheck = jest.fn();
const primary = createKnex(successfulCheckup({pg_is_in_recovery: false}));
const replica = createKnex(successfulCheckup({pg_is_in_recovery: true}));
const unavailable = createKnex(failedCheckup);
const {dispatcher, logger} = createDispatcher(
[primary, replica, unavailable],
{},
onHealthcheck,
);

await dispatcher.ready();
await new Promise((resolve) => setImmediate(resolve));

expect(onHealthcheck).toHaveBeenCalledWith({
topologyMode: 'primary-replica',
connections: [
{
host: 'database-0.example',
role: 'primary',
healthy: true,
latency: expect.any(Number),
},
{
host: 'database-1.example',
role: 'replica',
healthy: true,
latency: expect.any(Number),
},
{
host: 'database-2.example',
role: 'unknown',
healthy: false,
latency: expect.any(Number),
},
],
});

const statusLog = logger.info.mock.calls.find(
([message]) => message === 'Database current status',
);
expect(statusLog[1].connections).toEqual([
expect.objectContaining({host: 'database-0.example', primary: true}),
expect.objectContaining({host: 'database-1.example', primary: false}),
expect.objectContaining({host: 'database-2.example', primary: false}),
]);
});

test('reports proxy status without primary/replica roles when status logs are suppressed', async () => {
const onHealthcheck = jest.fn();
const {dispatcher, logger} = createDispatcher(
[createKnex(successfulCheckup({value: 1}))],
{suppressStatusLogs: true, topologyMode: 'proxy'},
onHealthcheck,
);

await dispatcher.ready();
await new Promise((resolve) => setImmediate(resolve));

expect(logger.info).not.toHaveBeenCalled();
expect(onHealthcheck).toHaveBeenCalledWith({
topologyMode: 'proxy',
connections: [
{
host: 'database-0.example',
healthy: true,
latency: expect.any(Number),
},
],
});
expect(onHealthcheck.mock.calls[0][0].connections[0]).not.toHaveProperty('role');
});

test('isolates callback errors from database routing', async () => {
const callbackError = new Error('Healthcheck consumer failed');
const onHealthcheck = jest.fn(() => {
throw callbackError;
});
const primary = createKnex(successfulCheckup({pg_is_in_recovery: false}));
const {dispatcher, logger} = createDispatcher([primary], {}, onHealthcheck);

await dispatcher.ready();
await new Promise((resolve) => setImmediate(resolve));

expect(dispatcher.primary).toBe(primary);
expect(logger.error).toHaveBeenCalledWith('PGDispatcher error', callbackError, undefined);
});

test('does not notify after termination starts', async () => {
let finishCheckup;
const checkup = () =>
new Promise((resolve) => {
finishCheckup = resolve;
});
const onHealthcheck = jest.fn();
const {dispatcher} = createDispatcher([createKnex(checkup)], {}, onHealthcheck);

await new Promise((resolve) => setImmediate(resolve));
await dispatcher.terminate();
finishCheckup({rows: [{pg_is_in_recovery: false}]});
await new Promise((resolve) => setImmediate(resolve));

expect(onHealthcheck).not.toHaveBeenCalled();
});
});
Loading