feat(profiling): add profilesSampler (#109)

* feat(profiling): add profilesSampler

* ref(utils): remove unnecessar bool check

* ref(profiling): disable support for distributed tracing

* ref(profiling): remove parentSampled

Author: JonasBa

package.json

@@ -44,7 +44,7 @@
     "benchmark:server": "node benchmarks/cpu/benchmark.server.js",
     "benchmark:format": "node benchmarks/format/benchmark.format.js",
     "benchmark:integration": "node benchmarks/cpu/benchmark.integration.base.js && node benchmarks/cpu/benchmark.integration.disabled.js && node benchmarks/cpu/benchmark.integration.js",
-    "test:watch": "jest --watch --silent",
+    "test:watch": "jest --watch",
     "test": "jest --config jest.config.ts",
     "prettier": "prettier --config ./.prettierrc --write"
   },

src/hubextensions.test.ts

@@ -3,6 +3,7 @@ import type {
   ClientOptions,
   Hub,
   Context,
+  Client,
   Transaction,
   TransactionMetadata
 } from '@sentry/types';
@@ -36,15 +37,22 @@ function makeTransactionMock(options = {}): Transaction {
   } as unknown as Transaction;
 }
 
-function makeHubMock({ profilesSampleRate }: { profilesSampleRate: number | undefined }): Hub {
+function makeHubMock({
+  profilesSampleRate,
+  client
+}: {
+  profilesSampleRate: number | undefined;
+  client?: Partial<Client<ClientOptions<BaseTransportOptions>>>;
+}): Hub {
   return {
     getClient: jest.fn().mockImplementation(() => {
       return {
         getOptions: jest.fn().mockImplementation(() => {
           return {
             profilesSampleRate
           } as unknown as ClientOptions<BaseTransportOptions>;
-        })
+        }),
+        ...(client ?? {})
       };
     })
   } as unknown as Hub;
@@ -131,4 +139,98 @@ describe('hubextensions', () => {
     expect(startProfilingSpy).not.toHaveBeenCalledTimes(1);
     expect(stopProfilingSpy).not.toHaveBeenCalledTimes(1);
   });
+
+  it('disabled if neither profilesSampler and profilesSampleRate are not set', () => {
+    const hub = makeHubMock({ profilesSampleRate: undefined });
+    const startTransaction = jest.fn().mockImplementation(() => makeTransactionMock());
+
+    const maybeStartTransactionWithProfiling = __PRIVATE__wrapStartTransactionWithProfiling(startTransaction);
+    const samplingContext = { beep: 'boop' };
+    const transaction = maybeStartTransactionWithProfiling.call(hub, { name: '' }, samplingContext);
+    transaction.finish();
+
+    const startProfilingSpy = jest.spyOn(profiler, 'startProfiling');
+    expect(startProfilingSpy).not.toHaveBeenCalled();
+  });
+
+  it('does not call startProfiling if profilesSampler returns invalid rate', () => {
+    const startProfilingSpy = jest.spyOn(profiler, 'startProfiling');
+    const options = { profilesSampler: jest.fn().mockReturnValue(NaN) };
+    const hub = makeHubMock({
+      profilesSampleRate: undefined,
+      client: {
+        // @ts-expect-error mock this
+        getOptions: () => options
+      }
+    });
+    const startTransaction = jest.fn().mockImplementation(() => makeTransactionMock());
+
+    const maybeStartTransactionWithProfiling = __PRIVATE__wrapStartTransactionWithProfiling(startTransaction);
+    const samplingContext = { beep: 'boop' };
+    const transaction = maybeStartTransactionWithProfiling.call(hub, { name: '' }, samplingContext);
+    transaction.finish();
+
+    expect(options.profilesSampler).toHaveBeenCalledWith(samplingContext);
+    expect(startProfilingSpy).not.toHaveBeenCalled();
+  });
+
+  it('does not call startProfiling if profilesSampleRate is invalid', () => {
+    const startProfilingSpy = jest.spyOn(profiler, 'startProfiling');
+    const options = { profilesSampler: jest.fn().mockReturnValue(NaN) };
+    const hub = makeHubMock({
+      profilesSampleRate: NaN,
+      client: {
+        // @ts-expect-error mock this
+        getOptions: () => options
+      }
+    });
+    const startTransaction = jest.fn().mockImplementation(() => makeTransactionMock());
+
+    const maybeStartTransactionWithProfiling = __PRIVATE__wrapStartTransactionWithProfiling(startTransaction);
+    const samplingContext = { beep: 'boop' };
+    const transaction = maybeStartTransactionWithProfiling.call(hub, { name: '' }, samplingContext);
+    transaction.finish();
+
+    expect(options.profilesSampler).toHaveBeenCalledWith(samplingContext);
+    expect(startProfilingSpy).not.toHaveBeenCalled();
+  });
+
+  it('calls profilesSampler with sampling context', () => {
+    const options = { profilesSampler: jest.fn() };
+    const hub = makeHubMock({
+      profilesSampleRate: undefined,
+      client: {
+        // @ts-expect-error mock this
+        getOptions: () => options
+      }
+    });
+    const startTransaction = jest.fn().mockImplementation(() => makeTransactionMock());
+
+    const maybeStartTransactionWithProfiling = __PRIVATE__wrapStartTransactionWithProfiling(startTransaction);
+    const samplingContext = { beep: 'boop' };
+    const transaction = maybeStartTransactionWithProfiling.call(hub, { name: '' }, samplingContext);
+    transaction.finish();
+
+    expect(options.profilesSampler).toHaveBeenCalledWith(samplingContext);
+  });
+
+  it('prioritizes profilesSampler outcome over profilesSampleRate', () => {
+    const startProfilingSpy = jest.spyOn(profiler, 'startProfiling');
+    const options = { profilesSampler: jest.fn().mockReturnValue(1) };
+    const hub = makeHubMock({
+      profilesSampleRate: 0,
+      client: {
+        // @ts-expect-error mock this
+        getOptions: () => options
+      }
+    });
+    const startTransaction = jest.fn().mockImplementation(() => makeTransactionMock());
+
+    const maybeStartTransactionWithProfiling = __PRIVATE__wrapStartTransactionWithProfiling(startTransaction);
+    const samplingContext = { beep: 'boop' };
+    const transaction = maybeStartTransactionWithProfiling.call(hub, { name: '' }, samplingContext);
+    transaction.finish();
+
+    expect(startProfilingSpy).toHaveBeenCalled();
+  });
 });

src/hubextensions.ts

@@ -4,6 +4,7 @@ import { logger, uuid4 } from '@sentry/utils';
 
 import { isDebugBuild } from './env';
 import { CpuProfilerBindings } from './cpu_profiler';
+import { isValidSampleRate } from './utils';
 
 const MAX_PROFILE_DURATION_MS = 30 * 1000;
 
@@ -30,26 +31,72 @@ export function __PRIVATE__wrapStartTransactionWithProfiling(startTransaction: S
     const profile_id = uuid4();
     const uniqueTransactionName = `${transactionContext.name} ${profile_id}`;
 
-    // profilesSampleRate is multiplied with tracesSampleRate to get the final sampling rate.
+    // profilesSampleRate is multiplied with tracesSampleRate to get the final sampling rate. We dont perform
+    // the actual multiplication to get the final rate, but we discard the profile if the transaction was sampled,
+    // so anything after this block from here is based on the transaction sampling.
     if (!transaction.sampled) {
       return transaction;
     }
 
-    // @ts-expect-error profilesSampleRate is not part of the options type yet
-    const profilesSampleRate = this.getClient()?.getOptions?.()?.profilesSampleRate;
-    if (profilesSampleRate === undefined) {
+    const client = this.getClient();
+    if (!client) {
+      if (isDebugBuild()) {
+        logger.log('[Profiling] Profiling disabled, no client found.');
+      }
+      return transaction;
+    }
+    const options = client.getOptions();
+    if (!options) {
+      if (isDebugBuild()) {
+        logger.log('[Profiling] Profiling disabled, no options found.');
+      }
+      return transaction;
+    }
+
+    // @ts-expect-error sampler is not yer part of the sdk options
+    const profilesSampler = options.profilesSampler;
+    // @ts-expect-error sampler is not yer part of the sdk options
+    let profilesSampleRate: number | undefined = options.profilesSampleRate;
+
+    // Prefer sampler to sample rate if both are provided.
+    if (typeof profilesSampler === 'function') {
+      profilesSampleRate = profilesSampler(customSamplingContext);
+    }
+
+    // Since this is coming from the user (or from a function provided by the user), who knows what we might get. (The
+    // only valid values are booleans or numbers between 0 and 1.)
+    if (!isValidSampleRate(profilesSampleRate)) {
+      if (isDebugBuild()) {
+        logger.warn('[Profiling] Discarding profile because of invalid sample rate.');
+      }
+      return transaction;
+    }
+
+    // if the function returned 0 (or false), or if `profileSampleRate` is 0, it's a sign the profile should be dropped
+    if (!profilesSampleRate) {
       if (isDebugBuild()) {
         logger.log(
-          '[Profiling] Profiling disabled, enable it by setting `profilesSampleRate` option to SDK init call.'
+          `[Profiling] Discarding profile because ${
+            typeof profilesSampler === 'function'
+              ? 'profileSampler returned 0 or false'
+              : 'a negative sampling decision was inherited or profileSampleRate is set to 0'
+          }`
         );
       }
       return transaction;
     }
 
+    // Now we roll the dice. Math.random is inclusive of 0, but not of 1, so strict < is safe here. In case sampleRate is
+    // a boolean, the < comparison will cause it to be automatically cast to 1 if it's true and 0 if it's false.
+    const sampled = Math.random() < (profilesSampleRate as number | boolean);
     // Check if we should sample this profile
-    if (Math.random() > profilesSampleRate) {
+    if (!sampled) {
       if (isDebugBuild()) {
-        logger.log('[Profiling] Skip profiling transaction due to sampling.');
+        logger.log(
+          `[Profiling] Discarding profile because it's not included in the random sample (sampling rate = ${Number(
+            profilesSampleRate
+          )})`
+        );
       }
       return transaction;
     }

src/utils.test.ts

@@ -1,5 +1,6 @@
 import type { SdkMetadata, DsnComponents } from '@sentry/types';
 import type { ProfiledEvent } from './utils';
+import { isValidSampleRate } from './utils';
 
 import {
   maybeRemoveProfileFromSdkMetadata,
@@ -255,3 +256,27 @@ describe('createProfilingEventEnvelope', () => {
     expect(profile.transaction.trace_id).toBe('trace_id');
   });
 });
+
+describe('isValidSampleRate', () => {
+  it.each([
+    [0, true],
+    [0.1, true],
+    [1, true],
+    [true, true],
+    [false, true],
+    // invalid values
+    [1.1, false],
+    [-0.1, false],
+    [NaN, false],
+    [Infinity, false],
+    [null, false],
+    [undefined, false],
+    ['', false],
+    [' ', false],
+    [{}, false],
+    [[], false],
+    [() => null, false]
+  ])('value %s is %s', (input, expected) => {
+    expect(isValidSampleRate(input)).toBe(expected);
+  });
+});

src/utils.ts

@@ -274,3 +274,29 @@ export function getProjectRootDirectory(): string | null {
   const components = path.resolve(root_directory).split('/node_modules');
   return components?.[0] ?? null;
 }
+
+/**
+ * Checks the given sample rate to make sure it is valid type and value (a boolean, or a number between 0 and 1).
+ */
+export function isValidSampleRate(rate: unknown): boolean {
+  // we need to check NaN explicitly because it's of type 'number' and therefore wouldn't get caught by this typecheck
+  if ((typeof rate !== 'number' && typeof rate !== 'boolean') || (typeof rate === 'number' && isNaN(rate))) {
+    if (isDebugBuild()) {
+      logger.warn(
+        `[Profiling] Invalid sample rate. Sample rate must be a boolean or a number between 0 and 1. Got ${JSON.stringify(
+          rate
+        )} of type ${JSON.stringify(typeof rate)}.`
+      );
+    }
+    return false;
+  }
+
+  // in case sampleRate is a boolean, it will get automatically cast to 1 if it's true and 0 if it's false
+  if (rate < 0 || rate > 1) {
+    if (isDebugBuild()) {
+      logger.warn(`[Profiling] Invalid sample rate. Sample rate must be between 0 and 1. Got ${rate}.`);
+    }
+    return false;
+  }
+  return true;
+}