feat(core): Ensure startSpan & startSpanManual fork scope

Author: mydea

packages/core/src/tracing/trace.ts

@@ -2,7 +2,7 @@ import type { TransactionContext } from '@sentry/types';
 import { dropUndefinedKeys, isThenable, logger, tracingContextFromHeaders } from '@sentry/utils';
 
 import { DEBUG_BUILD } from '../debug-build';
-import { getCurrentScope } from '../exports';
+import { getCurrentScope, withScope } from '../exports';
 import type { Hub } from '../hub';
 import { getCurrentHub } from '../hub';
 import { hasTracingEnabled } from '../utils/hasTracingEnabled';
@@ -89,42 +89,42 @@ export function trace<T>(
 export function startSpan<T>(context: TransactionContext, callback: (span: Span | undefined) => T): T {
   const ctx = normalizeContext(context);
 
-  const hub = getCurrentHub();
-  const scope = getCurrentScope();
-  const parentSpan = scope.getSpan();
+  return withScope(scope => {
+    const hub = getCurrentHub();
+    const parentSpan = scope.getSpan();
 
-  const activeSpan = createChildSpanOrTransaction(hub, parentSpan, ctx);
-  scope.setSpan(activeSpan);
-
-  function finishAndSetSpan(): void {
-    activeSpan && activeSpan.finish();
-    scope.setSpan(parentSpan);
-  }
+    const activeSpan = createChildSpanOrTransaction(hub, parentSpan, ctx);
+    scope.setSpan(activeSpan);
 
-  let maybePromiseResult: T;
-  try {
-    maybePromiseResult = callback(activeSpan);
-  } catch (e) {
-    activeSpan && activeSpan.setStatus('internal_error');
-    finishAndSetSpan();
-    throw e;
-  }
+    function finishAndSetSpan(): void {
+      activeSpan && activeSpan.finish();
+    }
 
-  if (isThenable(maybePromiseResult)) {
-    Promise.resolve(maybePromiseResult).then(
-      () => {
-        finishAndSetSpan();
-      },
-      () => {
-        activeSpan && activeSpan.setStatus('internal_error');
-        finishAndSetSpan();
-      },
-    );
-  } else {
-    finishAndSetSpan();
-  }
-
-  return maybePromiseResult;
+    let maybePromiseResult: T;
+    try {
+      maybePromiseResult = callback(activeSpan);
+    } catch (e) {
+      activeSpan && activeSpan.setStatus('internal_error');
+      finishAndSetSpan();
+      throw e;
+    }
+
+    if (isThenable(maybePromiseResult)) {
+      Promise.resolve(maybePromiseResult).then(
+        () => {
+          finishAndSetSpan();
+        },
+        () => {
+          activeSpan && activeSpan.setStatus('internal_error');
+          finishAndSetSpan();
+        },
+      );
+    } else {
+      finishAndSetSpan();
+    }
+
+    return maybePromiseResult;
+  });
 }
 
 /**
@@ -149,33 +149,33 @@ export function startSpanManual<T>(
 ): T {
   const ctx = normalizeContext(context);
 
-  const hub = getCurrentHub();
-  const scope = getCurrentScope();
-  const parentSpan = scope.getSpan();
+  return withScope(scope => {
+    const hub = getCurrentHub();
+    const parentSpan = scope.getSpan();
 
-  const activeSpan = createChildSpanOrTransaction(hub, parentSpan, ctx);
-  scope.setSpan(activeSpan);
+    const activeSpan = createChildSpanOrTransaction(hub, parentSpan, ctx);
+    scope.setSpan(activeSpan);
 
-  function finishAndSetSpan(): void {
-    activeSpan && activeSpan.finish();
-    scope.setSpan(parentSpan);
-  }
+    function finishAndSetSpan(): void {
+      activeSpan && activeSpan.finish();
+    }
 
-  let maybePromiseResult: T;
-  try {
-    maybePromiseResult = callback(activeSpan, finishAndSetSpan);
-  } catch (e) {
-    activeSpan && activeSpan.setStatus('internal_error');
-    throw e;
-  }
-
-  if (isThenable(maybePromiseResult)) {
-    Promise.resolve(maybePromiseResult).then(undefined, () => {
+    let maybePromiseResult: T;
+    try {
+      maybePromiseResult = callback(activeSpan, finishAndSetSpan);
+    } catch (e) {
       activeSpan && activeSpan.setStatus('internal_error');
-    });
-  }
+      throw e;
+    }
 
-  return maybePromiseResult;
+    if (isThenable(maybePromiseResult)) {
+      Promise.resolve(maybePromiseResult).then(undefined, () => {
+        activeSpan && activeSpan.setStatus('internal_error');
+      });
+    }
+
+    return maybePromiseResult;
+  });
 }
 
 /**

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

@@ -1,5 +1,6 @@
-import { Hub, addTracingExtensions, makeMain } from '../../../src';
-import { continueTrace, startSpan } from '../../../src/tracing';
+import type { Span } from '@sentry/types';
+import { Hub, addTracingExtensions, getCurrentScope, makeMain } from '../../../src';
+import { continueTrace, startInactiveSpan, startSpan, startSpanManual } from '../../../src/tracing';
 import { TestClient, getDefaultTestClientOptions } from '../../mocks/client';
 
 beforeAll(() => {
@@ -80,6 +81,18 @@ describe('startSpan', () => {
       expect(ref.status).toEqual(isError ? 'internal_error' : undefined);
     });
 
+    it('creates & finishes span', async () => {
+      let _span: Span | undefined;
+      startSpan({ name: 'GET users/[id]' }, span => {
+        expect(span).toBeDefined();
+        expect(span?.endTimestamp).toBeUndefined();
+        _span = span;
+      });
+
+      expect(_span).toBeDefined();
+      expect(_span?.endTimestamp).toBeDefined();
+    });
+
     it('allows traceparent information to be overriden', async () => {
       let ref: any = undefined;
       client.on('finishTransaction', transaction => {
@@ -168,6 +181,71 @@ describe('startSpan', () => {
       expect(ref.spanRecorder.spans).toHaveLength(2);
       expect(ref.spanRecorder.spans[1].op).toEqual('db.query');
     });
+
+    it('forks the scope', () => {
+      const initialScope = getCurrentScope();
+
+      startSpan({ name: 'GET users/[id]' }, span => {
+        expect(getCurrentScope()).not.toBe(initialScope);
+        expect(getCurrentScope().getSpan()).toBe(span);
+      });
+
+      expect(getCurrentScope()).toBe(initialScope);
+      expect(initialScope.getSpan()).toBe(undefined);
+    });
+  });
+});
+
+describe('startSpanManual', () => {
+  it('creates & finishes span', async () => {
+    startSpanManual({ name: 'GET users/[id]' }, (span, finish) => {
+      expect(span).toBeDefined();
+      expect(span?.endTimestamp).toBeUndefined();
+      finish();
+      expect(span?.endTimestamp).toBeDefined();
+    });
+  });
+
+  it('forks the scope automatically', () => {
+    const initialScope = getCurrentScope();
+
+    startSpanManual({ name: 'GET users/[id]' }, (span, finish) => {
+      expect(getCurrentScope()).not.toBe(initialScope);
+      expect(getCurrentScope().getSpan()).toBe(span);
+
+      finish();
+
+      // Is still the active span
+      expect(getCurrentScope().getSpan()).toBe(span);
+    });
+
+    expect(getCurrentScope()).toBe(initialScope);
+    expect(initialScope.getSpan()).toBe(undefined);
+  });
+});
+
+describe('startInactiveSpan', () => {
+  it('creates & finishes span', async () => {
+    const span = startInactiveSpan({ name: 'GET users/[id]' });
+
+    expect(span).toBeDefined();
+    expect(span?.endTimestamp).toBeUndefined();
+
+    span?.finish();
+
+    expect(span?.endTimestamp).toBeDefined();
+  });
+
+  it('does not set span on scope', () => {
+    const initialScope = getCurrentScope();
+
+    const span = startInactiveSpan({ name: 'GET users/[id]' });
+
+    expect(initialScope.getSpan()).toBeUndefined();
+
+    span?.finish();
+
+    expect(initialScope.getSpan()).toBeUndefined();
   });
 });