In distributed systems, tracking requests across multiple services can be challenging. To address this, we use TraceIDs, unique identifiers attached to each request, allowing us to trace the request flow through different services. In this article, we'll explore how to implement TraceID in a NestJS application, focusing on both HTTP requests and SQS consumers.
Why TraceID?
TraceIDs help in:
Debugging: Identifying issues across services by following a single request.
Monitoring: Tracking performance bottlenecks and latency issues.
Logging: Correlating logs from different services to get a complete picture of request processing.
Setting Up TraceID in NestJS
Prerequisites
Ensure you have a NestJS project set up. If not, create one using the NestJS CLI:
npm i -g @nestjs/cli
nest new traceid-demo
Adding TraceID to HTTP Requests
- Install Dependencies
We'll use the cls-hooked
library to manage context propagation for TraceIDs.
npm install cls-hooked
- Create a TraceID Middleware
Create a middleware to generate or extract a TraceID from incoming requests.
import { Injectable, NestMiddleware } from '@nestjs/common';
import { v4 as uuidv4 } from 'uuid';
import { Request, Response, NextFunction } from 'express';
import { createNamespace } from 'cls-hooked';
const traceIdNamespace = createNamespace('traceIdNamespace');
@Injectable()
export class TraceIdMiddleware implements NestMiddleware {
use(req: Request, res: Response, next: NextFunction) {
traceIdNamespace.run(() => {
const traceId = req.headers['x-trace-id'] || uuidv4();
traceIdNamespace.set('traceId', traceId);
res.setHeader('x-trace-id', traceId);
next();
});
}
}
- Apply Middleware
Apply the middleware in the main application module.
import { Module, NestModule, MiddlewareConsumer } from '@nestjs/common';
import { TraceIdMiddleware } from './trace-id.middleware';
@Module({
imports: [],
controllers: [],
providers: [],
})
export class AppModule implements NestModule {
configure(consumer: MiddlewareConsumer) {
consumer.apply(TraceIdMiddleware).forRoutes('*');
}
}
- Access TraceID in Services and Controllers
Access the TraceID using cls-hooked
wherever needed.
import { Injectable } from '@nestjs/common';
import { getNamespace } from 'cls-hooked';
@Injectable()
export class AppService {
getHello(): string {
const traceIdNamespace = getNamespace('traceIdNamespace');
const traceId = traceIdNamespace ? traceIdNamespace.get('traceId') : 'unknown';
console.log(`TraceID: ${traceId}`);
return 'Hello World!';
}
}
Adding TraceID to SQS Consumers
- Setup SQS Consumer
Assuming you have an SQS consumer setup using NestJS.
import { Injectable } from '@nestjs/common';
import { SqsMessageHandler, SqsMessageHandlerMeta } from '@nestjs-packages/sqs';
@Injectable()
export class SqsConsumerService {
@SqsMessageHandler({ batch: false })
async handleMessage(message: AWS.SQS.Message) {
// Process the message
console.log(message);
}
}
- Extract TraceID from SQS Message
Modify the SQS consumer to extract and set the TraceID.
import { Injectable } from '@nestjs/common';
import { SqsMessageHandler, SqsMessageHandlerMeta } from '@nestjs-packages/sqs';
import { getNamespace, createNamespace } from 'cls-hooked';
const traceIdNamespace = createNamespace('traceIdNamespace');
@Injectable()
export class SqsConsumerService {
@SqsMessageHandler({ batch: false })
async handleMessage(message: AWS.SQS.Message) {
traceIdNamespace.run(() => {
const traceId = message.MessageAttributes?.traceId?.StringValue || uuidv4();
traceIdNamespace.set('traceId', traceId);
// Process the message with TraceID
console.log(`TraceID: ${traceId}`, message);
});
}
}
Sending Messages with TraceID
When sending messages to SQS, ensure to include the TraceID.
import { SQS } from 'aws-sdk';
import { getNamespace } from 'cls-hooked';
const sqs = new SQS();
async function sendMessageToSqs(messageBody: string) {
const traceIdNamespace = getNamespace('traceIdNamespace');
const traceId = traceIdNamespace ? traceIdNamespace.get('traceId') : uuidv4();
const params = {
MessageBody: messageBody,
QueueUrl: 'YOUR_SQS_QUEUE_URL',
MessageAttributes: {
traceId: {
DataType: 'String',
StringValue: traceId,
},
},
};
await sqs.sendMessage(params).promise();
}
Conclusion
Implementing TraceID in your NestJS application helps in better tracking and debugging of requests across distributed systems. By integrating TraceIDs into both HTTP requests and SQS consumers, you can ensure comprehensive traceability and monitoring. This setup improves your ability to diagnose issues, enhance performance, and maintain a robust logging system.