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

apps/api/src/health/health.e2e-spec.ts

@@ -14,6 +14,7 @@ import { ErrorCode, ok } from "@drift/contracts";
 import { Type } from "class-transformer";
 import { IsInt, IsString } from "class-validator";
 import { createServer, type Socket } from "node:net";
+import { EventEmitter } from "node:events";
 import request from "supertest";
 import { afterAll, beforeAll, describe, expect, it, vi } from "vitest";
 import { AppModule } from "../app.module.js";
@@ -24,6 +25,26 @@ import { PrismaService } from "../database/prisma.service.js";
 import { bootstrap, configureApp } from "../main.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 {
   @IsString()
   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 () => {
     const previousTimeout = process.env.REDIS_PROBE_TIMEOUT_MS;
     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 () => {
     const module = await Test.createTestingModule({ imports: [AppModule] })
       .overrideProvider(PrismaService)

apps/api/src/health/health.service.ts

@@ -1,78 +1,122 @@
 import {
+  HttpStatus,
   Inject,
   Injectable,
   Logger,
+  Optional,
   type OnModuleDestroy,
 } from "@nestjs/common";
 import { createClient } from "redis";
-import { PrismaService } from "../database/prisma.service.js";
-import { DomainException } from "../common/domain.exception.js";
 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()
 export class HealthService implements OnModuleDestroy {
-  private redis: ReturnType<typeof createClient>;
+  private redis: RedisClient | undefined;
+  private lastRedis: RedisClient | undefined;
   private readonly probeTimeoutMs: number;
   private readonly logger = new Logger(HealthService.name);
-  private redisDestroyed = false;
-  private redisErrorLogged = false;
+  private readonly destroyedClients = new WeakSet<object>();
+  private readonly connections = new WeakMap<object, Promise<void>>();
+  private generation = 0;
+  private shuttingDown = false;
 
-  get redisClient(): ReturnType<typeof createClient> {
-    return this.redis;
+  get redisClient(): RedisClient {
+    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.redis = this.createRedisClient();
+    this.lastRedis = this.redis;
   }
 
-  private createRedisClient(): ReturnType<typeof createClient> {
-    const redis = createClient({
+  private createRedisClient(): RedisClient {
+    const redis = this.redisClientFactory({
       url: process.env.REDIS_URL ?? "redis://127.0.0.1:6379",
       socket: { connectTimeout: this.probeTimeoutMs, reconnectStrategy: false },
     });
+    let errorLogged = false;
     redis.on("error", () => {
-      if (this.redisErrorLogged) return;
-      this.redisErrorLogged = true;
+      if (errorLogged) return;
+      errorLogged = true;
       this.logger.warn({ event: "Redis client error", code: "REDIS_ERROR" });
     });
     return redis;
   }
 
-  private getRedisClient(): ReturnType<typeof createClient> {
-    if (this.redisDestroyed) {
-      this.redisErrorLogged = false;
+  private getRedisClient(): { client: RedisClient; generation: number } {
+    if (this.shuttingDown) throw this.unavailable();
+    if (!this.redis) {
       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" }> {
-    const redis = this.getRedisClient();
+    if (this.shuttingDown) throw this.unavailable();
+    const { client, generation } = this.getRedisClient();
     try {
-      await this.withDeadline(async () => {
-        if (!redis.isOpen) await redis.connect();
-        await Promise.all([this.prisma.$queryRaw`SELECT 1`, redis.ping()]);
+      await this.withDeadline(client, async () => {
+        await this.connectRedis(client);
+        await Promise.all([this.prisma.$queryRaw`SELECT 1`, client.ping()]);
       });
       return { status: "ok", postgres: "up", redis: "up" };
     } catch {
-      this.destroyRedis();
-      throw new DomainException(
-        ErrorCode.SERVICE_UNAVAILABLE,
-        "Dependencies unavailable",
-        HttpStatus.SERVICE_UNAVAILABLE,
-      );
+      this.destroyRedis(client, generation);
+      throw this.unavailable();
     }
   }
 
   async onModuleDestroy(): Promise<void> {
-    if (!this.redis.isOpen) return;
+    this.shuttingDown = true;
+    const client = this.redis;
+    const generation = this.generation;
+    if (!client || !client.isOpen) return;
     try {
-      await this.withDeadline(() => this.redis.quit());
+      await this.withDeadline(client, () => client.quit());
     } catch {
-      this.destroyRedis();
+      this.destroyRedis(client, generation);
+    }
+  }
+
+  private unavailable(): DomainException {
+    return new DomainException(
+      ErrorCode.SERVICE_UNAVAILABLE,
+      "Dependencies unavailable",
+      HttpStatus.SERVICE_UNAVAILABLE,
+    );
+  }
+
+  private async connectRedis(client: RedisClient): Promise<void> {
+    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 {
+      await connection;
+    } finally {
+      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;
   }
 
-  private withDeadline<T>(operation: () => Promise<T>): Promise<T> {
+  private withDeadline<T>(
+    client: RedisClient,
+    operation: () => Promise<T>,
+  ): Promise<T> {
     return new Promise<T>((resolve, reject) => {
       const timer = setTimeout(() => {
-        this.destroyRedis();
+        this.destroyRedis(client);
         reject(new Error("Redis operation timed out"));
       }, this.probeTimeoutMs);
       void operation().then(
@@ -104,13 +151,20 @@ export class HealthService implements OnModuleDestroy {
     });
   }
 
-  private destroyRedis(): void {
-    if (this.redisDestroyed) return;
-    this.redisDestroyed = true;
-    try {
-      this.redis.destroy();
-    } catch {
-      // Destruction is best-effort and must never mask readiness/shutdown errors.
+  private destroyRedis(client: RedisClient, generation?: number): void {
+    if (!this.destroyedClients.has(client)) {
+      this.destroyedClients.add(client);
+      try {
+        client.destroy();
+      } catch {
+        // Destruction is best-effort and must never mask readiness/shutdown errors.
+      }
+    }
+    if (
+      this.redis === client &&
+      (generation === undefined || generation === this.generation)
+    ) {
+      this.redis = undefined;
     }
   }
 }