From 28693aba965e602e12fa14a4a5f3ba3648fb64dd Mon Sep 17 00:00:00 2001 From: Udit Takkar <53316345+Udit-takkar@users.noreply.github.com> Date: Thu, 30 Oct 2025 19:42:22 +0530 Subject: [PATCH] feat: distributed tracking (#24717) --- packages/lib/tracing/error.ts | 53 ++++ packages/lib/tracing/factory.ts | 14 + packages/lib/tracing/index.test.ts | 463 +++++++++++++++++++++++++++++ packages/lib/tracing/index.ts | 97 ++++++ 4 files changed, 627 insertions(+) create mode 100644 packages/lib/tracing/error.ts create mode 100644 packages/lib/tracing/factory.ts create mode 100644 packages/lib/tracing/index.test.ts create mode 100644 packages/lib/tracing/index.ts diff --git a/packages/lib/tracing/error.ts b/packages/lib/tracing/error.ts new file mode 100644 index 0000000000..d89a2af05f --- /dev/null +++ b/packages/lib/tracing/error.ts @@ -0,0 +1,53 @@ +import type { TraceContext } from "./index"; + +export class TracedError extends Error { + public readonly traceId: string; + public readonly data?: Record; + public readonly originalError: unknown; + + constructor(error: unknown, traceContext: TraceContext, additionalData?: Record) { + const message = error instanceof Error ? error.message : String(error); + super(message); + + this.name = error instanceof Error ? error.name : "TracedError"; + this.traceId = traceContext.traceId; + this.originalError = error; + + if (error instanceof Error && error.stack) { + this.stack = error.stack; + } + + if (error && typeof error === "object" && "data" in error) { + const errorWithData = error as { data: Record }; + this.data = { + ...(errorWithData.data && typeof errorWithData.data === "object" + ? (errorWithData.data as Record) + : {}), + ...additionalData, + }; + } else if (additionalData) { + this.data = additionalData; + } + } + + toJSON() { + return { + name: this.name, + message: this.message, + traceId: this.traceId, + data: this.data, + stack: this.stack, + }; + } + + static createFromError( + error: unknown, + traceContext: TraceContext, + additionalData?: Record + ): TracedError { + if (error instanceof TracedError) { + return error; + } + return new TracedError(error, traceContext, additionalData); + } +} diff --git a/packages/lib/tracing/factory.ts b/packages/lib/tracing/factory.ts new file mode 100644 index 0000000000..5ed60f27e5 --- /dev/null +++ b/packages/lib/tracing/factory.ts @@ -0,0 +1,14 @@ +import { nanoid } from "nanoid"; + +import logger from "../logger"; +import { DistributedTracing, type IdGenerator } from "./index"; + +const idGenerator: IdGenerator = { + generate: () => nanoid(), +}; + +export function createDistributedTracing(): DistributedTracing { + return new DistributedTracing(idGenerator, logger); +} + +export const distributedTracing = createDistributedTracing(); diff --git a/packages/lib/tracing/index.test.ts b/packages/lib/tracing/index.test.ts new file mode 100644 index 0000000000..16c097622f --- /dev/null +++ b/packages/lib/tracing/index.test.ts @@ -0,0 +1,463 @@ +import { describe, expect, it, beforeEach, vi } from "vitest"; +import type { Logger } from "tslog"; + +import { DistributedTracing } from "./index"; +import type { IdGenerator } from "./index"; + +describe("DistributedTracing", () => { + let idCounter: number; + let mockIdGenerator: IdGenerator; + let mockLogger: Logger; + let tracing: DistributedTracing; + + beforeEach(() => { + idCounter = 0; + mockIdGenerator = { + generate: () => { + idCounter++; + return `id_${idCounter}`; + }, + }; + + mockLogger = { + getSubLogger: vi.fn().mockReturnValue({ + info: vi.fn(), + error: vi.fn(), + warn: vi.fn(), + debug: vi.fn(), + }), + } as unknown as Logger; + + tracing = new DistributedTracing(mockIdGenerator, mockLogger); + }); + + describe("createTrace", () => { + it("should create a new trace with generated IDs when no context is provided", () => { + const result = tracing.createTrace("test_operation"); + + expect(result).toEqual({ + traceId: "trace_id_1", + spanId: "span_id_2", + parentSpanId: undefined, + operation: "test_operation", + meta: undefined, + }); + }); + + it("should preserve existing traceId when provided in context", () => { + const result = tracing.createTrace("test_operation", { + traceId: "existing_trace_123", + }); + + expect(result.traceId).toBe("existing_trace_123"); + expect(result.spanId).toBe("span_id_1"); + expect(result.operation).toBe("test_operation"); + }); + + it("should set parentSpanId to incoming spanId when provided", () => { + const result = tracing.createTrace("child_operation", { + spanId: "parent_span_456", + }); + + expect(result.parentSpanId).toBe("parent_span_456"); + expect(result.spanId).toMatch(/^span_id_\d+$/); + expect(result.spanId).not.toBe(result.parentSpanId); + }); + + it("should preserve meta when provided in context", () => { + const meta = { + userId: "123", + requestId: "req_456", + }; + + const result = tracing.createTrace("test_operation", { meta }); + + expect(result.meta).toEqual(meta); + }); + + it("should NOT allow context to override the new operation", () => { + const contextWithOperation = { + operation: "old_operation", + }; + const result = tracing.createTrace("new_operation", contextWithOperation); + + expect(result.operation).toBe("new_operation"); + }); + + it("should NOT allow context to override the new parentSpanId", () => { + const contextWithParentSpan = { + spanId: "correct_parent_span", + parentSpanId: "incorrect_parent_span", + }; + const result = tracing.createTrace("test_operation", contextWithParentSpan); + + expect(result.parentSpanId).toBe("correct_parent_span"); + }); + + it("should handle complex context with multiple properties", () => { + const result = tracing.createTrace("complex_operation", { + traceId: "trace_999", + spanId: "span_888", + meta: { + env: "production", + version: "1.2.3", + }, + }); + + expect(result).toEqual({ + traceId: "trace_999", + spanId: "span_id_1", + parentSpanId: "span_888", + operation: "complex_operation", + meta: { + env: "production", + version: "1.2.3", + }, + }); + }); + + it("should generate new traceId when traceId is explicitly undefined", () => { + const result = tracing.createTrace("test_operation", { + traceId: undefined, + }); + + expect(result.traceId).toBe("trace_id_1"); + }); + + it("should handle empty context object", () => { + const result = tracing.createTrace("test_operation", {}); + + expect(result).toEqual({ + traceId: "trace_id_1", + spanId: "span_id_2", + parentSpanId: undefined, + operation: "test_operation", + meta: undefined, + }); + }); + }); + + describe("createSpan", () => { + it("should create a child span from parent context", () => { + const parentContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "parent_operation", + meta: { + userId: "user_789", + }, + }; + + const result = tracing.createSpan(parentContext, "child_operation"); + + expect(result).toEqual({ + traceId: "trace_123", + spanId: "span_id_1", + parentSpanId: "span_456", + operation: "child_operation", + meta: { + userId: "user_789", + }, + }); + }); + + it("should merge additional meta with parent meta", () => { + const parentContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "parent_operation", + meta: { + userId: "user_789", + }, + }; + + const result = tracing.createSpan(parentContext, "child_operation", { + requestId: "req_999", + }); + + expect(result.meta).toEqual({ + userId: "user_789", + requestId: "req_999", + }); + }); + + it("should allow additional meta to override parent meta keys", () => { + const parentContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "parent_operation", + meta: { + userId: "old_user", + env: "staging", + }, + }; + + const result = tracing.createSpan(parentContext, "child_operation", { + userId: "new_user", + }); + + expect(result.meta).toEqual({ + userId: "new_user", + env: "staging", + }); + }); + + it("should handle parent context without meta", () => { + const parentContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "parent_operation", + }; + + const result = tracing.createSpan(parentContext, "child_operation"); + + expect(result.meta).toEqual({}); + }); + }); + + describe("getTracingLogger", () => { + it("should create logger with correct prefixes for basic context", () => { + const context = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + }; + + tracing.getTracingLogger(context); + + expect(mockLogger.getSubLogger).toHaveBeenCalledWith({ + prefix: ["distributed-trace", "trace:trace_123", "span:span_456", "op:test_operation"], + }); + }); + + it("should include meta in prefixes when provided", () => { + const context = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "user_789", + requestId: "req_999", + }, + }; + + tracing.getTracingLogger(context); + + expect(mockLogger.getSubLogger).toHaveBeenCalledWith({ + prefix: [ + "distributed-trace", + "trace:trace_123", + "span:span_456", + "op:test_operation", + "userId:user_789", + "requestId:req_999", + ], + }); + }); + + it("should skip undefined meta values", () => { + const context = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "user_789", + requestId: undefined, + }, + }; + + tracing.getTracingLogger(context); + + expect(mockLogger.getSubLogger).toHaveBeenCalledWith({ + prefix: ["distributed-trace", "trace:trace_123", "span:span_456", "op:test_operation", "userId:user_789"], + }); + }); + + it("should skip null meta values", () => { + const context = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "user_789", + requestId: null, + } as Record, + }; + + tracing.getTracingLogger(context); + + expect(mockLogger.getSubLogger).toHaveBeenCalledWith({ + prefix: ["distributed-trace", "trace:trace_123", "span:span_456", "op:test_operation", "userId:user_789"], + }); + }); + }); + + describe("extractTraceFromPayload", () => { + it("should extract trace context from valid JSON payload", () => { + const traceContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + }; + + const payload = JSON.stringify({ + _traceContext: traceContext, + otherData: "value", + }); + + const result = tracing.extractTraceFromPayload(payload); + + expect(result).toEqual(traceContext); + }); + + it("should return null for payload without trace context", () => { + const payload = JSON.stringify({ + otherData: "value", + }); + + const result = tracing.extractTraceFromPayload(payload); + + expect(result).toBeNull(); + }); + + it("should return null for invalid JSON", () => { + const payload = "invalid json {"; + + const result = tracing.extractTraceFromPayload(payload); + + expect(result).toBeNull(); + }); + + it("should return null for payload with _traceContext but no traceId", () => { + const payload = JSON.stringify({ + _traceContext: { + spanId: "span_456", + }, + }); + + const result = tracing.extractTraceFromPayload(payload); + + expect(result).toBeNull(); + }); + }); + + describe("injectTraceIntoPayload", () => { + it("should inject trace context into payload", () => { + const payload = { + data: "value", + }; + + const traceContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + }; + + const result = tracing.injectTraceIntoPayload(payload, traceContext); + + expect(result).toEqual({ + data: "value", + _traceContext: traceContext, + }); + }); + + it("should return original payload when no trace context provided", () => { + const payload = { + data: "value", + }; + + const result = tracing.injectTraceIntoPayload(payload); + + expect(result).toEqual(payload); + }); + + it("should return original payload when trace context is undefined", () => { + const payload = { + data: "value", + }; + + const result = tracing.injectTraceIntoPayload(payload, undefined); + + expect(result).toEqual(payload); + }); + }); + + describe("updateTrace", () => { + it("should update trace context with additional meta", () => { + const traceContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "user_789", + }, + }; + + const result = tracing.updateTrace(traceContext, { + requestId: "req_999", + }); + + expect(result).toEqual({ + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "user_789", + requestId: "req_999", + }, + }); + }); + + it("should override existing meta keys", () => { + const traceContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "old_user", + }, + }; + + const result = tracing.updateTrace(traceContext, { + userId: "new_user", + }); + + expect(result.meta).toEqual({ + userId: "new_user", + }); + }); + + it("should handle trace context without meta", () => { + const traceContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + }; + + const result = tracing.updateTrace(traceContext, { + userId: "user_789", + }); + + expect(result.meta).toEqual({ + userId: "user_789", + }); + }); + + it("should handle undefined additional meta", () => { + const traceContext = { + traceId: "trace_123", + spanId: "span_456", + operation: "test_operation", + meta: { + userId: "user_789", + }, + }; + + const result = tracing.updateTrace(traceContext); + + expect(result.meta).toEqual({ + userId: "user_789", + }); + }); + }); +}); diff --git a/packages/lib/tracing/index.ts b/packages/lib/tracing/index.ts new file mode 100644 index 0000000000..73750e72e3 --- /dev/null +++ b/packages/lib/tracing/index.ts @@ -0,0 +1,97 @@ +import type { Logger } from "tslog"; + +export interface TraceContext { + traceId: string; + spanId: string; + parentSpanId?: string; + operation: string; + // So that we don't violate open closed principle, we allow meta to be added to the trace context + meta?: Record; +} + +export interface IdGenerator { + generate: () => string; +} + +export class DistributedTracing { + constructor(private idGenerator: IdGenerator, private loggerInstance: Logger) {} + + createTrace(operation: string, context?: Partial): TraceContext { + const { traceId, spanId, meta, ...otherContext } = context || {}; + return { + ...otherContext, + traceId: traceId || `trace_${this.idGenerator.generate()}`, + spanId: `span_${this.idGenerator.generate()}`, + parentSpanId: spanId, + operation, + meta, + }; + } + + createSpan( + parentContext: TraceContext, + operation: string, + additionalMeta?: Record + ): TraceContext { + const mergedMeta = { + ...parentContext.meta, + ...additionalMeta, + }; + + return { + ...parentContext, + spanId: `span_${this.idGenerator.generate()}`, + parentSpanId: parentContext.spanId, + operation, + meta: mergedMeta, + }; + } + + getTracingLogger(context: TraceContext) { + const prefixes = [ + "distributed-trace", + `trace:${context.traceId}`, + `span:${context.spanId}`, + `op:${context.operation}`, + ]; + + if (context.meta) { + Object.entries(context.meta).forEach(([key, value]) => { + if (value !== undefined && value !== null) { + prefixes.push(`${key}:${value}`); + } + }); + } + + return this.loggerInstance.getSubLogger({ prefix: prefixes }); + } + + extractTraceFromPayload(payload: string): TraceContext | null { + try { + const parsed = JSON.parse(payload); + if (parsed._traceContext && parsed._traceContext.traceId) { + return parsed._traceContext; + } + } catch {} + return null; + } + + injectTraceIntoPayload(payload: any, traceContext?: TraceContext): any { + if (!traceContext) return payload; + + return { + ...payload, + _traceContext: traceContext, + }; + } + + updateTrace(traceContext: TraceContext, additionalMeta?: Record): TraceContext { + return { + ...traceContext, + meta: { + ...traceContext.meta, + ...additionalMeta, + }, + }; + } +}