feat(core): Allow to pass scope to startSpan APIs (#10076)

Also allow to pass this to `withScope()`, which actually aligns with
OTEL as well, where you can also do that. If you pass a scope there, it
will use this instead of forking a new one.

This should be the last thing needed to refactor some
`span.startChild()` occurrences - you can now store the scope you want
to fork off, and pass this into `startSpan` as needed.

Author: mydea

packages/core/src/exports.ts

@@ -12,6 +12,7 @@ import type {
   FinishedCheckIn,
   MonitorConfig,
   Primitive,
+  Scope as ScopeInterface,
   Session,
   SessionContext,
   Severity,
@@ -167,11 +168,33 @@ export function setUser(user: User | null): ReturnType<Hub['setUser']> {
  *     pushScope();
  *     callback();
  *     popScope();
- *
- * @param callback that will be enclosed into push/popScope.
  */
-export function withScope<T>(callback: (scope: Scope) => T): T {
-  return getCurrentHub().withScope(callback);
+export function withScope<T>(callback: (scope: Scope) => T): T;
+/**
+ * Set the given scope as the active scope in the callback.
+ */
+export function withScope<T>(scope: ScopeInterface | undefined, callback: (scope: Scope) => T): T;
+/**
+ * Either creates a new active scope, or sets the given scope as active scope in the given callback.
+ */
+export function withScope<T>(
+  ...rest: [callback: (scope: Scope) => T] | [scope: ScopeInterface | undefined, callback: (scope: Scope) => T]
+): T {
+  // If a scope is defined, we want to make this the active scope instead of the default one
+  if (rest.length === 2) {
+    const [scope, callback] = rest;
+    if (!scope) {
+      return getCurrentHub().withScope(callback);
+    }
+
+    const hub = getCurrentHub();
+    return hub.withScope(() => {
+      hub.getStackTop().scope = scope as Scope;
+      return callback(scope as Scope);
+    });
+  }
+
+  return getCurrentHub().withScope(rest[0]);
 }
 
 /**

packages/core/src/tracing/trace.ts

@@ -1,4 +1,4 @@
-import type { Span, SpanTimeInput, TransactionContext } from '@sentry/types';
+import type { Scope, Span, SpanTimeInput, TransactionContext } from '@sentry/types';
 import { dropUndefinedKeys, logger, tracingContextFromHeaders } from '@sentry/utils';
 
 import { DEBUG_BUILD } from '../debug-build';
@@ -12,6 +12,9 @@ import { spanTimeInputToSeconds } from '../utils/spanUtils';
 interface StartSpanOptions extends TransactionContext {
   /** 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. */
+  scope?: Scope;
 }
 
 /**
@@ -74,9 +77,10 @@ export function trace<T>(
 export function startSpan<T>(context: StartSpanOptions, callback: (span: Span | undefined) => T): T {
   const ctx = normalizeContext(context);
 
-  return withScope(scope => {
+  return withScope(context.scope, scope => {
     const hub = getCurrentHub();
-    const parentSpan = scope.getSpan();
+    const scopeForSpan = context.scope || scope;
+    const parentSpan = scopeForSpan.getSpan();
 
     const activeSpan = createChildSpanOrTransaction(hub, parentSpan, ctx);
     scope.setSpan(activeSpan);
@@ -116,7 +120,7 @@ export function startSpanManual<T>(
 ): T {
   const ctx = normalizeContext(context);
 
-  return withScope(scope => {
+  return withScope(context.scope, scope => {
     const hub = getCurrentHub();
     const parentSpan = scope.getSpan();
 
@@ -156,7 +160,7 @@ export function startInactiveSpan(context: StartSpanOptions): Span | undefined {
 
   const ctx = normalizeContext(context);
   const hub = getCurrentHub();
-  const parentSpan = getActiveSpan();
+  const parentSpan = context.scope ? context.scope.getSpan() : getActiveSpan();
   return parentSpan
     ? // eslint-disable-next-line deprecation/deprecation
       parentSpan.startChild(ctx)

packages/core/test/lib/exports.test.ts

@@ -134,6 +134,30 @@ describe('withScope', () => {
 
     expect(getCurrentScope()).toBe(scope1);
   });
+
+  it('allows to pass a custom scope', () => {
+    const scope1 = getCurrentScope();
+    scope1.setExtra('x1', 'x1');
+
+    const customScope = new Scope();
+    customScope.setExtra('x2', 'x2');
+
+    withScope(customScope, scope2 => {
+      expect(scope2).not.toBe(scope1);
+      expect(scope2).toBe(customScope);
+      expect(getCurrentScope()).toBe(scope2);
+      expect(scope2['_extra']).toEqual({ x2: 'x2' });
+    });
+
+    withScope(customScope, scope3 => {
+      expect(scope3).not.toBe(scope1);
+      expect(scope3).toBe(customScope);
+      expect(getCurrentScope()).toBe(scope3);
+      expect(scope3['_extra']).toEqual({ x2: 'x2' });
+    });
+
+    expect(getCurrentScope()).toBe(scope1);
+  });
 });
 
 describe('session APIs', () => {

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

@@ -1,6 +1,6 @@
-import type { Span } from '@sentry/types';
 import { Hub, addTracingExtensions, getCurrentScope, makeMain } from '../../../src';
-import { continueTrace, startInactiveSpan, startSpan, startSpanManual } from '../../../src/tracing';
+import { Scope } from '../../../src/scope';
+import { Span, continueTrace, startInactiveSpan, startSpan, startSpanManual } from '../../../src/tracing';
 import { TestClient, getDefaultTestClientOptions } from '../../mocks/client';
 
 beforeAll(() => {
@@ -81,18 +81,6 @@ 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 => {
@@ -160,14 +148,6 @@ describe('startSpan', () => {
       expect(ref.spanRecorder.spans[1].status).toEqual(isError ? 'internal_error' : undefined);
     });
 
-    it('allows to pass a `startTime`', () => {
-      const start = startSpan({ name: 'outer', startTime: [1234, 0] }, span => {
-        return span?.startTimestamp;
-      });
-
-      expect(start).toEqual(1234);
-    });
-
     it('allows for span to be mutated', async () => {
       let ref: any = undefined;
       client.on('finishTransaction', transaction => {
@@ -189,18 +169,57 @@ describe('startSpan', () => {
       expect(ref.spanRecorder.spans).toHaveLength(2);
       expect(ref.spanRecorder.spans[1].op).toEqual('db.query');
     });
+  });
 
-    it('forks the scope', () => {
-      const initialScope = getCurrentScope();
+  it('creates & finishes span', async () => {
+    let _span: Span | undefined;
+    startSpan({ name: 'GET users/[id]' }, span => {
+      expect(span).toBeDefined();
+      expect(span?.endTimestamp).toBeUndefined();
+      _span = span as Span;
+    });
 
-      startSpan({ name: 'GET users/[id]' }, span => {
-        expect(getCurrentScope()).not.toBe(initialScope);
-        expect(getCurrentScope().getSpan()).toBe(span);
-      });
+    expect(_span).toBeDefined();
+    expect(_span?.endTimestamp).toBeDefined();
+  });
 
-      expect(getCurrentScope()).toBe(initialScope);
-      expect(initialScope.getSpan()).toBe(undefined);
+  it('allows to pass a `startTime`', () => {
+    const start = startSpan({ name: 'outer', startTime: [1234, 0] }, span => {
+      return span?.startTimestamp;
     });
+
+    expect(start).toEqual(1234);
+  });
+
+  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);
+  });
+
+  it('allows to pass a scope', () => {
+    const initialScope = getCurrentScope();
+
+    const manualScope = new Scope();
+    const parentSpan = new Span({ spanId: 'parent-span-id' });
+    manualScope.setSpan(parentSpan);
+
+    startSpan({ name: 'GET users/[id]', scope: manualScope }, span => {
+      expect(getCurrentScope()).not.toBe(initialScope);
+      expect(getCurrentScope()).toBe(manualScope);
+      expect(getCurrentScope().getSpan()).toBe(span);
+
+      expect(span?.parentSpanId).toBe('parent-span-id');
+    });
+
+    expect(getCurrentScope()).toBe(initialScope);
+    expect(initialScope.getSpan()).toBe(undefined);
   });
 });
 
@@ -231,6 +250,29 @@ describe('startSpanManual', () => {
     expect(initialScope.getSpan()).toBe(undefined);
   });
 
+  it('allows to pass a scope', () => {
+    const initialScope = getCurrentScope();
+
+    const manualScope = new Scope();
+    const parentSpan = new Span({ spanId: 'parent-span-id' });
+    manualScope.setSpan(parentSpan);
+
+    startSpanManual({ name: 'GET users/[id]', scope: manualScope }, (span, finish) => {
+      expect(getCurrentScope()).not.toBe(initialScope);
+      expect(getCurrentScope()).toBe(manualScope);
+      expect(getCurrentScope().getSpan()).toBe(span);
+      expect(span?.parentSpanId).toBe('parent-span-id');
+
+      finish();
+
+      // Is still the active span
+      expect(getCurrentScope().getSpan()).toBe(span);
+    });
+
+    expect(getCurrentScope()).toBe(initialScope);
+    expect(initialScope.getSpan()).toBe(undefined);
+  });
+
   it('allows to pass a `startTime`', () => {
     const start = startSpanManual({ name: 'outer', startTime: [1234, 0] }, span => {
       span?.end();
@@ -266,6 +308,24 @@ describe('startInactiveSpan', () => {
     expect(initialScope.getSpan()).toBeUndefined();
   });
 
+  it('allows to pass a scope', () => {
+    const initialScope = getCurrentScope();
+
+    const manualScope = new Scope();
+    const parentSpan = new Span({ spanId: 'parent-span-id' });
+    manualScope.setSpan(parentSpan);
+
+    const span = startInactiveSpan({ name: 'GET users/[id]', scope: manualScope });
+
+    expect(span).toBeDefined();
+    expect(span?.parentSpanId).toBe('parent-span-id');
+    expect(initialScope.getSpan()).toBeUndefined();
+
+    span?.end();
+
+    expect(initialScope.getSpan()).toBeUndefined();
+  });
+
   it('allows to pass a `startTime`', () => {
     const span = startInactiveSpan({ name: 'outer', startTime: [1234, 0] });
     expect(span?.startTimestamp).toEqual(1234);

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

@@ -16,7 +16,7 @@ import type {
   User,
 } from '@sentry/types';
 import { consoleSandbox, dateTimestampInSeconds } from '@sentry/utils';
-import { getScopesFromContext, setScopesOnContext } from '../utils/contextData';
+import { getContextFromScope, getScopesFromContext, setScopesOnContext } from '../utils/contextData';
 
 import type { ExclusiveEventHintOrCaptureContext } from '../utils/prepareEvent';
 import { parseEventHintOrCaptureContext } from '../utils/prepareEvent';
@@ -27,9 +27,39 @@ export { getCurrentScope, getGlobalScope, getIsolationScope, getClient };
 export { setCurrentScope, setIsolationScope } from './scope';
 
 /**
- * Fork a scope from the current scope, and make it the current scope in the given callback
+ * Creates a new scope with and executes the given operation within.
+ * The scope is automatically removed once the operation
+ * finishes or throws.
+ *
+ * This is essentially a convenience function for:
+ *
+ *     pushScope();
+ *     callback();
+ *     popScope();
  */
-export function withScope<T>(callback: (scope: Scope) => T): T {
+export function withScope<T>(callback: (scope: Scope) => T): T;
+/**
+ * Set the given scope as the active scope in the callback.
+ */
+export function withScope<T>(scope: Scope | undefined, callback: (scope: Scope) => T): T;
+/**
+ * Either creates a new active scope, or sets the given scope as active scope in the given callback.
+ */
+export function withScope<T>(
+  ...rest: [callback: (scope: Scope) => T] | [scope: Scope | undefined, callback: (scope: Scope) => T]
+): T {
+  // If a scope is defined, we want to make this the active scope instead of the default one
+  if (rest.length === 2) {
+    const [scope, callback] = rest;
+    if (!scope) {
+      return context.with(context.active(), () => callback(getCurrentScope()));
+    }
+
+    const ctx = getContextFromScope(scope);
+    return context.with(ctx || context.active(), () => callback(getCurrentScope()));
+  }
+
+  const callback = rest[0];
   return context.with(context.active(), () => callback(getCurrentScope()));
 }
 

packages/node-experimental/src/utils/contextData.ts

@@ -1,10 +1,13 @@
 import type { Context } from '@opentelemetry/api';
 import { createContextKey } from '@opentelemetry/api';
+import type { Scope } from '@sentry/types';
 
 import type { CurrentScopes } from '../sdk/types';
 
 export const SENTRY_SCOPES_CONTEXT_KEY = createContextKey('sentry_scopes');
 
+const SCOPE_CONTEXT_MAP = new WeakMap<Scope, Context>();
+
 /**
  * Try to get the current scopes from the given OTEL context.
  * This requires a Context Manager that was wrapped with getWrappedContextManager.
@@ -18,5 +21,17 @@ export function getScopesFromContext(context: Context): CurrentScopes | undefine
  * This will return a forked context with the Propagation Context set.
  */
 export function setScopesOnContext(context: Context, scopes: CurrentScopes): Context {
+  // So we can look up the context from the scope later
+  SCOPE_CONTEXT_MAP.set(scopes.scope, context);
+  SCOPE_CONTEXT_MAP.set(scopes.isolationScope, context);
+
   return context.setValue(SENTRY_SCOPES_CONTEXT_KEY, scopes);
 }
+
+/**
+ * Get the context related to a scope.
+ * TODO v8: Use this for the `trace` functions.
+ * */
+export function getContextFromScope(scope: Scope): Context | undefined {
+  return SCOPE_CONTEXT_MAP.get(scope);
+}

packages/opentelemetry/src/types.ts

@@ -1,6 +1,6 @@
 import type { Attributes, Span as WriteableSpan, SpanKind, TimeInput, Tracer } from '@opentelemetry/api';
 import type { BasicTracerProvider, ReadableSpan, Span } from '@opentelemetry/sdk-trace-base';
-import type { SpanOrigin, TransactionMetadata, TransactionSource } from '@sentry/types';
+import type { Scope, SpanOrigin, TransactionMetadata, TransactionSource } from '@sentry/types';
 
 export interface OpenTelemetryClient {
   tracer: Tracer;
@@ -13,6 +13,7 @@ export interface OpenTelemetrySpanContext {
   metadata?: Partial<TransactionMetadata>;
   origin?: SpanOrigin;
   source?: TransactionSource;
+  scope?: Scope;
 
   // Base SpanOptions we support
   attributes?: Attributes;