rework implementation to fork scope instead of manual cleanup

Author: Lms24

packages/core/src/tracing/trace.ts

@@ -46,9 +46,11 @@ export function startSpan<T>(options: StartSpanOptions, callback: (span: Span) =
   const spanArguments = parseSentrySpanArguments(options);
   const { forceTransaction, parentSpan: customParentSpan, scope: customScope } = options;
 
-  const previousActiveSpanOnCustomScope = customScope && _getSpanForScope(customScope);
+  // We still need to fork a potentially passed scope, as we set the active span on it
+  // and we need to ensure that it is cleaned up properly once the span ends.
+  const customForkedScope = customScope?.clone();
 
-  return withScope(customScope, () => {
+  return withScope(customForkedScope, () => {
     // If `options.parentSpan` is defined, we want to wrap the callback in `withActiveSpan`
     const wrapper = getActiveSpanWrapper<T>(customParentSpan);
 
@@ -79,11 +81,6 @@ export function startSpan<T>(options: StartSpanOptions, callback: (span: Span) =
         },
         () => {
           activeSpan.end();
-          if (customScope) {
-            // If a custom scope is passed, we don't fork the scope. Therefore, we need to reset the
-            // active span on the scope to the previous active span (or to undefined if there was none)
-            _setSpanForScope(customScope, previousActiveSpanOnCustomScope);
-          }
         },
       );
     });

packages/core/src/types-hoist/startSpanOptions.ts

@@ -5,7 +5,16 @@ export interface StartSpanOptions {
   /** A manually specified start time for the created `Span` object. */
   startTime?: SpanTimeInput;
 
-  /** If defined, start this span off this scope instead off the current scope. */
+  /**
+   * If set, start the span on a fork of this scope instead of on the current scope.
+   * To ensure proper span cleanup, the passed scope is cloned for the duration of the span.
+   *
+   * If you want to modify the passed scope inside the callback, be aware that:
+   * - calling `getCurrentScope()` will return the cloned scope, meaning all modifications
+   *   will be reset once the callback finishes
+   * - if you want to modify the passed scope and have the changes persist after the callback
+   *   ends, modify the scope directly and don't use `getCurrentScope()`
+   */
   scope?: Scope;
 
   /** The name of the span. */

packages/core/test/lib/tracing/trace.test.ts

@@ -25,7 +25,7 @@ import {
 import { SentryNonRecordingSpan } from '../../../src/tracing/sentryNonRecordingSpan';
 import { startNewTrace } from '../../../src/tracing/trace';
 import type { Event, Span, StartSpanOptions } from '../../../src/types-hoist';
-import { _setSpanForScope } from '../../../src/utils/spanOnScope';
+import { _getSpanForScope, _setSpanForScope } from '../../../src/utils/spanOnScope';
 import { getActiveSpan, getRootSpan, getSpanDescendants, spanIsSampled } from '../../../src/utils/spanUtils';
 import { TestClient, getDefaultTestClientOptions } from '../../mocks/client';
 
@@ -251,20 +251,44 @@ describe('startSpan', () => {
     expect(getActiveSpan()).toBe(undefined);
   });
 
-  it('allows to pass a scope', () => {
+  it('starts the span on the fork of a passed custom scope', () => {
     const initialScope = getCurrentScope();
 
-    const manualScope = initialScope.clone();
+    const customScope = initialScope.clone();
+    customScope.setTag('dogs', 'great');
+
     const parentSpan = new SentrySpan({ spanId: 'parent-span-id', sampled: true });
-    _setSpanForScope(manualScope, parentSpan);
+    _setSpanForScope(customScope, parentSpan);
 
-    startSpan({ name: 'GET users/[id]', scope: manualScope }, span => {
+    startSpan({ name: 'GET users/[id]', scope: customScope }, span => {
+      // current scope is forked from the customScope
       expect(getCurrentScope()).not.toBe(initialScope);
-      expect(getCurrentScope()).toBe(manualScope);
+      expect(getCurrentScope()).not.toBe(customScope);
+      expect(getCurrentScope().getScopeData().tags).toEqual({ dogs: 'great' });
+
+      // active span is set correctly
       expect(getActiveSpan()).toBe(span);
+
+      // span has the correct parent span
       expect(spanToJSON(span).parent_span_id).toBe('parent-span-id');
+
+      // scope data modifications
+      getCurrentScope().setTag('cats', 'great');
+      customScope.setTag('bears', 'great');
+
+      expect(getCurrentScope().getScopeData().tags).toEqual({ dogs: 'great', cats: 'great' });
+      expect(customScope.getScopeData().tags).toEqual({ dogs: 'great', bears: 'great' });
     });
 
+    // customScope modifications are persisted
+    expect(customScope.getScopeData().tags).toEqual({ dogs: 'great', bears: 'great' });
+
+    // span is parent span again on customScope
+    withScope(customScope, () => {
+      expect(getActiveSpan()).toBe(parentSpan);
+    });
+
+    // but activeSpan and currentScope are reset, since customScope was never active
     expect(getCurrentScope()).toBe(initialScope);
     expect(getActiveSpan()).toBe(undefined);
   });
@@ -273,29 +297,35 @@ describe('startSpan', () => {
     it('with parent span', () => {
       const initialScope = getCurrentScope();
 
-      const manualScope = initialScope.clone();
+      const customScope = initialScope.clone();
       const parentSpan = new SentrySpan({ spanId: 'parent-span-id', sampled: true });
-      _setSpanForScope(manualScope, parentSpan);
+      _setSpanForScope(customScope, parentSpan);
 
-      startSpan({ name: 'span 1', scope: manualScope }, span1 => {
+      startSpan({ name: 'span 1', scope: customScope }, span1 => {
+        // current scope is forked from the customScope
         expect(getCurrentScope()).not.toBe(initialScope);
-        expect(getCurrentScope()).toBe(manualScope);
+        expect(getCurrentScope()).not.toBe(customScope);
+
         expect(getActiveSpan()).toBe(span1);
         expect(spanToJSON(span1).parent_span_id).toBe('parent-span-id');
       });
 
-      withScope(manualScope, () => {
+      // active span on customScope is reset
+      withScope(customScope, () => {
         expect(getActiveSpan()).toBe(parentSpan);
       });
 
-      startSpan({ name: 'span 2', scope: manualScope }, span2 => {
+      startSpan({ name: 'span 2', scope: customScope }, span2 => {
+        // current scope is forked from the customScope
         expect(getCurrentScope()).not.toBe(initialScope);
-        expect(getCurrentScope()).toBe(manualScope);
+        expect(getCurrentScope()).not.toBe(customScope);
+
         expect(getActiveSpan()).toBe(span2);
+        // both, span1 and span2 are children of the parent span
         expect(spanToJSON(span2).parent_span_id).toBe('parent-span-id');
       });
 
-      withScope(manualScope, () => {
+      withScope(customScope, () => {
         expect(getActiveSpan()).toBe(parentSpan);
       });
 
@@ -305,31 +335,35 @@ describe('startSpan', () => {
 
     it('without parent span', () => {
       const initialScope = getCurrentScope();
-      const manualScope = initialScope.clone();
+      const customScope = initialScope.clone();
 
-      const traceId = manualScope.getPropagationContext()?.traceId;
+      const traceId = customScope.getPropagationContext()?.traceId;
 
-      startSpan({ name: 'span 1', scope: manualScope }, span1 => {
+      startSpan({ name: 'span 1', scope: customScope }, span1 => {
         expect(getCurrentScope()).not.toBe(initialScope);
-        expect(getCurrentScope()).toBe(manualScope);
+        expect(getCurrentScope()).not.toBe(customScope);
+
         expect(getActiveSpan()).toBe(span1);
-        expect(spanToJSON(span1).parent_span_id).toBe(undefined);
+        expect(getRootSpan(getActiveSpan()!)).toBe(span1);
+
         expect(span1.spanContext().traceId).toBe(traceId);
       });
 
-      withScope(manualScope, () => {
+      withScope(customScope, () => {
         expect(getActiveSpan()).toBe(undefined);
       });
 
-      startSpan({ name: 'span 2', scope: manualScope }, span2 => {
+      startSpan({ name: 'span 2', scope: customScope }, span2 => {
         expect(getCurrentScope()).not.toBe(initialScope);
-        expect(getCurrentScope()).toBe(manualScope);
+        expect(getCurrentScope()).not.toBe(customScope);
+
         expect(getActiveSpan()).toBe(span2);
-        expect(spanToJSON(span2).parent_span_id).toBe(undefined);
+        expect(getRootSpan(getActiveSpan()!)).toBe(span2);
+
         expect(span2.spanContext().traceId).toBe(traceId);
       });
 
-      withScope(manualScope, () => {
+      withScope(customScope, () => {
         expect(getActiveSpan()).toBe(undefined);
       });
 

packages/opentelemetry/test/trace.test.ts

@@ -290,24 +290,46 @@ describe('trace', () => {
       let manualScope: Scope;
       let parentSpan: Span;
 
+      // "hack" to create a manual scope with a parent span
       startSpanManual({ name: 'detached' }, span => {
         parentSpan = span;
         manualScope = getCurrentScope();
         manualScope.setTag('manual', 'tag');
       });
 
+      expect(manualScope!.getScopeData().tags).toEqual({ manual: 'tag' });
+      expect(getCurrentScope()).not.toBe(manualScope!);
+
       getCurrentScope().setTag('outer', 'tag');
 
       startSpan({ name: 'GET users/[id]', scope: manualScope! }, span => {
+        // the current scope in the callback is a fork of the manual scope
         expect(getCurrentScope()).not.toBe(initialScope);
+        expect(getCurrentScope()).not.toBe(manualScope);
+        expect(getCurrentScope().getScopeData().tags).toEqual({ manual: 'tag' });
 
-        expect(getCurrentScope()).toEqual(manualScope);
+        // getActiveSpan returns the correct span
         expect(getActiveSpan()).toBe(span);
 
+        // span hierarchy is correct
         expect(getSpanParentSpanId(span)).toBe(parentSpan.spanContext().spanId);
+
+        // scope data modifications are isolated between original and forked manual scope
+        getCurrentScope().setTag('inner', 'tag');
+        manualScope!.setTag('manual-scope-inner', 'tag');
+
+        expect(getCurrentScope().getScopeData().tags).toEqual({ manual: 'tag', inner: 'tag' });
+        expect(manualScope!.getScopeData().tags).toEqual({ manual: 'tag', 'manual-scope-inner': 'tag' });
       });
 
+      // manualScope modifications remain set outside the callback
+      expect(manualScope!.getScopeData().tags).toEqual({ manual: 'tag', 'manual-scope-inner': 'tag' });
+
+      // current scope is reset back to initial scope
       expect(getCurrentScope()).toBe(initialScope);
+      expect(getCurrentScope().getScopeData().tags).toEqual({ outer: 'tag' });
+
+      // although the manual span is still running, it's no longer active due to being outside of the callback
       expect(getActiveSpan()).toBe(undefined);
     });