ref: Refactor span map handling

Author: mydea

packages/node-experimental/src/sdk/initOtel.ts

@@ -29,7 +29,7 @@ export function initOtel(): () => void {
   const provider = new NodeTracerProvider({
     sampler: new AlwaysOnSampler(),
   });
-  provider.addSpanProcessor(new SentrySpanProcessor({ strictSpanParentHandling: true }));
+  provider.addSpanProcessor(new SentrySpanProcessor());
 
   // We use a custom context manager to keep context in sync with sentry scope
   const contextManager = new SentryContextManager();

packages/opentelemetry-node/src/index.ts

@@ -1,4 +1,4 @@
-import { getSentrySpan } from './spanprocessor';
+import { getSentrySpan } from './utils/spanMap';
 
 export { SentrySpanProcessor } from './spanprocessor';
 export { SentryPropagator } from './propagator';

packages/opentelemetry-node/src/propagator.ts

@@ -13,7 +13,7 @@ import {
   SENTRY_TRACE_HEADER,
   SENTRY_TRACE_PARENT_CONTEXT_KEY,
 } from './constants';
-import { SENTRY_SPAN_PROCESSOR_MAP } from './spanprocessor';
+import { getSentrySpan } from './utils/spanMap';
 
 /**
  * Injects and extracts `sentry-trace` and `baggage` headers from carriers.
@@ -30,7 +30,7 @@ export class SentryPropagator extends W3CBaggagePropagator {
 
     let baggage = propagation.getBaggage(context) || propagation.createBaggage({});
 
-    const span = SENTRY_SPAN_PROCESSOR_MAP.get(spanContext.spanId);
+    const span = getSentrySpan(spanContext.spanId);
     if (span) {
       setter.set(carrier, SENTRY_TRACE_HEADER, span.toTraceparent());
 

packages/opentelemetry-node/src/spanprocessor.ts

@@ -10,46 +10,14 @@ import { SENTRY_DYNAMIC_SAMPLING_CONTEXT_KEY, SENTRY_TRACE_PARENT_CONTEXT_KEY }
 import { isSentryRequestSpan } from './utils/isSentryRequest';
 import { mapOtelStatus } from './utils/mapOtelStatus';
 import { parseSpanDescription } from './utils/parseOtelSpanDescription';
-
-interface SpanProcessorOptions {
-  /**
-   * By default, if a span is started and we cannot find a Sentry parent span for it,
-   * even if the OTEL span has a parent reference, we will still create the Sentry span as a root span.
-   *
-   * While this is more tolerant of errors, it means that the generated Spans in Sentry may have an incorrect hierarchy.
-   *
-   * When opting into strict span parent handling, we will discard any Spans where we can't find the corresponding parent.
-   * This also requires that we defer clearing of references to the point where the root span is finished -
-   * as sometimes these are not fired in correct order, leading to spans being dropped.
-   *
-   * Note that enabling this is the more correct option
-   * and will probably eventually become the default in a future version.
-   */
-  strictSpanParentHandling: boolean;
-}
-
-export const SENTRY_SPAN_PROCESSOR_MAP: Map<string, SentrySpan> = new Map<string, SentrySpan>();
-
-// A map of a sentry span ID to a list of otel span IDs
-// When the sentry span is finished, clear all references of the given otel spans
-export const SCHEDULE_TO_CLEAR: Map<string, string[]> = new Map<string, string[]>();
-
-/** Get a Sentry span for an otel span ID. */
-export function getSentrySpan(otelSpanId: string): SentrySpan | undefined {
-  return SENTRY_SPAN_PROCESSOR_MAP.get(otelSpanId);
-}
+import { clearSpan, getSentrySpan, setSentrySpan } from './utils/spanMap';
 
 /**
  * Converts OpenTelemetry Spans to Sentry Spans and sends them to Sentry via
  * the Sentry SDK.
  */
 export class SentrySpanProcessor implements OtelSpanProcessor {
-  private _strictSpanParentHandling: boolean;
-
-  public constructor({ strictSpanParentHandling }: Partial<SpanProcessorOptions> = {}) {
-    // Default to false
-    this._strictSpanParentHandling = !!strictSpanParentHandling;
-
+  public constructor() {
     addTracingExtensions();
 
     addGlobalEventProcessor(event => {
@@ -83,7 +51,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
 
     // Otel supports having multiple non-nested spans at the same time
     // so we cannot use hub.getSpan(), as we cannot rely on this being on the current span
-    const sentryParentSpan = otelParentSpanId && SENTRY_SPAN_PROCESSOR_MAP.get(otelParentSpanId);
+    const sentryParentSpan = otelParentSpanId && getSentrySpan(otelParentSpanId);
 
     if (sentryParentSpan) {
       const sentryChildSpan = sentryParentSpan.startChild({
@@ -93,7 +61,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
         spanId: otelSpanId,
       });
 
-      SENTRY_SPAN_PROCESSOR_MAP.set(otelSpanId, sentryChildSpan);
+      setSentrySpan(otelSpanId, sentryChildSpan);
     } else {
       const traceCtx = getTraceData(otelSpan, parentContext);
       const transaction = getCurrentHub().startTransaction({
@@ -104,11 +72,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
         spanId: otelSpanId,
       });
 
-      SENTRY_SPAN_PROCESSOR_MAP.set(otelSpanId, transaction);
-
-      if (this._strictSpanParentHandling) {
-        SCHEDULE_TO_CLEAR.set(transaction.spanId, []);
-      }
+      setSentrySpan(otelSpanId, transaction);
     }
   }
 
@@ -122,7 +86,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
     if (!sentrySpan) {
       __DEBUG_BUILD__ &&
         logger.error(`SentrySpanProcessor could not find span with OTEL-spanId ${otelSpanId} to finish.`);
-      this._clearSpan(otelSpanId);
+      clearSpan(otelSpanId);
       return;
     }
 
@@ -131,7 +95,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
     // leading to an infinite loop.
     // In this case, we do not want to finish the span, in order to avoid sending it to Sentry
     if (isSentryRequestSpan(otelSpan)) {
-      this._clearSpan(otelSpanId);
+      clearSpan(otelSpanId);
       return;
     }
 
@@ -141,7 +105,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
     client && client.emit && client?.emit('otelSpanEnd', otelSpan, mutableOptions);
 
     if (mutableOptions.drop) {
-      this._clearSpan(otelSpanId);
+      clearSpan(otelSpanId);
       return;
     }
 
@@ -194,7 +158,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
 
     sentrySpan.finish(convertOtelTimeToSeconds(otelSpan.endTime));
 
-    this._clearSpan(otelSpanId);
+    clearSpan(otelSpanId);
   }
 
   /**
@@ -214,17 +178,6 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
     }
     return Promise.resolve();
   }
-
-  /**
-   * Clear all references for a given OTEL span.
-   */
-  private _clearSpan(otelSpanId: string): void {
-    if (this._strictSpanParentHandling) {
-      scheduleToClear(otelSpanId);
-    } else {
-      clearSpan(otelSpanId);
-    }
-  }
 }
 
 function getTraceData(otelSpan: OtelSpan, parentContext: Context): Partial<TransactionContext> {
@@ -300,50 +253,3 @@ function updateTransactionWithOtelData(transaction: Transaction, otelSpan: OtelS
 function convertOtelTimeToSeconds([seconds, nano]: [number, number]): number {
   return seconds + nano / 1_000_000_000;
 }
-
-function scheduleToClear(otelSpanId: string): void {
-  const span = SENTRY_SPAN_PROCESSOR_MAP.get(otelSpanId);
-
-  if (!span) {
-    // hmm, something is fishy here, but abort...
-    // But to be sure we still try to delete the SCHEDULE_TO_CLEAR, to avoid leaks
-    SCHEDULE_TO_CLEAR.delete(otelSpanId);
-    return;
-  }
-
-  const sentrySpanId = span.spanId;
-
-  // This is the root, clear all that have been scheduled
-  if (spanIsRoot(span) || !span.transaction) {
-    const toClear = SCHEDULE_TO_CLEAR.get(sentrySpanId) || [];
-    toClear.push(otelSpanId);
-
-    toClear.forEach(otelSpanIdToClear => clearSpan(otelSpanIdToClear));
-    SCHEDULE_TO_CLEAR.delete(sentrySpanId);
-    return;
-  }
-
-  // Clear when root span is cleared
-  const root = span.transaction;
-  const rootSentrySpanId = root.spanId;
-
-  const toClear = SCHEDULE_TO_CLEAR.get(rootSentrySpanId);
-
-  // If this does not exist, it means we prob. already cleaned it up before
-  // So we ignore the parent and just clean this span up right now
-  if (!toClear) {
-    clearSpan(otelSpanId);
-    return;
-  }
-
-  toClear.push(otelSpanId);
-}
-
-function spanIsRoot(span: SentrySpan): span is Transaction {
-  return span.transaction === span;
-}
-
-// make sure to remove references in maps, to ensure this can be GCed
-function clearSpan(otelSpanId: string): void {
-  SENTRY_SPAN_PROCESSOR_MAP.delete(otelSpanId);
-}

packages/opentelemetry-node/src/utils/spanData.ts

@@ -2,7 +2,7 @@
 import { Transaction } from '@sentry/core';
 import type { Context, SpanOrigin } from '@sentry/types';
 
-import { getSentrySpan } from '../spanprocessor';
+import { getSentrySpan } from './spanMap';
 
 type SentryTags = Record<string, string>;
 type SentryData = Record<string, unknown>;

packages/opentelemetry-node/src/utils/spanMap.ts

@@ -0,0 +1,92 @@
+import type { Span as SentrySpan } from '@sentry/types';
+
+interface SpanMapEntry {
+  sentrySpan: SentrySpan;
+  ref: SpanRefType;
+  // These are not direct children, but all spans under the tree of a root span.
+  subSpans: string[];
+}
+
+const SPAN_REF_ROOT = Symbol('root');
+const SPAN_REF_CHILD = Symbol('child');
+const SPAN_REF_CHILD_ENDED = Symbol('child_ended');
+type SpanRefType = typeof SPAN_REF_ROOT | typeof SPAN_REF_CHILD | typeof SPAN_REF_CHILD_ENDED;
+
+/** Exported only for tests. */
+export const SPAN_MAP = new Map<string, SpanMapEntry>();
+
+/**
+ * Get a Sentry span for a given span ID.
+ */
+export function getSentrySpan(spanId: string): SentrySpan | undefined {
+  const entry = SPAN_MAP.get(spanId);
+  return entry ? entry.sentrySpan : undefined;
+}
+
+/**
+ * Set a Sentry span for a given span ID.
+ * This is necessary so we can lookup parent spans later.
+ * We also keep a list of children for root spans only, in order to be able to clean them up together.
+ */
+export function setSentrySpan(spanId: string, sentrySpan: SentrySpan): void {
+  let ref: SpanRefType = SPAN_REF_ROOT;
+
+  const rootSpanId = sentrySpan.transaction?.spanId;
+
+  if (rootSpanId && rootSpanId !== spanId) {
+    const root = SPAN_MAP.get(rootSpanId);
+    if (root) {
+      root.subSpans.push(spanId);
+      ref = SPAN_REF_CHILD;
+    }
+  }
+
+  SPAN_MAP.set(spanId, {
+    sentrySpan,
+    ref,
+    subSpans: [],
+  });
+}
+
+/**
+ * Clear references of the given span ID.
+ */
+export function clearSpan(spanId: string): void {
+  const entry = SPAN_MAP.get(spanId);
+  if (!entry) {
+    return;
+  }
+
+  const { ref, subSpans } = entry;
+
+  // If this is a child, mark it as ended.
+  if (ref === SPAN_REF_CHILD) {
+    entry.ref = SPAN_REF_CHILD_ENDED;
+    return;
+  }
+
+  // If this is a root span, clear all (ended) children
+  if (ref === SPAN_REF_ROOT) {
+    for (const childId of subSpans) {
+      const child = SPAN_MAP.get(childId);
+      if (!child) {
+        continue;
+      }
+
+      if (child.ref === SPAN_REF_CHILD_ENDED) {
+        // if the child has already ended, just clear it
+        SPAN_MAP.delete(childId);
+      } else if (child.ref === SPAN_REF_CHILD) {
+        // If the child has not ended yet, mark it as a root span so it is cleared when it ends.
+        child.ref = SPAN_REF_ROOT;
+      }
+    }
+
+    SPAN_MAP.delete(spanId);
+    return;
+  }
+
+  // Generally, `clearSpan` should never be called for ref === SPAN_REF_CHILD_ENDED
+  // But if it does, just clear the span
+  SPAN_MAP.delete(spanId);
+}

packages/opentelemetry-node/test/propagator.test.ts

@@ -17,7 +17,7 @@ import {
   SENTRY_TRACE_PARENT_CONTEXT_KEY,
 } from '../src/constants';
 import { SentryPropagator } from '../src/propagator';
-import { SENTRY_SPAN_PROCESSOR_MAP } from '../src/spanprocessor';
+import { setSentrySpan, SPAN_MAP } from '../src/utils/spanMap';
 
 beforeAll(() => {
   addTracingExtensions();
@@ -51,7 +51,7 @@ describe('SentryPropagator', () => {
       makeMain(hub);
 
       afterEach(() => {
-        SENTRY_SPAN_PROCESSOR_MAP.clear();
+        SPAN_MAP.clear();
       });
 
       enum PerfType {
@@ -61,12 +61,12 @@ describe('SentryPropagator', () => {
 
       function createTransactionAndMaybeSpan(type: PerfType, transactionContext: TransactionContext) {
         const transaction = new Transaction(transactionContext, hub);
-        SENTRY_SPAN_PROCESSOR_MAP.set(transaction.spanId, transaction);
+        setSentrySpan(transaction.spanId, transaction);
         if (type === PerfType.Span) {
           // eslint-disable-next-line @typescript-eslint/no-unused-vars
           const { spanId, ...ctx } = transactionContext;
           const span = transaction.startChild({ ...ctx, description: transaction.name });
-          SENTRY_SPAN_PROCESSOR_MAP.set(span.spanId, span);
+          setSentrySpan(span.spanId, span);
         }
       }