autotel-edge 3.0.0

package/src/core/spanprocessor.ts

@@ -0,0 +1,232 @@
+/**
+ * Span processor with flush and tail sampling support
+ */
+
+import type { Context } from '@opentelemetry/api';
+import type {
+  ReadableSpan,
+  Span,
+  SpanExporter,
+  SpanProcessor,
+} from '@opentelemetry/sdk-trace-base';
+import type { PostProcessorFn, TailSampleFn, LocalTrace } from '../types';
+
+/**
+ * Span processor that supports flush by trace ID and tail sampling
+ */
+export class SpanProcessorWithFlush implements SpanProcessor {
+  private exporter: SpanExporter;
+  private postProcessor?: PostProcessorFn;
+  private spans: Map<string, ReadableSpan[]> = new Map();
+
+  constructor(exporter: SpanExporter, postProcessor?: PostProcessorFn) {
+    this.exporter = exporter;
+    this.postProcessor = postProcessor;
+  }
+
+  onStart(_span: Span, _parentContext: Context): void {
+    // No-op for now
+  }
+
+  onEnd(span: ReadableSpan): void {
+    const traceId = span.spanContext().traceId;
+
+    if (!this.spans.has(traceId)) {
+      this.spans.set(traceId, []);
+    }
+
+    this.spans.get(traceId)!.push(span);
+  }
+
+  /**
+   * Force flush spans for a specific trace
+   */
+  async forceFlush(traceId?: string): Promise<void> {
+    if (traceId) {
+      const spans = this.spans.get(traceId);
+      if (spans && spans.length > 0) {
+        await this.exportSpans(spans);
+        this.spans.delete(traceId);
+      }
+    } else {
+      // Flush all traces
+      const promises: Promise<void>[] = [];
+      for (const [id, spans] of this.spans.entries()) {
+        promises.push(this.exportSpans(spans));
+        this.spans.delete(id);
+      }
+      await Promise.all(promises);
+    }
+  }
+
+  async shutdown(): Promise<void> {
+    await this.forceFlush();
+    if (this.exporter) {
+      await this.exporter.shutdown();
+    }
+  }
+
+  /**
+   * Export spans with post-processing
+   * Errors are caught and logged but don't throw to prevent worker instability
+   */
+  private async exportSpans(spans: ReadableSpan[]): Promise<void> {
+    if (spans.length === 0) return;
+    if (!this.exporter) return; // No exporter configured (e.g., in tests)
+
+    let processedSpans = spans;
+
+    if (this.postProcessor) {
+      try {
+        processedSpans = this.postProcessor(spans);
+      } catch (error) {
+        // Post-processor errors should not prevent export
+        console.error('[autotel-edge] Post-processor error:', error);
+        // Continue with original spans
+        processedSpans = spans;
+      }
+    }
+
+    return new Promise((resolve) => {
+      this.exporter.export(processedSpans, (result) => {
+        if (result.code === 0) {
+          // SUCCESS
+          resolve();
+        } else {
+          // Log but don't reject - exporter failures shouldn't crash the worker
+          console.error(
+            '[autotel-edge] Exporter error:',
+            result.error?.message || 'Unknown error',
+          );
+          resolve(); // Resolve instead of reject to prevent unhandled promise rejection
+        }
+      });
+    });
+  }
+}
+
+/**
+ * Span processor that supports tail sampling decisions
+ */
+export class TailSamplingSpanProcessor implements SpanProcessor {
+  private wrapped: SpanProcessorWithFlush;
+  private tailSampler?: TailSampleFn;
+  private traces: Map<string, LocalTrace> = new Map();
+
+  constructor(
+    exporter: SpanExporter,
+    postProcessor?: PostProcessorFn,
+    tailSampler?: TailSampleFn,
+  ) {
+    this.wrapped = new SpanProcessorWithFlush(exporter, postProcessor);
+    this.tailSampler = tailSampler;
+  }
+
+  onStart(span: Span, parentContext: Context): void {
+    this.wrapped.onStart(span, parentContext);
+  }
+
+  onEnd(span: ReadableSpan): void {
+    const traceId = span.spanContext().traceId;
+    const spanId = span.spanContext().spanId;
+    const parentSpanId = 'parentSpanId' in span ? span.parentSpanId : undefined;
+
+    // Initialize trace if not exists
+    if (!this.traces.has(traceId)) {
+      this.traces.set(traceId, {
+        traceId,
+        spans: [],
+        localRootSpan: undefined as any, // Will be set when we identify the local root
+      });
+    }
+
+    const trace = this.traces.get(traceId)!;
+
+    // Determine if this span is a local root by checking if its parent is in buffered spans
+    // A span is a local root if:
+    // 1. It has no parentSpanId (definitive root)
+    // 2. Its parentSpanId doesn't match any already-buffered span (remote parent = distributed trace entry)
+    const hasLocalParent = parentSpanId &&
+                            trace.spans.some(s => s.spanContext().spanId === parentSpanId);
+
+    // Set localRootSpan if this is the local root (no local parent found in buffer)
+    if (!hasLocalParent) {
+      trace.localRootSpan = span;
+    }
+
+    trace.spans.push(span); // Buffer the span AFTER checking parent relationships
+
+    // Auto-flush decision: only auto-flush for normal traces (no parentSpanId at all)
+    // For distributed traces (parentSpanId present), we rely on explicit forceFlush() from instrument.ts
+    // This ensures we don't trigger before all spans have been buffered
+    const isDefinitiveRoot = !parentSpanId;
+    const shouldAutoFlush = isDefinitiveRoot && trace.localRootSpan &&
+                             trace.localRootSpan.spanContext().spanId === spanId;
+
+    if (shouldAutoFlush) {
+      if (this.tailSampler) {
+        const shouldKeep = this.tailSampler(trace);
+
+        if (shouldKeep) {
+          // Export ALL buffered spans in the trace
+          for (const bufferedSpan of trace.spans) {
+            this.wrapped.onEnd(bufferedSpan);
+          }
+          // Force flush to actually export the spans
+          void this.wrapped.forceFlush(traceId);
+        }
+        // If not keeping, just drop all spans (don't export)
+      } else {
+        // No tail sampler, export all buffered spans
+        for (const bufferedSpan of trace.spans) {
+          this.wrapped.onEnd(bufferedSpan);
+        }
+        // Force flush to actually export the spans
+        void this.wrapped.forceFlush(traceId);
+      }
+
+      // Clean up trace after decision
+      this.traces.delete(traceId);
+    }
+    // If not local root span, just buffer it - don't export yet
+  }
+
+  async forceFlush(traceId?: string): Promise<void> {
+    if (traceId) {
+      // Make tail sampling decision for this specific trace before flushing
+      const trace = this.traces.get(traceId);
+      if (trace) {
+        // Ensure localRootSpan is set (fallback to first span if not)
+        // This handles distributed traces where no span has undefined parentSpanId
+        if (!trace.localRootSpan && trace.spans.length > 0) {
+          trace.localRootSpan = trace.spans[0];
+        }
+
+        if (this.tailSampler) {
+          const shouldKeep = this.tailSampler(trace);
+
+          if (shouldKeep) {
+            // Export ALL buffered spans in the trace
+            for (const bufferedSpan of trace.spans) {
+              this.wrapped.onEnd(bufferedSpan);
+            }
+          }
+        } else {
+          // No tail sampler, export all buffered spans
+          for (const bufferedSpan of trace.spans) {
+            this.wrapped.onEnd(bufferedSpan);
+          }
+        }
+
+        // Clean up trace after decision
+        this.traces.delete(traceId);
+      }
+    }
+    return this.wrapped.forceFlush(traceId);
+  }
+
+  async shutdown(): Promise<void> {
+    this.traces.clear();
+    return this.wrapped.shutdown();
+  }
+}

package/src/core/trace-context.ts

@@ -0,0 +1,66 @@
+/**
+ * Trace context types and utilities
+ */
+
+import type { Span, SpanStatusCode } from '@opentelemetry/api';
+
+/**
+ * WeakMap to store span names for active spans.
+ * Enables retrieving span names for correlation helpers.
+ */
+const spanNameMap = new WeakMap<Span, string>();
+
+/**
+ * Base trace context containing trace identifiers
+ */
+export interface TraceContextBase {
+  traceId: string;
+  spanId: string;
+  correlationId: string;
+  'code.function'?: string;
+}
+
+/**
+ * Span methods available on trace context
+ */
+export interface SpanMethods {
+  setAttribute(key: string, value: string | number | boolean): void;
+  setAttributes(attrs: Record<string, string | number | boolean>): void;
+  setStatus(status: { code: SpanStatusCode; message?: string }): void;
+  recordException(exception: Error): void;
+}
+
+/**
+ * Complete trace context that merges base context and span methods
+ *
+ * This is the ctx parameter passed to factory functions in trace().
+ * It provides access to trace IDs and span manipulation methods.
+ */
+export type TraceContext = TraceContextBase & SpanMethods;
+
+/**
+ * Create a TraceContext from an OpenTelemetry Span
+ *
+ * This utility extracts trace context information from a span
+ * and provides span manipulation methods in a consistent format.
+ */
+export function createTraceContext(span: Span): TraceContext {
+  const spanContext = span.spanContext();
+  return {
+    traceId: spanContext.traceId,
+    spanId: spanContext.spanId,
+    correlationId: spanContext.traceId.slice(0, 16),
+    'code.function': spanNameMap.get(span),
+    setAttribute: span.setAttribute.bind(span),
+    setAttributes: span.setAttributes.bind(span),
+    setStatus: span.setStatus.bind(span),
+    recordException: span.recordException.bind(span),
+  };
+}
+
+/**
+ * Store the span name for later retrieval via trace context helpers.
+ */
+export function setSpanName(span: Span, name: string): void {
+  spanNameMap.set(span, name);
+}

package/src/core/tracer.test.ts

@@ -0,0 +1,123 @@
+import { describe, it, expect, beforeEach, vi } from 'vitest';
+import { WorkerTracer } from './tracer';
+import { resourceFromAttributes } from '@opentelemetry/resources';
+import { SamplingDecision, type SpanProcessor } from '@opentelemetry/sdk-trace-base';
+
+describe('WorkerTracer', () => {
+  let tracer: WorkerTracer;
+  let mockProcessor: SpanProcessor;
+  let mockHeadSampler: any;
+
+  beforeEach(() => {
+    mockProcessor = {
+      onStart: vi.fn(),
+      onEnd: vi.fn(),
+      shutdown: vi.fn(async () => {}),
+      forceFlush: vi.fn(async () => {}),
+    };
+
+    mockHeadSampler = {
+      shouldSample: vi.fn(() => ({
+        decision: SamplingDecision.RECORD_AND_SAMPLED,
+        attributes: {},
+        traceState: undefined,
+      })),
+      toString: () => 'MockHeadSampler',
+    };
+
+    const resource = resourceFromAttributes({});
+    tracer = new WorkerTracer([mockProcessor], resource);
+    tracer.setHeadSampler(mockHeadSampler);
+  });
+
+  describe('Per-span sampler', () => {
+    it('should use per-span sampler when provided in options', () => {
+      const perSpanSampler = {
+        shouldSample: vi.fn(() => ({
+          decision: SamplingDecision.NOT_RECORD,
+          attributes: {},
+          traceState: undefined,
+        })),
+        toString: () => 'PerSpanSampler',
+      };
+
+      const options: any = { sampler: perSpanSampler };
+      tracer.startSpan('test.span', options);
+
+      // Per-span sampler should be called, NOT head sampler
+      expect(perSpanSampler.shouldSample).toHaveBeenCalledTimes(1);
+      expect(mockHeadSampler.shouldSample).not.toHaveBeenCalled();
+    });
+
+    it('should use head sampler when no per-span sampler provided', () => {
+      tracer.startSpan('test.span', {});
+
+      // Head sampler should be called
+      expect(mockHeadSampler.shouldSample).toHaveBeenCalledTimes(1);
+    });
+
+    it('should pass correct arguments to per-span sampler', () => {
+      const perSpanSampler = {
+        shouldSample: vi.fn(() => ({
+          decision: SamplingDecision.RECORD_AND_SAMPLED,
+          attributes: {},
+          traceState: undefined,
+        })),
+        toString: () => 'PerSpanSampler',
+      };
+
+      const options: any = {
+        sampler: perSpanSampler,
+        attributes: { 'test.attr': 'value' },
+      };
+
+      tracer.startSpan('test.span', options);
+
+      // Verify sampler was called with correct arguments
+      expect(perSpanSampler.shouldSample).toHaveBeenCalledWith(
+        expect.anything(), // context
+        expect.any(String), // traceId
+        'test.span', // span name
+        expect.any(Number), // spanKind
+        expect.objectContaining({ 'test.attr': 'value' }), // attributes
+        [], // links
+      );
+    });
+
+    it('should respect per-span sampler decision to NOT_RECORD', () => {
+      const rejectingSampler = {
+        shouldSample: vi.fn(() => ({
+          decision: SamplingDecision.NOT_RECORD,
+          attributes: {},
+          traceState: undefined,
+        })),
+        toString: () => 'RejectingSampler',
+      };
+
+      const options: any = { sampler: rejectingSampler };
+      const span = tracer.startSpan('test.span', options);
+
+      // Span should be created but not sampled
+      expect(span).toBeDefined();
+      expect(span.spanContext().traceFlags & 1).toBe(0); // Not sampled
+    });
+
+    it('should respect per-span sampler decision to RECORD_AND_SAMPLED', () => {
+      const acceptingSampler = {
+        shouldSample: vi.fn(() => ({
+          decision: SamplingDecision.RECORD_AND_SAMPLED,
+          attributes: {},
+          traceState: undefined,
+        })),
+        toString: () => 'AcceptingSampler',
+      };
+
+      const options: any = { sampler: acceptingSampler };
+      const span = tracer.startSpan('test.span', options);
+
+      // Span should be created and sampled
+      expect(span).toBeDefined();
+      expect(span.spanContext().traceFlags & 1).toBe(1); // Sampled
+    });
+  });
+});

package/src/core/tracer.ts

@@ -0,0 +1,216 @@
+/**
+ * Lightweight WorkerTracer for edge environments
+ */
+
+import type {
+  Attributes,
+  Tracer,
+  Span,
+  SpanKind,
+  SpanOptions,
+  Context,
+} from '@opentelemetry/api';
+import {
+  context as api_context,
+  trace,
+  type SpanContext,
+} from '@opentelemetry/api';
+import { sanitizeAttributes } from '@opentelemetry/core';
+import type { Resource } from '@opentelemetry/resources';
+import {
+  type SpanProcessor,
+  RandomIdGenerator,
+  type ReadableSpan,
+  SamplingDecision,
+} from '@opentelemetry/sdk-trace-base';
+
+import { SpanImpl } from './span';
+import type { TraceFlushableSpanProcessor } from '../types';
+
+const NewTraceFlags = {
+  RANDOM_TRACE_ID_SET: 2,
+  RANDOM_TRACE_ID_UNSET: 0,
+} as const;
+
+type NewTraceFlagValues =
+  | typeof NewTraceFlags.RANDOM_TRACE_ID_SET
+  | typeof NewTraceFlags.RANDOM_TRACE_ID_UNSET;
+
+const idGenerator: RandomIdGenerator = new RandomIdGenerator();
+
+let withNextSpanAttributes: Attributes;
+
+function getFlagAt(flagSequence: number, position: number): number {
+  return ((flagSequence >> (position - 1)) & 1) * position;
+}
+
+/**
+ * WorkerTracer - Lightweight tracer for edge environments
+ */
+export class WorkerTracer implements Tracer {
+  private readonly spanProcessors: TraceFlushableSpanProcessor[];
+  private readonly resource: Resource;
+  private headSampler: any; // Will be set via setHeadSampler
+
+  constructor(spanProcessors: SpanProcessor[], resource: Resource) {
+    this.spanProcessors = spanProcessors as TraceFlushableSpanProcessor[];
+    this.resource = resource;
+  }
+
+  /**
+   * Set the head sampler (called from config)
+   */
+  setHeadSampler(sampler: any): void {
+    this.headSampler = sampler;
+  }
+
+  /**
+   * Force flush spans for a specific trace
+   */
+  async forceFlush(traceId?: string) {
+    const promises = this.spanProcessors.map(async (spanProcessor) => {
+      await spanProcessor.forceFlush(traceId);
+    });
+    await Promise.allSettled(promises);
+  }
+
+  /**
+   * Add extra resource attributes
+   */
+  addToResource(extra: Resource) {
+    this.resource.merge(extra);
+  }
+
+  /**
+   * Start a new span
+   */
+  startSpan(
+    name: string,
+    options: SpanOptions = {},
+    context = api_context.active(),
+  ): Span {
+    if (options.root) {
+      context = trace.deleteSpan(context);
+    }
+
+    if (!this.headSampler) {
+      throw new Error(
+        'Head sampler not configured. This is a bug in the instrumentation logic',
+      );
+    }
+
+    const parentSpanContext = trace.getSpan(context)?.spanContext();
+    const { traceId, randomTraceFlag } = getTraceInfo(parentSpanContext);
+
+    const spanKind = options.kind || (0 as SpanKind); // SpanKind.INTERNAL
+    const sanitisedAttrs = sanitizeAttributes(options.attributes);
+
+    // Use per-span sampler if provided, otherwise use head sampler
+    const optionsWithSampler = options as any;
+    const sampler = optionsWithSampler.sampler || this.headSampler;
+
+    const samplingDecision = sampler.shouldSample(
+      context,
+      traceId,
+      name,
+      spanKind,
+      sanitisedAttrs,
+      [],
+    );
+    const { decision, traceState, attributes: attrs } = samplingDecision;
+
+    const attributes = Object.assign(
+      {},
+      options.attributes,
+      attrs,
+      withNextSpanAttributes,
+    );
+    withNextSpanAttributes = {};
+
+    const spanId = idGenerator.generateSpanId();
+    const parentSpanId = parentSpanContext?.spanId;
+
+    const sampleFlag =
+      decision === SamplingDecision.RECORD_AND_SAMPLED ? 1 : 0; // TraceFlags.SAMPLED : TraceFlags.NONE
+    const traceFlags = sampleFlag + randomTraceFlag;
+    const spanContext: SpanContext = { traceId, spanId, traceFlags, traceState };
+
+    const span = new SpanImpl({
+      attributes: sanitizeAttributes(attributes),
+      name,
+      onEnd: (span) => {
+        for (const sp of this.spanProcessors) {
+          sp.onEnd(span as unknown as ReadableSpan);
+        }
+      },
+      resource: this.resource,
+      spanContext,
+      parentSpanContext,
+      parentSpanId,
+      spanKind,
+      startTime: options.startTime,
+    });
+
+    for (const sp of this.spanProcessors) {
+      //@ts-ignore - OTel type quirk
+      sp.onStart(span, context);
+    }
+
+    return span;
+  }
+
+  /**
+   * Start an active span (with automatic context management)
+   */
+  startActiveSpan<F extends (span: Span) => ReturnType<F>>(
+    name: string,
+    fn: F,
+  ): ReturnType<F>;
+  startActiveSpan<F extends (span: Span) => ReturnType<F>>(
+    name: string,
+    options: SpanOptions,
+    fn: F,
+  ): ReturnType<F>;
+  startActiveSpan<F extends (span: Span) => ReturnType<F>>(
+    name: string,
+    options: SpanOptions,
+    context: Context,
+    fn: F,
+  ): ReturnType<F>;
+  startActiveSpan<F extends (span: Span) => ReturnType<F>>(
+    name: string,
+    ...args: unknown[]
+  ): ReturnType<F> {
+    const options = args.length > 1 ? (args[0] as SpanOptions) : undefined;
+    const parentContext =
+      args.length > 2 ? (args[1] as Context) : api_context.active();
+    const fn = args.at(-1) as F;
+
+    const span = this.startSpan(name, options, parentContext);
+    const contextWithSpanSet = trace.setSpan(parentContext, span);
+
+    return api_context.with(contextWithSpanSet, fn, undefined, span);
+  }
+}
+
+/**
+ * Set attributes for the next span created
+ */
+export function withNextSpan(attrs: Attributes) {
+  withNextSpanAttributes = Object.assign({}, withNextSpanAttributes, attrs);
+}
+
+function getTraceInfo(parentSpanContext?: SpanContext): {
+  traceId: string;
+  randomTraceFlag: NewTraceFlagValues;
+} {
+  if (parentSpanContext && trace.isSpanContextValid(parentSpanContext)) {
+    const { traceId, traceFlags } = parentSpanContext;
+    return { traceId, randomTraceFlag: getFlagAt(traceFlags, 2) as NewTraceFlagValues };
+  } else {
+    return {
+      traceId: idGenerator.generateTraceId(),
+      randomTraceFlag: NewTraceFlags.RANDOM_TRACE_ID_SET,
+    };
+  }
+}