NestJS 中优雅实现 RequestID 和 TaskID 追踪

背景介绍

• 在微服务架构中，请求追踪是一个重要的需求
• 需要在整个请求链路中（包括 HTTP 请求、数据库操作等）保持统一的追踪标识
• 本文将介绍如何在 NestJS 应用的各个组件中实现 RequestID 和 TaskID 的追踪机制

技术方案

1. 安装依赖

npm install nestjs-cls uuid @nestjs/axios @prisma/client

2. 配置 ClsModule

• 在 app.module.ts 中配置:

import { ClsModule } from 'nestjs-cls';

@Module({
  imports: [
    ClsModule.forRoot({
      global: true,
      middleware: { mount: true }
    }),
  ],
})
export class AppModule {}

3. 在 Axios 服务中集成追踪 ID

• 创建 axios.service.ts:

@Injectable()
export class AxiosService {
  constructor(
    private readonly httpService: HttpService,
    private readonly cls: ClsService,
  ) {
    this.httpService.axiosRef.interceptors.request.use(
      (config) => {
        const requestId = this.cls.getId() as string;
        const taskId = this.cls.get<string>('taskId');
        const traceId = taskId || requestId;
        
        config.headers = new AxiosHeaders({
          ...DEFAULT_HEADERS,
          'X-Trace-ID': traceId,
          ...config.headers,
        });
        
        return config;
      },
    );
  }
}

4. 在 Prisma 服务中集成追踪 ID

• 创建 prisma.service.ts:

export const prismaExtendedClient = (
  prismaClient: PrismaClient,
  cls: ClsService,
  prismaService: PrismaService,
) =>
  prismaClient.$extends({
    query: {
      $allModels: {
        async $allOperations({ operation, args, query }) {
          const result = await query(args);

          const requestId = cls.getId() as string;
          const taskId = cls.get<string>('taskId');
          const traceId = taskId || requestId;

          prismaService.logger.debug(
            `[TraceID: ${traceId}] SQL Query:\n${prismaService.queryInfo}\n`,
          );
          prismaService.queryInfo = '';
          return result;
        },
      },
    },
  });

@Injectable()
export class PrismaService extends PrismaClient {
  constructor(private readonly cls: ClsService) {
    super({
      log: [{ emit: 'event', level: 'query' }],
    });
    // ... 其他初始化代码
  }
}

5. 定时任务中的 TaskID 装饰器

export function WithTaskId() {
  return function (
    target: any,
    propertyKey: string,
    descriptor: PropertyDescriptor,
  ) {
    const originalMethod = descriptor.value;

    descriptor.value = async function (...args: any[]) {
      const cls = this.cls as ClsService;
      const taskId = `task-${Date.now()}-${uuidv4().slice(0, 8)}`;

      return cls.run(async () => {
        cls.set('taskId', taskId);
        return originalMethod.apply(this, args);
      });
    };

    return descriptor;
  };
}

实现效果

HTTP 请求链路追踪

• 入站请求自动获得 RequestID
• 出站 HTTP 请求携带相同的 ID
• 数据库操作日志包含相同的 ID

定时任务链路追踪

• 定时任务自动生成 TaskID
• TaskID 在整个任务执行过程中保持不变
• 包括 HTTP 请求和数据库操作在内的所有操作都使用相同的 TaskID

日志输出示例

[Nest] 28889  - 2024/03/21 14:30:45  [HTTP] Incoming request - TraceID: req-1234-5678
[Nest] 28889  - 2024/03/21 14:30:45  [Prisma] SQL Query - TraceID: req-1234-5678
[Nest] 28889  - 2024/03/21 14:30:46  [Axios] Outgoing request - TraceID: req-1234-5678

[Nest] 28889  - 2024/03/21 14:31:00  [Schedule] Starting task - TraceID: task-1234-5678
[Nest] 28889  - 2024/03/21 14:31:00  [Prisma] SQL Query - TraceID: task-1234-5678
[Nest] 28889  - 2024/03/21 14:31:01  [Axios] Outgoing request - TraceID: task-1234-5678

最佳实践

• 类型安全

  • 使用泛型类型确保 ID 的类型安全
  • 避免使用 any 类型
  • 使用类型断言明确类型

• 日志格式统一

  • 使用统一的 TraceID 标识符
  • 保持日志格式的一致性
  • 包含必要的上下文信息

• 错误处理

  • 在所有错误日志中包含追踪 ID
  • 确保异常情况下 ID 传递不中断

总结

• 通过在各个组件中集成 RequestID 和 TaskID：
  • 实现了完整的请求链路追踪
  • 提高了系统的可观测性
  • 便于问题定位和性能分析
• 统一的追踪机制让系统维护更加高效
• 类型安全的实现确保了代码的可靠性