feat(profiling) add global profile context while profiler is running (#12394)

Attach profile context to events while a profiler is active and assign
thread id and name to the trace context.

This will be required by the profiling API so that we can figure out the
links between spans and profiles. In the future (span streaming), tid
and thread name will be assigned to each span

Author: JonasBa

packages/profiling-node/src/integration.ts

@@ -1,17 +1,20 @@
 import { defineIntegration, getCurrentScope, getIsolationScope, getRootSpan, spanToJSON } from '@sentry/core';
 import type { NodeClient } from '@sentry/node';
-import type { Integration, IntegrationFn, Profile, ProfileChunk, Span } from '@sentry/types';
+import type { Event, Integration, IntegrationFn, Profile, ProfileChunk, Span } from '@sentry/types';
 
 import { LRUMap, logger, timestampInSeconds, uuid4 } from '@sentry/utils';
 
+import { getGlobalScope } from '../../core/src/currentScopes';
 import { CpuProfilerBindings } from './cpu_profiler';
 import { DEBUG_BUILD } from './debug-build';
 import { NODE_MAJOR, NODE_VERSION } from './nodeVersion';
 import { MAX_PROFILE_DURATION_MS, maybeProfileSpan, stopSpanProfile } from './spanProfileUtils';
-import type { RawThreadCpuProfile } from './types';
+import type { RawChunkCpuProfile, RawThreadCpuProfile } from './types';
 import { ProfileFormat } from './types';
+import { PROFILER_THREAD_NAME } from './utils';
 
 import {
+  PROFILER_THREAD_ID_STRING,
   addProfilesToEnvelope,
   createProfilingChunkEvent,
   createProfilingEvent,
@@ -211,7 +214,9 @@ class ContinuousProfiler {
         logger.log(`[Profiling] Failed to collect profile for: ${this._chunkData?.id}, the chunk_id is missing.`);
       return;
     }
-    const profile = CpuProfilerBindings.stopProfiling(this._chunkData.id, ProfileFormat.CHUNK);
+
+    const profile = this._stopChunkProfiling(this._chunkData);
+
     if (!profile || !this._chunkData.startTimestampMS) {
       DEBUG_BUILD && logger.log(`[Profiling] _chunkiledStartTraceID to collect profile for: ${this._chunkData.id}`);
       return;
@@ -274,12 +279,22 @@ class ContinuousProfiler {
     });
   }
 
+  /**
+   * Stops the profile and clears chunk instrumentation from global scope
+   * @returns void
+   */
+  private _stopChunkProfiling(chunk: ChunkData): RawChunkCpuProfile | null {
+    this._teardownSpanChunkInstrumentation();
+    return CpuProfilerBindings.stopProfiling(chunk.id, ProfileFormat.CHUNK);
+  }
+
   /**
    * Starts the profiler and registers the flush timer for a given chunk.
    * @param chunk
    */
   private _startChunkProfiling(chunk: ChunkData): void {
-    CpuProfilerBindings.startProfiling(chunk.id!);
+    this._setupSpanChunkInstrumentation();
+    CpuProfilerBindings.startProfiling(chunk.id);
     DEBUG_BUILD && logger.log(`[Profiling] starting profiling chunk: ${chunk.id}`);
 
     chunk.timer = global.setTimeout(() => {
@@ -293,6 +308,32 @@ class ContinuousProfiler {
     chunk.timer.unref();
   }
 
+  /**
+   * Attaches profiling information to spans that were started
+   * during a profiling session.
+   */
+  private _setupSpanChunkInstrumentation(): void {
+    if (!this._client) {
+      DEBUG_BUILD &&
+        logger.log('[Profiling] Failed to collect profile, sentry client was never attached to the profiler.');
+      return;
+    }
+
+    getGlobalScope().setContext('profile', {
+      profiler_id: this._profilerId,
+    });
+
+    this._client.on('beforeSendEvent', e => this._assignThreadIdContext(e));
+  }
+
+  /**
+   * Clear profiling information from global context when a profile is not running.
+   */
+  private _teardownSpanChunkInstrumentation(): void {
+    const globalScope = getGlobalScope();
+    globalScope.setContext('profile', {});
+  }
+
   /**
    * Initializes new profile chunk metadata
    */
@@ -305,6 +346,30 @@ class ContinuousProfiler {
     };
   }
 
+  /**
+   * Assigns thread_id and thread name context to a profiled event.
+   */
+  private _assignThreadIdContext(event: Event): any {
+    if (!event?.['contexts']?.['profile']) {
+      return;
+    }
+
+    if (!event.contexts) {
+      return;
+    }
+
+    // @ts-expect-error the trace fallback value is wrong, though it should never happen
+    // and in case it does, we dont want to override whatever was passed initially.
+    event.contexts['trace'] = {
+      ...(event.contexts?.['trace'] ?? {}),
+      data: {
+        ...(event.contexts?.['trace']?.['data'] ?? {}),
+        ['thread.id']: PROFILER_THREAD_ID_STRING,
+        ['thread.name']: PROFILER_THREAD_NAME,
+      },
+    };
+  }
+
   /**
    * Resets the current chunk state.
    */

packages/profiling-node/src/spanProfileUtils.ts

@@ -115,7 +115,6 @@ export function stopSpanProfile(span: Span, profile_id: string | undefined): Raw
   }
 
   const profile = CpuProfilerBindings.stopProfiling(profile_id, 0);
-
   DEBUG_BUILD && logger.log(`[Profiling] stopped profiling of transaction: ${spanToJSON(span).description}`);
 
   // In case of an overlapping span, stopProfiling may return null and silently ignore the overlapping profile.

packages/profiling-node/src/utils.ts

@@ -28,8 +28,8 @@ import type { RawChunkCpuProfile, RawThreadCpuProfile } from './types';
 
 // We require the file because if we import it, it will be included in the bundle.
 // I guess tsc does not check file contents when it's imported.
-const THREAD_ID_STRING = String(threadId);
-const THREAD_NAME = isMainThread ? 'main' : 'worker';
+export const PROFILER_THREAD_ID_STRING = String(threadId);
+export const PROFILER_THREAD_NAME = isMainThread ? 'main' : 'worker';
 const FORMAT_VERSION = '1';
 const CONTINUOUS_FORMAT_VERSION = '2';
 
@@ -75,8 +75,8 @@ export function enrichWithThreadInformation(
     frames: profile.frames,
     stacks: profile.stacks,
     thread_metadata: {
-      [THREAD_ID_STRING]: {
-        name: THREAD_NAME,
+      [PROFILER_THREAD_ID_STRING]: {
+        name: PROFILER_THREAD_NAME,
       },
     },
   } as ThreadCpuProfile | ContinuousThreadCpuProfile;
@@ -172,7 +172,7 @@ function createProfilePayload(
       name: transaction,
       id: event_id,
       trace_id: trace_id || '',
-      active_thread_id: THREAD_ID_STRING,
+      active_thread_id: PROFILER_THREAD_ID_STRING,
     },
   };
 

packages/profiling-node/test/spanProfileUtils.test.ts

@@ -7,8 +7,6 @@ import { GLOBAL_OBJ, createEnvelope, logger } from '@sentry/utils';
 import { CpuProfilerBindings } from '../src/cpu_profiler';
 import { type ProfilingIntegration, _nodeProfilingIntegration } from '../src/integration';
 
-jest.setTimeout(10000);
-
 function makeClientWithHooks(): [Sentry.NodeClient, Transport] {
   const integration = _nodeProfilingIntegration();
   const client = new Sentry.NodeClient({
@@ -322,7 +320,6 @@ describe('automated span instrumentation', () => {
     transaction.end();
     expect(stopProfilingSpy).toHaveBeenCalledTimes(1);
   });
-
   it('enriches profile with debug_id', async () => {
     GLOBAL_OBJ._sentryDebugIds = {
       'Error\n    at filename.js (filename.js:36:15)': 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaa',
@@ -574,6 +571,47 @@ describe('continuous profiling', () => {
 
     expect(transportSpy.mock.calls?.[0]?.[0]?.[1]?.[0]?.[0].type).toBe('profile_chunk');
   });
+
+  it('sets global profile context', async () => {
+    const [client, transport] = makeContinuousProfilingClient();
+    Sentry.setCurrentClient(client);
+    client.init();
+
+    const transportSpy = jest.spyOn(transport, 'send').mockReturnValue(Promise.resolve({}));
+
+    const nonProfiledTransaction = Sentry.startInactiveSpan({ forceTransaction: true, name: 'profile_hub' });
+    nonProfiledTransaction.end();
+
+    expect(transportSpy.mock.calls?.[0]?.[0]?.[1]?.[0]?.[1]).not.toMatchObject({
+      contexts: {
+        profile: {},
+      },
+    });
+
+    const integration = client.getIntegrationByName<ProfilingIntegration>('ProfilingIntegration');
+    if (!integration) {
+      throw new Error('Profiling integration not found');
+    }
+
+    integration._profiler.start();
+    const profiledTransaction = Sentry.startInactiveSpan({ forceTransaction: true, name: 'profile_hub' });
+    profiledTransaction.end();
+    integration._profiler.stop();
+
+    expect(transportSpy.mock.calls?.[1]?.[0]?.[1]?.[0]?.[1]).toMatchObject({
+      contexts: {
+        trace: {
+          data: expect.objectContaining({
+            ['thread.id']: expect.any(String),
+            ['thread.name']: expect.any(String),
+          }),
+        },
+        profile: {
+          profiler_id: expect.any(String),
+        },
+      },
+    });
+  });
 });
 
 describe('span profiling mode', () => {
@@ -610,7 +648,6 @@ describe('span profiling mode', () => {
     Sentry.startInactiveSpan({ forceTransaction: true, name: 'profile_hub' });
 
     expect(startProfilingSpy).toHaveBeenCalled();
-
     const integration = client.getIntegrationByName<ProfilingIntegration>('ProfilingIntegration');
 
     if (!integration) {

packages/profiling-node/test/spanProfileUtils.worker.test.ts

@@ -0,0 +1,75 @@
+// Mock the modules before the import, so that the value is initialized before the module is loaded
+jest.mock('worker_threads', () => {
+  return {
+    isMainThread: false,
+    threadId: 9999,
+  };
+});
+jest.setTimeout(10000);
+
+import * as Sentry from '@sentry/node';
+import type { Transport } from '@sentry/types';
+import { type ProfilingIntegration, _nodeProfilingIntegration } from '../src/integration';
+
+function makeContinuousProfilingClient(): [Sentry.NodeClient, Transport] {
+  const integration = _nodeProfilingIntegration();
+  const client = new Sentry.NodeClient({
+    stackParser: Sentry.defaultStackParser,
+    tracesSampleRate: 1,
+    profilesSampleRate: undefined,
+    debug: true,
+    environment: 'test-environment',
+    dsn: 'https://[email protected]/6625302',
+    integrations: [integration],
+    transport: _opts =>
+      Sentry.makeNodeTransport({
+        url: 'https://[email protected]/6625302',
+        recordDroppedEvent: () => {
+          return undefined;
+        },
+      }),
+  });
+
+  return [client, client.getTransport() as Transport];
+}
+
+it('worker threads context', () => {
+  const [client, transport] = makeContinuousProfilingClient();
+  Sentry.setCurrentClient(client);
+  client.init();
+
+  const transportSpy = jest.spyOn(transport, 'send').mockReturnValue(Promise.resolve({}));
+
+  const nonProfiledTransaction = Sentry.startInactiveSpan({ forceTransaction: true, name: 'profile_hub' });
+  nonProfiledTransaction.end();
+
+  expect(transportSpy.mock.calls?.[0]?.[0]?.[1]?.[0]?.[1]).not.toMatchObject({
+    contexts: {
+      profile: {},
+    },
+  });
+
+  const integration = client.getIntegrationByName<ProfilingIntegration>('ProfilingIntegration');
+  if (!integration) {
+    throw new Error('Profiling integration not found');
+  }
+
+  integration._profiler.start();
+  const profiledTransaction = Sentry.startInactiveSpan({ forceTransaction: true, name: 'profile_hub' });
+  profiledTransaction.end();
+  integration._profiler.stop();
+
+  expect(transportSpy.mock.calls?.[1]?.[0]?.[1]?.[0]?.[1]).toMatchObject({
+    contexts: {
+      trace: {
+        data: expect.objectContaining({
+          ['thread.id']: '9999',
+          ['thread.name']: 'worker',
+        }),
+      },
+      profile: {
+        profiler_id: expect.any(String),
+      },
+    },
+  });
+});