tracer

index.ts

import { Tracer } from './tracer';
import express, { Request, Response } from 'express';
import { ErrorHandler, Handler, wrap } from 'async-middleware';
import { spanProcessor } from './span-processor';
import { SpanOptions } from '@opentelemetry/api';

async function main() {
  const tracer = new Tracer({ name: 'example-tracer-node' });
  const app = express();

  tracer.provider.addSpanProcessor(spanProcessor);

  app.use('/users/:id', wrapSpan<Request, Response>((req, res) => {
    res.send('Hello Users!');
  }));

  app.use(wrapSpan((req, res) => {
    res.send('Hello World!');
  }));

  app.use(wrapSpan((err, req, res, next) => {
    res.status(500).send('Something broke!');
  }));

  app.listen(3000, () => {
    console.log('Example app listening on port 3000!');
  });

  function wrapSpan<T, U>(fn: Handler<T, U>, options?: Pick<SpanOptions, 'attributes'>): Handler<T, U>;
  function wrapSpan<T, U>(fn: ErrorHandler<T, U>, options?: Pick<SpanOptions, 'attributes'>): ErrorHandler<T, U>;
  function wrapSpan<T, U>(fn: Handler<T, U> | ErrorHandler<T, U>, options?: Pick<SpanOptions, 'attributes'>) {
    const isErrorHandler = (x: any): x is ErrorHandler<T, U> => x.length > 3;

    if (isErrorHandler(fn)) {
      return wrap<T, U>((err, req, res, next) => {
        const context = tracer.extract(req.headers);
        return tracer.startActiveSpan('ErrorHandler', () => {
          return fn(err, req, res, next);
        }, {
          options,
          context
        });
      });
    }

    return wrap<T, U>((req: any, res, next) => {
      const context = tracer.extract(req.headers);
      return tracer.startActiveSpan('Handler', () => {
        return fn(req, res, next);
      }, {
        options,
        context
      });
    });
  }

}

main();

propagator.ts

import {
  Context,
  isSpanContextValid,
  isValidSpanId,
  isValidTraceId,
  trace,
  TextMapGetter,
  TextMapPropagator,
  TextMapSetter,
  TraceFlags,
} from '@opentelemetry/api';
import { isTracingSuppressed } from '@opentelemetry/core';

export const X_TRACE_ID = 'x-traceid';
export const X_SPAN_ID = 'x-spanid';
export const X_SAMPLED = 'x-sampled';
export const X_PARENT_SPAN_ID = 'x-parentspanid';
export const X_FLAGS = 'x-flags';

const VALID_SAMPLED_VALUES = new Set([true, 'true', 'True', '1', 1]);
const VALID_UNSAMPLED_VALUES = new Set([false, 'false', 'False', '0', 0]);

function parseHeader(header: unknown) {
  return Array.isArray(header) ? header[0] : header;
}

function getHeaderValue(carrier: unknown, getter: TextMapGetter, key: string) {
  const header = getter.get(carrier, key);
  return parseHeader(header);
}

function getTraceId(carrier: unknown, getter: TextMapGetter): string {
  const traceId = getHeaderValue(carrier, getter, X_TRACE_ID);
  if (typeof traceId === 'string') {
    return traceId.padStart(32, '0');
  }
  return '';
}

function getSpanId(carrier: unknown, getter: TextMapGetter): string {
  const spanId = getHeaderValue(carrier, getter, X_SPAN_ID);
  if (typeof spanId === 'string') {
    return spanId;
  }
  return '';
}

function getDebug(carrier: unknown, getter: TextMapGetter): string | undefined {
  const debug = getHeaderValue(carrier, getter, X_FLAGS);
  return debug === '1' ? '1' : undefined;
}

function getTraceFlags(
  carrier: unknown,
  getter: TextMapGetter
): TraceFlags | undefined {
  const traceFlags = getHeaderValue(carrier, getter, X_SAMPLED);
  const debug = getDebug(carrier, getter);
  if (debug === '1' || VALID_SAMPLED_VALUES.has(traceFlags)) {
    return TraceFlags.SAMPLED;
  }
  if (VALID_UNSAMPLED_VALUES.has(traceFlags)) {
    return TraceFlags.NONE;
  }
  if (traceFlags === undefined) {
    return TraceFlags.SAMPLED;
  }
  return;
}

export class MultiPropagator implements TextMapPropagator {
  inject(context: Context, carrier: unknown, setter: TextMapSetter): void {
    const spanContext = trace.getSpanContext(context);
    if (
      !spanContext ||
      !isSpanContextValid(spanContext) ||
      isTracingSuppressed(context)
    )
      return;

    setter.set(carrier, X_TRACE_ID, spanContext.traceId);
    setter.set(carrier, X_SPAN_ID, spanContext.spanId);
    if (spanContext.traceFlags !== undefined) {
      setter.set(
        carrier,
        X_SAMPLED,
        (TraceFlags.SAMPLED & spanContext.traceFlags) === TraceFlags.SAMPLED
          ? '1'
          : '0'
      );
    }
  }

  extract(context: Context, carrier: unknown, getter: TextMapGetter): Context {
    const traceId = getTraceId(carrier, getter);
    const spanId = getSpanId(carrier, getter);
    const traceFlags = getTraceFlags(carrier, getter) as TraceFlags;

    if (
      isValidTraceId(traceId) &&
      isValidSpanId(spanId)
    ) {
      return trace.setSpanContext(context, {
        traceId,
        spanId,
        isRemote: true,
        traceFlags,
      });
    }
    return context;
  }

  fields(): string[] {
    return [
      X_TRACE_ID,
      X_SPAN_ID,
      X_FLAGS,
      X_SAMPLED,
      X_PARENT_SPAN_ID,
    ];
  }
}

span-processor.ts

import { ConsoleSpanExporter, SimpleSpanProcessor } from '@opentelemetry/sdk-trace-base';

/**
 * The span processor responsible for processing spans
 * @example provider.addSpanProcessor(spanProcessor)
 */
export const spanProcessor = new SimpleSpanProcessor(new ConsoleSpanExporter())

tracer.ts

import { Span, trace, Tracer as OpenTracer, TextMapPropagator, SpanOptions, Context, defaultTextMapGetter, defaultTextMapSetter } from '@opentelemetry/api';
import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks';
import { BasicTracerProvider } from '@opentelemetry/sdk-trace-base';
import { MultiPropagator } from './propagator';

interface SpanParams {
  options?: SpanOptions;
  context?: Context;
}

interface TraceConstructorParams {
  name: string;
}

export class Tracer {
  public provider: BasicTracerProvider;
  protected contextManager: AsyncLocalStorageContextManager;
  protected propagator: TextMapPropagator;
  protected tracer: OpenTracer;

  constructor({ name }: TraceConstructorParams) {
    this.contextManager = new AsyncLocalStorageContextManager();
    this.propagator = new MultiPropagator();
    this.provider = new BasicTracerProvider();
    this.provider.register({
      contextManager: this.contextManager,
      propagator: this.propagator
    });
    this.tracer = this.provider.getTracer(name);
  }

  /**
   * Wraps a function with a span for tracing purposes.
   * 
   * @param name - The name of the span.
   * @param fn - The function to be wrapped.
   * @returns The wrapped function.
   */
  wrapSpan<F extends (...args: any[]) => any>(name: string, fn: F): F {
    return ((...args: any[]) => {
      return this.startActiveSpan(name, () => {
        return fn(...args);
      });
    }) as F;
  }

  /**
   * Starts an active span with the given name and executes the provided function.
   * @param name - The name of the span.
   * @param fn - The function to be executed with the active span.
   * @returns The return value of the provided function.
   */
  startActiveSpan<F extends (span: Span) => unknown>(name: string, fn: F, params?: SpanParams): ReturnType<F> {
    const spanOptions = params?.options;
    const activeContext = params?.context;

    return this.tracer.startActiveSpan(name, spanOptions as SpanOptions, activeContext as Context, (span: Span) => {
      try {
        const res = fn(span);
        return ((res instanceof Promise) ?
          res.finally(() => {
            return span.end();
          }) :
          (span.end(), res)) as ReturnType<F>;
      } catch (error) {
        span.end();
        throw error;
      }
    });
  }

  extract<Carrier>(input: Carrier) {
    return this.propagator.extract(this.activeContext, input, defaultTextMapGetter);
  }

  inject<Carrier>(output: Carrier) {
    return this.propagator.inject(this.activeContext, output, defaultTextMapSetter);
  }

  get activeSpan() {
    return trace.getSpan(this.activeContext);
  }

  get activeSpanContext() {
    return trace.getSpanContext(this.activeContext);
  }

  get activeContext() {
    return this.contextManager.active();
  }
}