From 95a76489c74725262476c3fec9d61a698565a42f Mon Sep 17 00:00:00 2001 From: Serena <94026305+serenakeyitan@users.noreply.github.com> Date: Thu, 23 Jul 2026 01:56:50 -0700 Subject: [PATCH] fix(server): bound /healthz database probes with a TTL cache MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit /healthz disables rate limiting and previously executed SELECT 1 on every request, so public traffic or aggressive probes translated into unlimited database round trips — including while the database was already unhealthy. Cache the probe result (healthy or not) for a short TTL and share a single in-flight probe across concurrent requests, bounding database load to at most one probe per window regardless of request volume. The endpoint contract (200/503 bodies) and rateLimit: false are unchanged so orchestrator probes are never rejected. Fixes #1716 --- .../__tests__/api-small-routes-extra.test.ts | 4 +- .../src/__tests__/healthz-probe-cache.test.ts | 152 ++++++++++++++++++ packages/server/src/api/healthz.ts | 56 ++++++- 3 files changed, 204 insertions(+), 8 deletions(-) create mode 100644 packages/server/src/__tests__/healthz-probe-cache.test.ts diff --git a/packages/server/src/__tests__/api-small-routes-extra.test.ts b/packages/server/src/__tests__/api-small-routes-extra.test.ts index e545e23dc..40730b2fc 100644 --- a/packages/server/src/__tests__/api-small-routes-extra.test.ts +++ b/packages/server/src/__tests__/api-small-routes-extra.test.ts @@ -327,7 +327,9 @@ describe("small API route handlers", () => { const { app, routes } = makeApp({ db: { execute } }); await healthRoutes(app as never); - await healthzRoutes(app as never); + // Disable the healthz probe cache so this test can exercise both the + // healthy and degraded branches through a single registration. + await healthzRoutes(app as never, { probeCacheTtlMs: 0 }); await expect(route(routes, "GET", "/health").handler()).resolves.toEqual({ db: "connected", status: "ok" }); const okReply = makeReply(); diff --git a/packages/server/src/__tests__/healthz-probe-cache.test.ts b/packages/server/src/__tests__/healthz-probe-cache.test.ts new file mode 100644 index 000000000..f42dadd00 --- /dev/null +++ b/packages/server/src/__tests__/healthz-probe-cache.test.ts @@ -0,0 +1,152 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { HealthzOptions } from "../api/healthz.js"; +import { DEFAULT_PROBE_CACHE_TTL_MS, healthzRoutes } from "../api/healthz.js"; + +type UnknownFn = (...args: unknown[]) => unknown; +type CapturedRoute = { + path: string; + options: unknown; + handler: UnknownFn; +}; + +type ReplyDouble = { + status: ReturnType; + send: ReturnType; +}; + +function replyDouble(): ReplyDouble { + const reply = { + status: vi.fn(() => reply), + send: vi.fn((body: unknown) => body), + }; + return reply; +} + +async function registerHealthz( + execute: UnknownFn, + opts?: HealthzOptions, +): Promise<{ handler: UnknownFn; options: unknown }> { + const routes: CapturedRoute[] = []; + const app = { + db: { execute }, + get: (path: string, options: unknown, handler: UnknownFn) => { + routes.push({ path, options, handler }); + return app; + }, + }; + await healthzRoutes(app as never, opts); + expect(routes).toHaveLength(1); + const route = routes[0]; + if (!route) throw new Error("Route was not captured"); + expect(route.path).toBe("/healthz"); + return { handler: route.handler, options: route.options }; +} + +describe("/healthz database probe cache", () => { + afterEach(() => { + vi.useRealTimers(); + vi.clearAllMocks(); + }); + + it("keeps rate limiting disabled so orchestrator probes are never rejected", async () => { + const { options } = await registerHealthz(vi.fn(async () => [{ one: 1 }])); + expect(options).toEqual({ config: { rateLimit: false } }); + }); + + it("serves repeated requests within the TTL from one database probe", async () => { + vi.useFakeTimers(); + const execute = vi.fn(async () => [{ one: 1 }]); + const { handler } = await registerHealthz(execute); + + for (let i = 0; i < 5; i++) { + const reply = replyDouble(); + await handler({}, reply); + expect(reply.status).toHaveBeenCalledWith(200); + expect(reply.send).toHaveBeenCalledWith({ status: "ok" }); + } + expect(execute).toHaveBeenCalledTimes(1); + }); + + it("probes the database again after the TTL expires", async () => { + vi.useFakeTimers(); + const execute = vi.fn(async () => [{ one: 1 }]); + const { handler } = await registerHealthz(execute); + + await handler({}, replyDouble()); + expect(execute).toHaveBeenCalledTimes(1); + + vi.advanceTimersByTime(DEFAULT_PROBE_CACHE_TTL_MS - 1); + await handler({}, replyDouble()); + expect(execute).toHaveBeenCalledTimes(1); + + vi.advanceTimersByTime(1); + await handler({}, replyDouble()); + expect(execute).toHaveBeenCalledTimes(2); + }); + + it("caches an unhealthy probe so a down database is not hammered", async () => { + vi.useFakeTimers(); + const execute = vi.fn(async () => { + throw new Error("connection refused"); + }); + const { handler } = await registerHealthz(execute); + + for (let i = 0; i < 5; i++) { + const reply = replyDouble(); + await handler({}, reply); + expect(reply.status).toHaveBeenCalledWith(503); + expect(reply.send).toHaveBeenCalledWith({ status: "error", message: "database unreachable" }); + } + expect(execute).toHaveBeenCalledTimes(1); + }); + + it("reports recovery on the first probe after the TTL expires", async () => { + vi.useFakeTimers(); + const execute = vi + .fn() + .mockRejectedValueOnce(new Error("connection refused")) + .mockResolvedValue([{ one: 1 }]); + const { handler } = await registerHealthz(execute); + + const downReply = replyDouble(); + await handler({}, downReply); + expect(downReply.status).toHaveBeenCalledWith(503); + + vi.advanceTimersByTime(DEFAULT_PROBE_CACHE_TTL_MS); + const upReply = replyDouble(); + await handler({}, upReply); + expect(upReply.status).toHaveBeenCalledWith(200); + expect(execute).toHaveBeenCalledTimes(2); + }); + + it("shares one in-flight probe across concurrent requests", async () => { + let resolveProbe: (() => void) | undefined; + const execute = vi.fn( + () => + new Promise((resolve) => { + resolveProbe = resolve; + }), + ); + const { handler } = await registerHealthz(execute); + + const first = replyDouble(); + const second = replyDouble(); + const pending = Promise.all([handler({}, first), handler({}, second)]); + expect(execute).toHaveBeenCalledTimes(1); + + resolveProbe?.(); + await pending; + expect(execute).toHaveBeenCalledTimes(1); + expect(first.status).toHaveBeenCalledWith(200); + expect(second.status).toHaveBeenCalledWith(200); + }); + + it("probes on every request when probeCacheTtlMs is 0", async () => { + const execute = vi.fn(async () => [{ one: 1 }]); + const { handler } = await registerHealthz(execute, { probeCacheTtlMs: 0 }); + + await handler({}, replyDouble()); + await handler({}, replyDouble()); + expect(execute).toHaveBeenCalledTimes(2); + }); +}); diff --git a/packages/server/src/api/healthz.ts b/packages/server/src/api/healthz.ts index ce87f39a9..af34a5743 100644 --- a/packages/server/src/api/healthz.ts +++ b/packages/server/src/api/healthz.ts @@ -1,18 +1,60 @@ import { sql } from "drizzle-orm"; import type { FastifyInstance } from "fastify"; +/** + * Default TTL for the cached database probe result. Bounds `/healthz` database + * round trips to at most one `SELECT 1` per window regardless of request + * volume, while staying far finer-grained than orchestrator probe cadence + * (the image `HEALTHCHECK` polls every 30s). + */ +export const DEFAULT_PROBE_CACHE_TTL_MS = 5_000; + +export type HealthzOptions = { + /** + * TTL in milliseconds for the cached database probe result. + * `0` disables caching and probes the database on every request. + */ + probeCacheTtlMs?: number; +}; + /** * Root-level health check endpoint for container orchestration. * Returns 200 when healthy, 503 when degraded. * Used by Docker HEALTHCHECK, Railway, Fly.io, Kubernetes liveness/readiness probes. + * + * The endpoint is public and deliberately keeps `rateLimit: false` so + * orchestrator probes can never be rejected with 429. Database pressure is + * bounded instead: the `SELECT 1` probe result (healthy or not) is cached for + * a short TTL and concurrent requests share a single in-flight probe, so + * public traffic or aggressive probing cannot translate into unlimited + * database round trips — including while the database is already unhealthy. */ -export async function healthzRoutes(app: FastifyInstance): Promise { - app.get("/healthz", { config: { rateLimit: false } }, async (_request, reply) => { - try { - await app.db.execute(sql`SELECT 1`); - return reply.status(200).send({ status: "ok" }); - } catch { - return reply.status(503).send({ status: "error", message: "database unreachable" }); +export async function healthzRoutes(app: FastifyInstance, opts: HealthzOptions = {}): Promise { + const probeCacheTtlMs = opts.probeCacheTtlMs ?? DEFAULT_PROBE_CACHE_TTL_MS; + + let cached: { healthy: boolean; expiresAt: number } | null = null; + let inflight: Promise | null = null; + + const probeDatabase = (): Promise => { + if (inflight === null) { + inflight = app.db + .execute(sql`SELECT 1`) + .then( + () => true, + () => false, + ) + .then((healthy) => { + cached = { healthy, expiresAt: Date.now() + probeCacheTtlMs }; + inflight = null; + return healthy; + }); } + return inflight; + }; + + app.get("/healthz", { config: { rateLimit: false } }, async (_request, reply) => { + const healthy = cached !== null && Date.now() < cached.expiresAt ? cached.healthy : await probeDatabase(); + if (healthy) return reply.status(200).send({ status: "ok" }); + return reply.status(503).send({ status: "error", message: "database unreachable" }); }); }