traces.ts

/* eslint-disable @typescript-eslint/no-unsafe-argument */
/* eslint-disable @typescript-eslint/no-unsafe-member-access */
/* eslint-disable @typescript-eslint/no-unsafe-call */
/* eslint-disable @typescript-eslint/no-unsafe-return */
/* eslint-disable @typescript-eslint/no-unsafe-assignment */
/* eslint-disable @typescript-eslint/no-require-imports */
import type { Span } from '@opentelemetry/sdk-trace-base';
import type { SpanContext } from '@opentelemetry/api';
import { TelemetryCollector } from './collector';
import { globalParams } from '../utils';
import type { BasicTracerProvider as BasicTracerProviderType } from '@opentelemetry/sdk-trace-base';

// As DBOS OTLP is optional, OTLP objects must only be dynamically imported
// and only when OTLP is enabled. Importing OTLP types is fine as long
// as signatures using those types are not exported from this file.

interface Attributes {
  [attributeKey: string]: AttributeValue | undefined;
}
/**
 * Attribute values may be any non-nullish primitive value except an object.
 *
 * null or undefined attribute values are invalid and will result in undefined behavior.
 */
declare type AttributeValue =
  | string
  | number
  | boolean
  | Array<null | undefined | string>
  | Array<null | undefined | number>
  | Array<null | undefined | boolean>;

export enum SpanStatusCode {
  /**
   * The default status.
   */
  UNSET = 0,
  /**
   * The operation has been validated by an Application developer or
   * Operator to have completed successfully.
   */
  OK = 1,
  /**
   * The operation contains an error.
   */
  ERROR = 2,
}

interface SpanStatus {
  /** The status code of this message. */
  code: SpanStatusCode;
  /** A developer-facing error message. */
  message?: string;
}

export type DBOSSpan = {
  setStatus(status: SpanStatus): DBOSSpan;
  attributes: Attributes;
  setAttribute(key: string, attribute: AttributeValue): DBOSSpan;
  addEvent(name: string, attributesOrStartTime?: Attributes, timeStamp?: number): DBOSSpan;
};

class StubSpan implements DBOSSpan {
  attributes: Attributes = {};

  setStatus(_status: SpanStatus): DBOSSpan {
    return this;
  }

  setAttribute(_key: string, _attribute: AttributeValue): DBOSSpan {
    return this;
  }

  addEvent(_name: string, _attributesOrStartTime?: Attributes, _timeStamp?: number): DBOSSpan {
    return this;
  }
}

export function runWithTrace<R>(span: DBOSSpan, func: () => Promise<R>): Promise<R> {
  if (!globalParams.tracingEnabled) {
    return func();
  }
  const { context, trace } = require('@opentelemetry/api');
  return context.with(trace.setSpan(context.active(), span as Span), func);
}

export function getActiveSpan() {
  if (!globalParams.tracingEnabled) {
    return undefined;
  }
  const { trace } = require('@opentelemetry/api');
  return trace.getActiveSpan() as DBOSSpan | undefined;
}

export function isTraceContextWorking(): boolean {
  if (!globalParams.tracingEnabled) {
    return false;
  }
  const { context, trace } = require('@opentelemetry/api');
  const span = trace.getTracer('otel-bootstrap-check').startSpan('probe');
  const testContext = trace.setSpan(context.active(), span);

  let visible: boolean | undefined;
  context.with(testContext, () => {
    visible = trace.getSpan(context.active()) === span;
  });

  span.end?.();
  return visible === true;
}

export function installTraceContextManager(appName: string = 'dbos'): void {
  if (!globalParams.tracingEnabled) {
    return;
  }
  const { AsyncLocalStorageContextManager } = require('@opentelemetry/context-async-hooks');
  const { context, trace } = require('@opentelemetry/api');
  const { BasicTracerProvider } = require('@opentelemetry/sdk-trace-base');

  // setGlobalTracerProvider and setGlobalContextManager are "first one wins."
  // If an external provider is already registered, these calls are safely ignored.
  const contextManager = new AsyncLocalStorageContextManager();
  contextManager.enable();
  context.setGlobalContextManager(contextManager);
  const provider: BasicTracerProviderType = new BasicTracerProvider({
    resource: {
      attributes: {
        'service.name': appName,
      },
    },
  });
  trace.setGlobalTracerProvider(provider);
}

export class Tracer {
  readonly applicationID: string;
  readonly executorID: string;
  constructor(private readonly telemetryCollector: TelemetryCollector) {
    this.applicationID = globalParams.appID;
    this.executorID = globalParams.executorID;
  }

  startSpanWithContext(spanContext: unknown, name: string, attributes?: Attributes): DBOSSpan {
    if (!globalParams.tracingEnabled) {
      return new StubSpan();
    }
    const opentelemetry = require('@opentelemetry/api');
    const tracer = opentelemetry.trace.getTracer('dbos-tracer');
    const ctx = opentelemetry.trace.setSpanContext(opentelemetry.context.active(), spanContext as SpanContext);
    return tracer.startSpan(name, { startTime: performance.now(), attributes: attributes }, ctx) as Span;
  }

  startSpan(name: string, attributes?: Attributes, inputSpan?: DBOSSpan): DBOSSpan {
    if (!globalParams.tracingEnabled) {
      return new StubSpan();
    }
    const parentSpan = inputSpan as Span;
    const opentelemetry = require('@opentelemetry/api');
    const { hrTime } = require('@opentelemetry/core');
    const tracer = opentelemetry.trace.getTracer('dbos-tracer');
    const startTime = hrTime(performance.now());
    if (parentSpan) {
      const ctx = opentelemetry.trace.setSpan(opentelemetry.context.active(), parentSpan);
      return tracer.startSpan(name, { startTime: startTime, attributes: attributes }, ctx) as Span;
    } else {
      return tracer.startSpan(name, { startTime: startTime, attributes: attributes }) as Span;
    }
  }

  endSpan(inputSpan: DBOSSpan) {
    if (!globalParams.tracingEnabled) {
      return;
    }
    const { hrTime } = require('@opentelemetry/core');
    const span = inputSpan as Span;
    span.setAttributes({
      applicationID: this.applicationID,
      applicationVersion: globalParams.appVersion,
    });
    if (span.attributes && !('executorID' in span.attributes)) {
      span.setAttribute('executorID', this.executorID);
    }
    span.end(hrTime(performance.now()));
    // Only push to DBOS's own collector when DBOS manages export.
    // When an external TracerProvider is used, span.end() triggers its processors.
    if (globalParams.enableOTLP) {
      this.telemetryCollector.push(span);
    }
  }
}