fix: 隔离 Redis 健康检查客户端代际

This commit is contained in:
root
2026-09-14 14:42:33 +08:00
parent bbd790d86d
commit 6b901e074e
2 changed files with 258 additions and 40 deletions
+164
View File
@@ -14,6 +14,7 @@ import { ErrorCode, ok } from "@drift/contracts";
import { Type } from "class-transformer"; import { Type } from "class-transformer";
import { IsInt, IsString } from "class-validator"; import { IsInt, IsString } from "class-validator";
import { createServer, type Socket } from "node:net"; import { createServer, type Socket } from "node:net";
import { EventEmitter } from "node:events";
import request from "supertest"; import request from "supertest";
import { afterAll, beforeAll, describe, expect, it, vi } from "vitest"; import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
import { AppModule } from "../app.module.js"; import { AppModule } from "../app.module.js";
@@ -24,6 +25,26 @@ import { PrismaService } from "../database/prisma.service.js";
import { bootstrap, configureApp } from "../main.js"; import { bootstrap, configureApp } from "../main.js";
import { HealthService } from "./health.service.js"; import { HealthService } from "./health.service.js";
type RedisClientDouble = EventEmitter & {
isOpen: boolean;
connect: ReturnType<typeof vi.fn>;
ping: ReturnType<typeof vi.fn>;
quit: ReturnType<typeof vi.fn>;
destroy: ReturnType<typeof vi.fn>;
};
function redisClientDouble(ping: () => Promise<string>): RedisClientDouble {
const client = new EventEmitter() as RedisClientDouble;
client.isOpen = true;
client.connect = vi.fn().mockResolvedValue(undefined);
client.ping = vi.fn(ping);
client.quit = vi.fn().mockResolvedValue(undefined);
client.destroy = vi.fn(() => {
client.isOpen = false;
});
return client;
}
class ProbeDto { class ProbeDto {
@IsString() @IsString()
name!: string; name!: string;
@@ -386,6 +407,136 @@ describe("API infrastructure", () => {
} }
}); });
it("does not let an old Redis probe deadline destroy the replacement client", async () => {
vi.useFakeTimers();
const previousTimeout = process.env.REDIS_PROBE_TIMEOUT_MS;
process.env.REDIS_PROBE_TIMEOUT_MS = "100";
let rejectSecondProbe!: (error: Error) => void;
const firstProbe = new Promise<string>(() => undefined);
const secondProbe = new Promise<string>((_resolve, reject) => {
rejectSecondProbe = reject;
});
const pingA = vi
.fn<() => Promise<string>>()
.mockReturnValueOnce(firstProbe)
.mockReturnValueOnce(secondProbe);
const clientA = redisClientDouble(pingA);
const clientB = redisClientDouble(() => Promise.resolve("PONG"));
const service = new HealthService(
{
$queryRaw: vi.fn().mockResolvedValue([{ value: 1 }]),
} as unknown as PrismaService,
vi
.fn()
.mockReturnValueOnce(clientA)
.mockReturnValueOnce(clientB) as never,
);
try {
const oldPendingCheck = service.check();
const oldPendingResult = expect(oldPendingCheck).rejects.toMatchObject({
code: ErrorCode.SERVICE_UNAVAILABLE,
});
const failingCheck = service.check();
rejectSecondProbe(new Error("A failed"));
await expect(failingCheck).rejects.toMatchObject({
code: ErrorCode.SERVICE_UNAVAILABLE,
});
await expect(service.check()).resolves.toEqual({
status: "ok",
postgres: "up",
redis: "up",
});
await vi.advanceTimersByTimeAsync(100);
await oldPendingResult;
expect(clientB.destroy).not.toHaveBeenCalled();
await expect(service.check()).resolves.toEqual({
status: "ok",
postgres: "up",
redis: "up",
});
} finally {
await service.onModuleDestroy();
process.env.REDIS_PROBE_TIMEOUT_MS = previousTimeout;
vi.useRealTimers();
}
});
it("logs the first error from a replacement Redis client independently", async () => {
const warning = vi
.spyOn(Logger.prototype, "warn")
.mockImplementation(() => undefined);
const clientA = redisClientDouble(() =>
Promise.reject(new Error("A failed")),
);
const clientB = redisClientDouble(() => Promise.resolve("PONG"));
const service = new HealthService(
{
$queryRaw: vi.fn().mockResolvedValue([{ value: 1 }]),
} as unknown as PrismaService,
vi
.fn()
.mockReturnValueOnce(clientA)
.mockReturnValueOnce(clientB) as never,
);
try {
await expect(service.check()).rejects.toMatchObject({
code: ErrorCode.SERVICE_UNAVAILABLE,
});
await expect(service.check()).resolves.toMatchObject({ status: "ok" });
clientA.emit("error", new Error("late A error"));
clientB.emit("error", new Error("first B error"));
expect(warning).toHaveBeenCalledTimes(2);
} finally {
await service.onModuleDestroy();
warning.mockRestore();
}
});
it("shares one replacement client and connection across concurrent checks", async () => {
let resolveConnect!: () => void;
const clientA = redisClientDouble(() =>
Promise.reject(new Error("A failed")),
);
const clientB = redisClientDouble(() => Promise.resolve("PONG"));
clientB.isOpen = false;
clientB.connect = vi.fn(
() =>
new Promise<void>((resolve) => {
resolveConnect = () => {
clientB.isOpen = true;
resolve();
};
}),
);
const factory = vi
.fn()
.mockReturnValueOnce(clientA)
.mockReturnValueOnce(clientB);
const service = new HealthService(
{
$queryRaw: vi.fn().mockResolvedValue([{ value: 1 }]),
} as unknown as PrismaService,
factory as never,
);
try {
await expect(service.check()).rejects.toMatchObject({
code: ErrorCode.SERVICE_UNAVAILABLE,
});
const checks = [service.check(), service.check(), service.check()];
expect(factory).toHaveBeenCalledTimes(2);
await vi.waitFor(() => expect(clientB.connect).toHaveBeenCalledOnce());
resolveConnect();
await expect(Promise.all(checks)).resolves.toHaveLength(3);
expect(clientB.connect).toHaveBeenCalledOnce();
} finally {
await service.onModuleDestroy();
}
});
it("destroys Redis when graceful module shutdown exceeds its deadline", async () => { it("destroys Redis when graceful module shutdown exceeds its deadline", async () => {
const previousTimeout = process.env.REDIS_PROBE_TIMEOUT_MS; const previousTimeout = process.env.REDIS_PROBE_TIMEOUT_MS;
process.env.REDIS_PROBE_TIMEOUT_MS = "25"; process.env.REDIS_PROBE_TIMEOUT_MS = "25";
@@ -406,6 +557,19 @@ describe("API infrastructure", () => {
} }
}); });
it("does not recreate Redis after module shutdown begins", async () => {
const client = redisClientDouble(() => Promise.resolve("PONG"));
client.isOpen = false;
const factory = vi.fn(() => client);
const service = new HealthService({} as PrismaService, factory as never);
await service.onModuleDestroy();
await expect(service.check()).rejects.toMatchObject({
code: ErrorCode.SERVICE_UNAVAILABLE,
});
expect(factory).toHaveBeenCalledOnce();
});
it("returns a sanitized 503 specifically when PostgreSQL is unavailable", async () => { it("returns a sanitized 503 specifically when PostgreSQL is unavailable", async () => {
const module = await Test.createTestingModule({ imports: [AppModule] }) const module = await Test.createTestingModule({ imports: [AppModule] })
.overrideProvider(PrismaService) .overrideProvider(PrismaService)
+90 -36
View File
@@ -1,78 +1,122 @@
import { import {
HttpStatus,
Inject, Inject,
Injectable, Injectable,
Logger, Logger,
Optional,
type OnModuleDestroy, type OnModuleDestroy,
} from "@nestjs/common"; } from "@nestjs/common";
import { createClient } from "redis"; import { createClient } from "redis";
import { PrismaService } from "../database/prisma.service.js";
import { DomainException } from "../common/domain.exception.js";
import { ErrorCode } from "@drift/contracts"; import { ErrorCode } from "@drift/contracts";
import { HttpStatus } from "@nestjs/common"; import { DomainException } from "../common/domain.exception.js";
import { PrismaService } from "../database/prisma.service.js";
type RedisClient = ReturnType<typeof createClient>;
type RedisClientFactory = (
options: Parameters<typeof createClient>[0],
) => RedisClient;
export const REDIS_CLIENT_FACTORY = Symbol("REDIS_CLIENT_FACTORY");
@Injectable() @Injectable()
export class HealthService implements OnModuleDestroy { export class HealthService implements OnModuleDestroy {
private redis: ReturnType<typeof createClient>; private redis: RedisClient | undefined;
private lastRedis: RedisClient | undefined;
private readonly probeTimeoutMs: number; private readonly probeTimeoutMs: number;
private readonly logger = new Logger(HealthService.name); private readonly logger = new Logger(HealthService.name);
private redisDestroyed = false; private readonly destroyedClients = new WeakSet<object>();
private redisErrorLogged = false; private readonly connections = new WeakMap<object, Promise<void>>();
private generation = 0;
private shuttingDown = false;
get redisClient(): ReturnType<typeof createClient> { get redisClient(): RedisClient {
return this.redis; const client = this.redis ?? this.lastRedis;
if (!client) throw new Error("Redis client is unavailable");
return client;
} }
constructor(@Inject(PrismaService) private readonly prisma: PrismaService) { constructor(
@Inject(PrismaService) private readonly prisma: PrismaService,
@Optional()
@Inject(REDIS_CLIENT_FACTORY)
private readonly redisClientFactory: RedisClientFactory = createClient,
) {
this.probeTimeoutMs = this.readTimeout(process.env.REDIS_PROBE_TIMEOUT_MS); this.probeTimeoutMs = this.readTimeout(process.env.REDIS_PROBE_TIMEOUT_MS);
this.redis = this.createRedisClient(); this.redis = this.createRedisClient();
this.lastRedis = this.redis;
} }
private createRedisClient(): ReturnType<typeof createClient> { private createRedisClient(): RedisClient {
const redis = createClient({ const redis = this.redisClientFactory({
url: process.env.REDIS_URL ?? "redis://127.0.0.1:6379", url: process.env.REDIS_URL ?? "redis://127.0.0.1:6379",
socket: { connectTimeout: this.probeTimeoutMs, reconnectStrategy: false }, socket: { connectTimeout: this.probeTimeoutMs, reconnectStrategy: false },
}); });
let errorLogged = false;
redis.on("error", () => { redis.on("error", () => {
if (this.redisErrorLogged) return; if (errorLogged) return;
this.redisErrorLogged = true; errorLogged = true;
this.logger.warn({ event: "Redis client error", code: "REDIS_ERROR" }); this.logger.warn({ event: "Redis client error", code: "REDIS_ERROR" });
}); });
return redis; return redis;
} }
private getRedisClient(): ReturnType<typeof createClient> { private getRedisClient(): { client: RedisClient; generation: number } {
if (this.redisDestroyed) { if (this.shuttingDown) throw this.unavailable();
this.redisErrorLogged = false; if (!this.redis) {
this.redis = this.createRedisClient(); this.redis = this.createRedisClient();
this.redisDestroyed = false; this.lastRedis = this.redis;
this.generation += 1;
} }
return this.redis; return { client: this.redis, generation: this.generation };
} }
async check(): Promise<{ status: "ok"; postgres: "up"; redis: "up" }> { async check(): Promise<{ status: "ok"; postgres: "up"; redis: "up" }> {
const redis = this.getRedisClient(); if (this.shuttingDown) throw this.unavailable();
const { client, generation } = this.getRedisClient();
try { try {
await this.withDeadline(async () => { await this.withDeadline(client, async () => {
if (!redis.isOpen) await redis.connect(); await this.connectRedis(client);
await Promise.all([this.prisma.$queryRaw`SELECT 1`, redis.ping()]); await Promise.all([this.prisma.$queryRaw`SELECT 1`, client.ping()]);
}); });
return { status: "ok", postgres: "up", redis: "up" }; return { status: "ok", postgres: "up", redis: "up" };
} catch { } catch {
this.destroyRedis(); this.destroyRedis(client, generation);
throw new DomainException( throw this.unavailable();
}
}
async onModuleDestroy(): Promise<void> {
this.shuttingDown = true;
const client = this.redis;
const generation = this.generation;
if (!client || !client.isOpen) return;
try {
await this.withDeadline(client, () => client.quit());
} catch {
this.destroyRedis(client, generation);
}
}
private unavailable(): DomainException {
return new DomainException(
ErrorCode.SERVICE_UNAVAILABLE, ErrorCode.SERVICE_UNAVAILABLE,
"Dependencies unavailable", "Dependencies unavailable",
HttpStatus.SERVICE_UNAVAILABLE, HttpStatus.SERVICE_UNAVAILABLE,
); );
} }
}
async onModuleDestroy(): Promise<void> { private async connectRedis(client: RedisClient): Promise<void> {
if (!this.redis.isOpen) return; if (client.isOpen) return;
const pending = this.connections.get(client);
if (pending) return pending;
const connection = client.connect().then(() => undefined);
this.connections.set(client, connection);
try { try {
await this.withDeadline(() => this.redis.quit()); await connection;
} catch { } finally {
this.destroyRedis(); if (this.connections.get(client) === connection) {
this.connections.delete(client);
}
} }
} }
@@ -81,10 +125,13 @@ export class HealthService implements OnModuleDestroy {
return Number.isFinite(timeout) && timeout > 0 ? timeout : 1000; return Number.isFinite(timeout) && timeout > 0 ? timeout : 1000;
} }
private withDeadline<T>(operation: () => Promise<T>): Promise<T> { private withDeadline<T>(
client: RedisClient,
operation: () => Promise<T>,
): Promise<T> {
return new Promise<T>((resolve, reject) => { return new Promise<T>((resolve, reject) => {
const timer = setTimeout(() => { const timer = setTimeout(() => {
this.destroyRedis(); this.destroyRedis(client);
reject(new Error("Redis operation timed out")); reject(new Error("Redis operation timed out"));
}, this.probeTimeoutMs); }, this.probeTimeoutMs);
void operation().then( void operation().then(
@@ -104,13 +151,20 @@ export class HealthService implements OnModuleDestroy {
}); });
} }
private destroyRedis(): void { private destroyRedis(client: RedisClient, generation?: number): void {
if (this.redisDestroyed) return; if (!this.destroyedClients.has(client)) {
this.redisDestroyed = true; this.destroyedClients.add(client);
try { try {
this.redis.destroy(); client.destroy();
} catch { } catch {
// Destruction is best-effort and must never mask readiness/shutdown errors. // Destruction is best-effort and must never mask readiness/shutdown errors.
} }
} }
if (
this.redis === client &&
(generation === undefined || generation === this.generation)
) {
this.redis = undefined;
}
}
} }