Elysia + OpenTelemetry `W3CTraceContextPropagator` Propagation

telemetry.ts

import { contextKeySpan, opentelemetry } from "@elysiajs/opentelemetry";
import {
  context as otelContext,
  propagation,
  trace,
  type Span,
} from "@opentelemetry/api";
import { W3CTraceContextPropagator } from "@opentelemetry/core";
import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-proto";
import { BatchSpanProcessor } from "@opentelemetry/sdk-trace-node";
import { PrismaInstrumentation } from "@prisma/instrumentation";
import Elysia from "elysia";

// new Elysia().use(telemetry('foo'));

export function telemetry(serviceName: string) {
  return new Elysia()
    .trace({ as: "global" }, ({ context: { request }, onAfterResponse }) => {
      const rootSpan = trace.getActiveSpan();
      if (!rootSpan) return;

      let parent = rootSpan;

      // This is quite hacky, as it touches the opentelemetry internals,
      // however it seems to be the only way to glue together the parent
      // span in a way that works well with the Elysia async behavior.
      // In fact this is lifted straight from the Elysia source code, see:
      // https://github.com/elysiajs/opentelemetry/blob/main/src/index.ts#L302-L315
      function setParent(span: Span) {
        const newContext = trace.setSpan(otelContext.active(), span);

        const currentContext: Map<Symbol, unknown> =
          // @ts-expect-error private property
          otelContext.active()._currentContext;

        currentContext?.set(
          contextKeySpan,
          newContext.getValue(contextKeySpan)
        );

        parent = span;
      }

      const extractedContext = propagation.extract(
        otelContext.active(),
        request.headers,
        {
          get(carrier, key) {
            return carrier.get(key) ?? undefined;
          },
          keys(carrier) {
            return [...carrier.keys()];
          },
        }
      );

      const tracer = trace.getTracer(serviceName);
      const span = tracer.startSpan("lifecycle", {}, extractedContext);
      setParent(span);

      onAfterResponse(() => {
        setImmediate(() => {
          span.end();
        });
      });
    })
    .use(
      opentelemetry({
        serviceName,
        textMapPropagator: new W3CTraceContextPropagator(),
        instrumentations: [new PrismaInstrumentation()],
        spanProcessors: [
          // @ts-ignore - The types are wrong
          new BatchSpanProcessor(
            new OTLPTraceExporter({
              url: process.env.OTLP_TRACE_URL,
            })
          ),
        ],
      })
    );
}

let originalFetch = globalThis.fetch;
globalThis.fetch = (...args) => {
  const headers = args[1]?.headers;

  if (headers) {
    propagation.inject(otelContext.active(), headers, {
      set(carrier, key, value) {
        if (carrier instanceof Headers) {
          carrier.set(key, value);
        } else {
          (carrier as any)[key] = value;
        }
      },
    });
  }

  return originalFetch(...args);
};