tracing.ts

import { Link, Span, SpanStatusCode, trace } from '@opentelemetry/api';
import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node';
import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks';
import {
  CompositePropagator,
  W3CBaggagePropagator,
  W3CTraceContextPropagator,
} from '@opentelemetry/core';
import { B3InjectEncoding, B3Propagator } from '@opentelemetry/propagator-b3';
import { NodeSDK, NodeSDKConfiguration } from '@opentelemetry/sdk-node';
import { addCleanupListener, exitAfterCleanup } from 'async-cleanup';
import { merge } from 'lodash';

const defaultOtelConfig: Partial<NodeSDKConfiguration> = {
  contextManager: new AsyncLocalStorageContextManager(),
  textMapPropagator: new CompositePropagator({
    propagators: [
      new W3CTraceContextPropagator(),
      new W3CBaggagePropagator(),
      new B3Propagator(),
      new B3Propagator({
        injectEncoding: B3InjectEncoding.MULTI_HEADER,
      }),
    ],
  }),
  // REVIEW: should we auto instrument just everything?
  instrumentations: [getNodeAutoInstrumentations()],
};

/**
 * Call this method as the very first before everything else.
 *
 * @param otelConfig Partial configuration to merge onto defaults
 */
export function instrument(otelConfig?: Partial<NodeSDKConfiguration>) {
  const mergedConfigs = merge({}, defaultOtelConfig, otelConfig);
  const otelSDK = new NodeSDK(mergedConfigs);

  const cleanup = async () => {
    try {
      // NOTE: hack to flush data upstream - https://github.com/open-telemetry/opentelemetry-js/issues/3310
      const flushes = [
        mergedConfigs.traceExporter?.forceFlush?.(),
        mergedConfigs.metricReader?.forceFlush?.(),
        ...(mergedConfigs.spanProcessors?.map((p) => p.forceFlush()) ?? []),
        ...(mergedConfigs.logRecordProcessors?.map((p) => p.forceFlush()) ?? []),
      ];
      await Promise.allSettled(flushes);

      await otelSDK.shutdown();
      // console.info('OTEL SDK shut down successfully');
      // eslint-disable-next-line @typescript-eslint/no-unused-vars
    } catch (error) {
      // REVIEW: Nothing really we can do at this point since we are closing the process,
      // and often this error is for unavailablity of the upstream OTEL collectors - unrecoverable
      // console.error('OTEL SDK encountered an error while shutting down', error);
    }
  };
  addCleanupListener(cleanup);

  otelSDK.start();
}

const handlerFactory = (name: string, parentSpan: Span, error: unknown) => {
  const parentLink: Link = {
    context: parentSpan.spanContext(),
  };

  trace.getTracer('default').startActiveSpan(name, { links: [parentLink] }, (span) => {
    if (error instanceof Error) span.recordException(error);
    const message = error instanceof Error ? error.message : new String(error).toString();
    span.setStatus({
      code: SpanStatusCode.ERROR,
      message,
    });
    span.end();
  });
};

/**
 * Run the bootstrap function instrumented with OpenTelemetry.
 * Automatically handles all uncaught exceptions and rejections as well.
 *
 * @param name Root trace name
 * @param bootstrap Main bootstrap function to run
 */
export async function runInstrumented(
  bootstrap: () => PromiseLike<void>,
  name: string = 'MainBootstrap',
): Promise<void> {
  await trace.getTracer('default').startActiveSpan(name, { root: true }, async (span) => {
    const handler = (err: unknown) => handlerFactory(name, span, err);
    process.addListener('unhandledRejection', handler);
    process.addListener('uncaughtException', handler);
    const cleanup = () => {
      process.removeListener('unhandledRejection', handler);
      process.removeListener('uncaughtException', handler);
    };
    addCleanupListener(cleanup);

    try {
      await bootstrap();
    } catch (error) {
      const message = error instanceof Error ? error.message : new String(error).toString();
      span.setStatus({
        code: SpanStatusCode.ERROR,
        message,
      });
      await exitAfterCleanup();
    } finally {
      span.end();
    }
  });
}