fix(core): Restore previously active span when passing a custom scope to startSpan

Author: Lms24

packages/core/src/tracing/trace.ts

@@ -16,7 +16,14 @@ import { propagationContextFromHeaders } from '../utils-hoist/tracing';
 import { handleCallbackErrors } from '../utils/handleCallbackErrors';
 import { hasTracingEnabled } from '../utils/hasTracingEnabled';
 import { _getSpanForScope, _setSpanForScope } from '../utils/spanOnScope';
-import { addChildSpanToSpan, getRootSpan, spanIsSampled, spanTimeInputToSeconds, spanToJSON } from '../utils/spanUtils';
+import {
+  addChildSpanToSpan,
+  getActiveSpan,
+  getRootSpan,
+  spanIsSampled,
+  spanTimeInputToSeconds,
+  spanToJSON,
+} from '../utils/spanUtils';
 import { freezeDscOnSpan, getDynamicSamplingContextFromSpan } from './dynamicSamplingContext';
 import { logSpanStart } from './logSpans';
 import { sampleSpan } from './sampling';
@@ -44,9 +51,11 @@ export function startSpan<T>(options: StartSpanOptions, callback: (span: Span) =
   }
 
   const spanArguments = parseSentrySpanArguments(options);
-  const { forceTransaction, parentSpan: customParentSpan } = options;
+  const { forceTransaction, parentSpan: customParentSpan, scope: customScope } = options;
 
-  return withScope(options.scope, () => {
+  const previousActiveSpanOnCustomScope = customScope && _getSpanForScope(customScope);
+
+  return withScope(customScope, () => {
     // If `options.parentSpan` is defined, we want to wrap the callback in `withActiveSpan`
     const wrapper = getActiveSpanWrapper<T>(customParentSpan);
 
@@ -75,7 +84,14 @@ export function startSpan<T>(options: StartSpanOptions, callback: (span: Span) =
             activeSpan.setStatus({ code: SPAN_STATUS_ERROR, message: 'internal_error' });
           }
         },
-        () => activeSpan.end(),
+        () => {
+          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/test/lib/tracing/trace.test.ts

@@ -269,6 +269,75 @@ describe('startSpan', () => {
     expect(getActiveSpan()).toBe(undefined);
   });
 
+  describe('handles multiple spans in sequence with a custom scope', () => {
+    it('with parent span', () => {
+      const initialScope = getCurrentScope();
+
+      const manualScope = initialScope.clone();
+      const parentSpan = new SentrySpan({ spanId: 'parent-span-id', sampled: true });
+      _setSpanForScope(manualScope, parentSpan);
+
+      startSpan({ name: 'span 1', scope: manualScope }, span1 => {
+        expect(getCurrentScope()).not.toBe(initialScope);
+        expect(getCurrentScope()).toBe(manualScope);
+        expect(getActiveSpan()).toBe(span1);
+        expect(spanToJSON(span1).parent_span_id).toBe('parent-span-id');
+      });
+
+      withScope(manualScope, () => {
+        expect(getActiveSpan()).toBe(parentSpan);
+      });
+
+      startSpan({ name: 'span 2', scope: manualScope }, span2 => {
+        expect(getCurrentScope()).not.toBe(initialScope);
+        expect(getCurrentScope()).toBe(manualScope);
+        expect(getActiveSpan()).toBe(span2);
+        expect(spanToJSON(span2).parent_span_id).toBe('parent-span-id');
+      });
+
+      withScope(manualScope, () => {
+        expect(getActiveSpan()).toBe(parentSpan);
+      });
+
+      expect(getCurrentScope()).toBe(initialScope);
+      expect(getActiveSpan()).toBe(undefined);
+    });
+
+    it('without parent span', () => {
+      const initialScope = getCurrentScope();
+      const manualScope = initialScope.clone();
+
+      const traceId = manualScope.getPropagationContext()?.traceId;
+
+      startSpan({ name: 'span 1', scope: manualScope }, span1 => {
+        expect(getCurrentScope()).not.toBe(initialScope);
+        expect(getCurrentScope()).toBe(manualScope);
+        expect(getActiveSpan()).toBe(span1);
+        expect(spanToJSON(span1).parent_span_id).toBe(undefined);
+        expect(span1.spanContext().traceId).toBe(traceId);
+      });
+
+      withScope(manualScope, () => {
+        expect(getActiveSpan()).toBe(undefined);
+      });
+
+      startSpan({ name: 'span 2', scope: manualScope }, span2 => {
+        expect(getCurrentScope()).not.toBe(initialScope);
+        expect(getCurrentScope()).toBe(manualScope);
+        expect(getActiveSpan()).toBe(span2);
+        expect(spanToJSON(span2).parent_span_id).toBe(undefined);
+        expect(span2.spanContext().traceId).toBe(traceId);
+      });
+
+      withScope(manualScope, () => {
+        expect(getActiveSpan()).toBe(undefined);
+      });
+
+      expect(getCurrentScope()).toBe(initialScope);
+      expect(getActiveSpan()).toBe(undefined);
+    });
+  });
+
   it('allows to pass a parentSpan', () => {
     const parentSpan = new SentrySpan({ spanId: 'parent-span-id', sampled: true, name: 'parent-span' });