Skip to content

Commit

Permalink
feat(core): Allow to pass scope to startSpan APIs
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
mydea committed Jan 8, 2024
1 parent 0a97514 commit 26d0440
Show file tree
Hide file tree
Showing 7 changed files with 200 additions and 43 deletions.
31 changes: 27 additions & 4 deletions packages/core/src/exports.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import type {
FinishedCheckIn,
MonitorConfig,
Primitive,
Scope as ScopeInterface,
Session,
SessionContext,
Severity,
Expand Down Expand Up @@ -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]);
}

/**
Expand Down
14 changes: 9 additions & 5 deletions packages/core/src/tracing/trace.ts
Original file line number Diff line number Diff line change
@@ -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';
Expand All @@ -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;
}

/**
Expand Down Expand Up @@ -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);
Expand Down Expand Up @@ -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();

Expand Down Expand Up @@ -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)
Expand Down
24 changes: 24 additions & 0 deletions packages/core/test/lib/exports.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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', () => {
Expand Down
120 changes: 90 additions & 30 deletions packages/core/test/lib/tracing/trace.test.ts
Original file line number Diff line number Diff line change
@@ -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(() => {
Expand Down Expand Up @@ -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 => {
Expand Down Expand Up @@ -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 => {
Expand All @@ -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);
});
});

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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);
Expand Down
36 changes: 33 additions & 3 deletions packages/node-experimental/src/sdk/api.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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()));
}

Expand Down
15 changes: 15 additions & 0 deletions packages/node-experimental/src/utils/contextData.ts
Original file line number Diff line number Diff line change
@@ -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.
Expand All @@ -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);
}
3 changes: 2 additions & 1 deletion packages/opentelemetry/src/types.ts
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -13,6 +13,7 @@ export interface OpenTelemetrySpanContext {
metadata?: Partial<TransactionMetadata>;
origin?: SpanOrigin;
source?: TransactionSource;
scope?: Scope;

// Base SpanOptions we support
attributes?: Attributes;
Expand Down

0 comments on commit 26d0440

Please sign in to comment.