feat(node-experimental): Add withActiveSpan (#10194)

packages/node-experimental/src/index.ts

@@ -32,6 +32,7 @@ export {
   setUser,
   withScope,
   withIsolationScope,
+  withActiveSpan,
   // eslint-disable-next-line deprecation/deprecation
   configureScope,
   getCurrentScope,

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

@@ -1,6 +1,7 @@
 // PUBLIC APIS
 
-import { context } from '@opentelemetry/api';
+import type { Span } from '@opentelemetry/api';
+import { context, trace } from '@opentelemetry/api';
 import type {
   Breadcrumb,
   BreadcrumbHint,
@@ -63,6 +64,18 @@ export function withScope<T>(
   return context.with(context.active(), () => callback(getCurrentScope()));
 }
 
+/**
+ * Forks the current scope and sets the provided span as active span in the context of the provided callback.
+ *
+ * @param span Spans started in the context of the provided callback will be children of this span.
+ * @param callback Execution context in which the provided span will be active. Is passed the newly forked scope.
+ * @returns the value returned from the provided callback function.
+ */
+export function withActiveSpan<T>(span: Span, callback: (scope: Scope) => T): T {
+  const newContextWithActiveSpan = trace.setSpan(context.active(), span);
+  return context.with(newContextWithActiveSpan, () => callback(getCurrentScope()));
+}
+
 /**
  * For a new isolation scope from the current isolation scope,
  * and make it the current isolation scope in the given callback.

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

@@ -0,0 +1,56 @@
+import { getActiveSpan, getClient, startInactiveSpan, startSpan, withActiveSpan } from '../../src';
+import { cleanupOtel, mockSdkInit } from '../helpers/mockSdkInit';
+
+afterEach(() => {
+  jest.restoreAllMocks();
+  cleanupOtel();
+});
+
+describe('withActiveSpan()', () => {
+  it('should set the active span within the callback', () => {
+    mockSdkInit();
+
+    const inactiveSpan = startInactiveSpan({ name: 'inactive-span' });
+
+    expect(getActiveSpan()).not.toBe(inactiveSpan);
+
+    withActiveSpan(inactiveSpan, () => {
+      expect(getActiveSpan()).toBe(inactiveSpan);
+    });
+  });
+
+  it('should create child spans when calling startSpan within the callback', async () => {
+    const beforeSendTransaction = jest.fn(() => null);
+    mockSdkInit({ enableTracing: true, beforeSendTransaction });
+    const client = getClient();
+
+    const inactiveSpan = startInactiveSpan({ name: 'inactive-span' });
+
+    withActiveSpan(inactiveSpan, () => {
+      startSpan({ name: 'child-span' }, () => {});
+    });
+
+    startSpan({ name: 'floating-span' }, () => {});
+
+    inactiveSpan.end();
+
+    await client.flush();
+
+    // The child span should be a child of the inactive span
+    expect(beforeSendTransaction).toHaveBeenCalledWith(
+      expect.objectContaining({
+        transaction: 'inactive-span',
+        spans: expect.arrayContaining([expect.objectContaining({ description: 'child-span' })]),
+      }),
+      expect.anything(),
+    );
+
+    // The floating span should be a separate transaction
+    expect(beforeSendTransaction).toHaveBeenCalledWith(
+      expect.objectContaining({
+        transaction: 'floating-span',
+      }),
+      expect.anything(),
+    );
+  });
+});