| 1 | import type { ContextManager, Span, SpanOptions, Telemetry, Tracer } from 'bullmq' |
| 2 | import { propagationContext } from '@yikart/common' |
| 3 | import { PinoLogger } from 'nestjs-pino' |
| 4 | import { storage, Store } from 'nestjs-pino/storage' |
| 5 | |
| 6 | interface PinoTelemetryContext { |
| 7 | requestId?: string |
| 8 | } |
| 9 | |
| 10 | class NoopSpan implements Span<PinoTelemetryContext> { |
| 11 | setSpanOnContext(ctx: PinoTelemetryContext): PinoTelemetryContext { |
| 12 | return ctx |
| 13 | } |
| 14 | |
| 15 | setAttribute(): void {} |
| 16 | setAttributes(): void {} |
| 17 | addEvent(): void {} |
| 18 | recordException(): void {} |
| 19 | end(): void {} |
| 20 | } |
| 21 | |
| 22 | class PinoTracer implements Tracer<PinoTelemetryContext> { |
| 23 | startSpan( |
| 24 | _name: string, |
| 25 | _options?: SpanOptions, |
| 26 | _context?: PinoTelemetryContext, |
| 27 | ): Span<PinoTelemetryContext> { |
| 28 | return new NoopSpan() |
| 29 | } |
| 30 | } |
| 31 | |
| 32 | class PinoContextManager implements ContextManager<PinoTelemetryContext> { |
| 33 | with<A extends (...args: unknown[]) => unknown>( |
| 34 | context: PinoTelemetryContext, |
| 35 | fn: A, |
| 36 | ): ReturnType<A> { |
| 37 | const bindings: Record<string, unknown> = {} |
| 38 | if (context.requestId) { |
| 39 | bindings['requestId'] = context.requestId |
| 40 | } |
| 41 | |
| 42 | const logger = PinoLogger.root.child(bindings) |
| 43 | const store = new Store(logger) |
| 44 | return propagationContext.run( |
| 45 | { headers: context.requestId ? { 'x-request-id': context.requestId } : {} }, |
| 46 | () => storage.run(store, fn), |
| 47 | ) as ReturnType<A> |
| 48 | } |
| 49 | |
| 50 | active(): PinoTelemetryContext { |
| 51 | const store = propagationContext.getStore() |
| 52 | const requestId = store?.headers?.['x-request-id'] |
| 53 | return { |
| 54 | requestId: typeof requestId === 'string' ? requestId : undefined, |
| 55 | } |
| 56 | } |
| 57 | |
| 58 | getMetadata(context: PinoTelemetryContext): string { |
| 59 | return JSON.stringify(context) |
| 60 | } |
| 61 | |
| 62 | fromMetadata( |
| 63 | activeContext: PinoTelemetryContext, |
| 64 | metadata: string, |
| 65 | ): PinoTelemetryContext { |
| 66 | try { |
| 67 | const parsed = JSON.parse(metadata) as PinoTelemetryContext |
| 68 | return { ...activeContext, ...parsed } |
| 69 | } |
| 70 | catch { |
| 71 | return activeContext |
| 72 | } |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | export function createPinoTelemetry(): Telemetry<PinoTelemetryContext> { |
| 77 | return { |
| 78 | tracer: new PinoTracer(), |
| 79 | contextManager: new PinoContextManager(), |
| 80 | } |
| 81 | } |
| 82 |