elasticsearch/ts/core/observability/tracing.ts

/**
 * Span attributes
 */
export type SpanAttributes = Record<string, string | number | boolean | null | undefined>;

/**
 * Span status
 */
export enum SpanStatus {
  OK = 'OK',
  ERROR = 'ERROR',
  UNSET = 'UNSET',
}

/**
 * Span context for distributed tracing
 */
export interface SpanContext {
  traceId: string;
  spanId: string;
  traceFlags: number;
  traceState?: string;
}

/**
 * Span interface
 */
export interface Span {
  /** Span name */
  name: string;

  /** Span start time */
  startTime: Date;

  /** Span end time (if ended) */
  endTime?: Date;

  /** Span status */
  status: SpanStatus;

  /** Span attributes */
  attributes: SpanAttributes;

  /** Span context */
  context: SpanContext;

  /** Parent span ID */
  parentSpanId?: string;

  /** Set an attribute */
  setAttribute(key: string, value: string | number | boolean): void;

  /** Set multiple attributes */
  setAttributes(attributes: SpanAttributes): void;

  /** Set span status */
  setStatus(status: SpanStatus, message?: string): void;

  /** Add an event to the span */
  addEvent(name: string, attributes?: SpanAttributes): void;

  /** End the span */
  end(): void;

  /** Record an exception */
  recordException(exception: Error): void;
}

/**
 * Tracer interface
 */
export interface Tracer {
  /** Start a new span */
  startSpan(name: string, attributes?: SpanAttributes): Span;

  /** Get the active span */
  getActiveSpan(): Span | null;

  /** Execute function within span context */
  withSpan<T>(name: string, fn: (span: Span) => Promise<T>): Promise<T>;
}

/**
 * Generate a random trace ID
 */
function generateTraceId(): string {
  return Array.from({ length: 32 }, () => Math.floor(Math.random() * 16).toString(16)).join('');
}

/**
 * Generate a random span ID
 */
function generateSpanId(): string {
  return Array.from({ length: 16 }, () => Math.floor(Math.random() * 16).toString(16)).join('');
}

/**
 * Basic in-memory span implementation
 */
export class InMemorySpan implements Span {
  public name: string;
  public startTime: Date;
  public endTime?: Date;
  public status: SpanStatus = SpanStatus.UNSET;
  public attributes: SpanAttributes;
  public context: SpanContext;
  public parentSpanId?: string;
  public events: Array<{ name: string; timestamp: Date; attributes?: SpanAttributes }> = [];

  constructor(
    name: string,
    attributes: SpanAttributes = {},
    parentContext?: SpanContext
  ) {
    this.name = name;
    this.startTime = new Date();
    this.attributes = { ...attributes };

    if (parentContext) {
      this.context = {
        traceId: parentContext.traceId,
        spanId: generateSpanId(),
        traceFlags: parentContext.traceFlags,
        traceState: parentContext.traceState,
      };
      this.parentSpanId = parentContext.spanId;
    } else {
      this.context = {
        traceId: generateTraceId(),
        spanId: generateSpanId(),
        traceFlags: 1,
      };
    }
  }

  setAttribute(key: string, value: string | number | boolean): void {
    this.attributes[key] = value;
  }

  setAttributes(attributes: SpanAttributes): void {
    Object.assign(this.attributes, attributes);
  }

  setStatus(status: SpanStatus, message?: string): void {
    this.status = status;
    if (message) {
      this.setAttribute('status.message', message);
    }
  }

  addEvent(name: string, attributes?: SpanAttributes): void {
    this.events.push({
      name,
      timestamp: new Date(),
      attributes,
    });
  }

  recordException(exception: Error): void {
    this.setStatus(SpanStatus.ERROR);
    this.setAttribute('exception.type', exception.name);
    this.setAttribute('exception.message', exception.message);
    if (exception.stack) {
      this.setAttribute('exception.stacktrace', exception.stack);
    }
  }

  end(): void {
    if (!this.endTime) {
      this.endTime = new Date();

      if (this.status === SpanStatus.UNSET) {
        this.status = SpanStatus.OK;
      }
    }
  }

  /**
   * Get duration in milliseconds
   */
  getDuration(): number | null {
    if (!this.endTime) return null;
    return this.endTime.getTime() - this.startTime.getTime();
  }

  /**
   * Convert span to JSON
   */
  toJSON(): Record<string, unknown> {
    return {
      name: this.name,
      traceId: this.context.traceId,
      spanId: this.context.spanId,
      parentSpanId: this.parentSpanId,
      startTime: this.startTime.toISOString(),
      endTime: this.endTime?.toISOString(),
      duration: this.getDuration(),
      status: this.status,
      attributes: this.attributes,
      events: this.events,
    };
  }
}

/**
 * Span exporter interface
 */
export interface SpanExporter {
  export(spans: Span[]): void | Promise<void>;
}

/**
 * Console span exporter for debugging
 */
export class ConsoleSpanExporter implements SpanExporter {
  export(spans: Span[]): void {
    for (const span of spans) {
      if (span instanceof InMemorySpan) {
        console.log('[TRACE]', JSON.stringify(span.toJSON(), null, 2));
      }
    }
  }
}

/**
 * In-memory tracer implementation
 */
export class InMemoryTracer implements Tracer {
  private activeSpan: Span | null = null;
  private spans: Span[] = [];
  private exporter: SpanExporter;

  constructor(exporter: SpanExporter = new ConsoleSpanExporter()) {
    this.exporter = exporter;
  }

  startSpan(name: string, attributes?: SpanAttributes): Span {
    const parentContext = this.activeSpan?.context;
    const span = new InMemorySpan(name, attributes, parentContext);
    this.spans.push(span);
    return span;
  }

  getActiveSpan(): Span | null {
    return this.activeSpan;
  }

  async withSpan<T>(name: string, fn: (span: Span) => Promise<T>): Promise<T> {
    const span = this.startSpan(name);
    const previousActiveSpan = this.activeSpan;
    this.activeSpan = span;

    try {
      const result = await fn(span);
      span.setStatus(SpanStatus.OK);
      return result;
    } catch (error) {
      span.recordException(error as Error);
      throw error;
    } finally {
      span.end();
      this.activeSpan = previousActiveSpan;
      this.exportSpan(span);
    }
  }

  private exportSpan(span: Span): void {
    try {
      const result = this.exporter.export([span]);
      if (result && typeof result.then === 'function') {
        result.catch((err) => {
          console.error('Span export error:', err);
        });
      }
    } catch (err) {
      console.error('Span export error:', err);
    }
  }

  /**
   * Get all recorded spans
   */
  getSpans(): Span[] {
    return [...this.spans];
  }

  /**
   * Clear all spans
   */
  clear(): void {
    this.spans = [];
  }

  /**
   * Create context carrier for propagation
   */
  inject(carrier: Record<string, string>): void {
    if (this.activeSpan) {
      const { traceId, spanId, traceFlags } = this.activeSpan.context;
      carrier['traceparent'] = `00-${traceId}-${spanId}-${traceFlags.toString(16).padStart(2, '0')}`;
    }
  }

  /**
   * Extract context from carrier
   */
  extract(carrier: Record<string, string>): SpanContext | null {
    const traceparent = carrier['traceparent'];
    if (!traceparent) return null;

    const parts = traceparent.split('-');
    if (parts.length !== 4) return null;

    return {
      traceId: parts[1],
      spanId: parts[2],
      traceFlags: parseInt(parts[3], 16),
    };
  }
}

/**
 * No-op tracer for when tracing is disabled
 */
class NoOpSpan implements Span {
  name = '';
  startTime = new Date();
  status = SpanStatus.UNSET;
  attributes = {};
  context: SpanContext = {
    traceId: '00000000000000000000000000000000',
    spanId: '0000000000000000',
    traceFlags: 0,
  };

  setAttribute(): void {}
  setAttributes(): void {}
  setStatus(): void {}
  addEvent(): void {}
  end(): void {}
  recordException(): void {}
}

class NoOpTracer implements Tracer {
  private noOpSpan = new NoOpSpan();

  startSpan(): Span {
    return this.noOpSpan;
  }

  getActiveSpan(): Span | null {
    return null;
  }

  async withSpan<T>(_name: string, fn: (span: Span) => Promise<T>): Promise<T> {
    return fn(this.noOpSpan);
  }
}

/**
 * Default tracer instance
 */
export const defaultTracer: Tracer = new InMemoryTracer();

/**
 * No-op tracer instance (for performance-sensitive scenarios)
 */
export const noOpTracer: Tracer = new NoOpTracer();

/**
 * Tracing provider configuration
 */
export interface TracingConfig {
  enabled: boolean;
  exporter?: SpanExporter;
  serviceName?: string;
  serviceVersion?: string;
}

/**
 * Tracing provider
 */
export class TracingProvider {
  private tracer: Tracer;
  private config: TracingConfig;

  constructor(config: Partial<TracingConfig> = {}) {
    this.config = {
      enabled: config.enabled !== false,
      exporter: config.exporter,
      serviceName: config.serviceName || 'elasticsearch-client',
      serviceVersion: config.serviceVersion,
    };

    this.tracer = this.config.enabled
      ? new InMemoryTracer(this.config.exporter)
      : noOpTracer;
  }

  getTracer(): Tracer {
    return this.tracer;
  }

  isEnabled(): boolean {
    return this.config.enabled;
  }

  createSpan(name: string, attributes?: SpanAttributes): Span {
    const span = this.tracer.startSpan(name, {
      ...attributes,
      'service.name': this.config.serviceName,
      ...(this.config.serviceVersion && { 'service.version': this.config.serviceVersion }),
    });
    return span;
  }

  async withSpan<T>(name: string, fn: (span: Span) => Promise<T>): Promise<T> {
    return this.tracer.withSpan(name, fn);
  }

  propagateContext(carrier: Record<string, string>): void {
    if (this.tracer instanceof InMemoryTracer) {
      this.tracer.inject(carrier);
    }
  }

  extractContext(carrier: Record<string, string>): SpanContext | null {
    if (this.tracer instanceof InMemoryTracer) {
      return this.tracer.extract(carrier);
    }
    return null;
  }
}

/**
 * Default tracing provider
 */
export const defaultTracingProvider = new TracingProvider();