fix(node-otel): Refactor OTEL span reference cleanup

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,18 +10,7 @@ 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';
-
-export const SENTRY_SPAN_PROCESSOR_MAP: Map<string, SentrySpan> = new Map<string, SentrySpan>();
-
-// make sure to remove references in maps, to ensure this can be GCed
-function clearSpan(otelSpanId: string): void {
-  SENTRY_SPAN_PROCESSOR_MAP.delete(otelSpanId);
-}
-
-/** 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
@@ -62,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({
@@ -72,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({
@@ -83,7 +72,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
         spanId: otelSpanId,
       });
 
-      SENTRY_SPAN_PROCESSOR_MAP.set(otelSpanId, transaction);
+      setSentrySpan(otelSpanId, transaction);
     }
   }
 
@@ -97,6 +86,7 @@ export class SentrySpanProcessor implements OtelSpanProcessor {
     if (!sentrySpan) {
       __DEBUG_BUILD__ &&
         logger.error(`SentrySpanProcessor could not find span with OTEL-spanId ${otelSpanId} to finish.`);
+      clearSpan(otelSpanId);
       return;
     }
 

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);
         }
       }