feat: distributed tracking (#24717)

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<string, unknown>;
+  public readonly originalError: unknown;
+
+  constructor(error: unknown, traceContext: TraceContext, additionalData?: Record<string, unknown>) {
+    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<string, unknown> };
+      this.data = {
+        ...(errorWithData.data && typeof errorWithData.data === "object"
+          ? (errorWithData.data as Record<string, unknown>)
+          : {}),
+        ...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<string, unknown>
+  ): TracedError {
+    if (error instanceof TracedError) {
+      return error;
+    }
+    return new TracedError(error, traceContext, additionalData);
+  }
+}

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();

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<unknown>;
+  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<unknown>;
+
+    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<string, string>,
+      };
+
+      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",
+      });
+    });
+  });
+});

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<string, string>;
+}
+
+export interface IdGenerator {
+  generate: () => string;
+}
+
+export class DistributedTracing {
+  constructor(private idGenerator: IdGenerator, private loggerInstance: Logger<unknown>) {}
+
+  createTrace(operation: string, context?: Partial<TraceContext>): 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<string, string>
+  ): 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<string, string>): TraceContext {
+    return {
+      ...traceContext,
+      meta: {
+        ...traceContext.meta,
+        ...additionalMeta,
+      },
+    };
+  }
+}