make tracesSampler able to return a boolean and make it override sampling inheritance

lobsterkatie · lobsterkatie · commit c7a44b9fa78e · 2020-09-03T20:12:08.000-07:00
diff --git a/packages/tracing/src/hubextensions.ts b/packages/tracing/src/hubextensions.ts
@@ -29,16 +29,28 @@ function traceHeaders(this: Hub): { [key: string]: string } {
 }
 
 /**
- * Use sample rate along with a random number generator to make a sampling decision, which all child spans and child
- * transactions inherit.
+ * Implements sampling inheritance and falls back to user-provided static rate if no parent decision is available.
  *
- * Sample rate is set in SDK config, either as a constant (`tracesSampleRate`) or a function to compute the rate
- * (`tracesSampler`).
+ * @param parentSampled: The parent transaction's sampling decision, if any.
+ * @param givenRate: The rate to use if no parental decision is available.
+ *
+ * @returns The parent's sampling decision (if one exists), or the provided static rate
+ */
+function _inheritOrUseGivenRate(parentSampled: boolean | undefined, givenRate: unknown): boolean | unknown {
+  return parentSampled !== undefined ? parentSampled : givenRate;
+}
+
+/**
+ * Makes a sampling decision for the given transaction and stores it on the transaction.
  *
  * Called every time a transaction is created. Only transactions which emerge with a `sampled` value of `true` will be
  * sent to Sentry.
  *
- * Mutates the given Transaction object and then returns the mutated object.
+ * @param hub: The hub off of which to read config options
+ * @param transaction: The transaction needing a sampling decision
+ * @param sampleContext: Default and user-provided data which may be used to help make the decision
+ *
+ * @returns The given transaction with its `sampled` value set
  */
 function sample<T extends Transaction>(hub: Hub, transaction: T, sampleContext: SampleContext = {}): T {
   const client = hub.getClient();
@@ -50,42 +62,46 @@ function sample<T extends Transaction>(hub: Hub, transaction: T, sampleContext:
     return transaction;
   }
 
-  // we have to test for a pre-existsing sampling decision, in case this transaction is a child transaction and has
-  // inherited its parent's decision
-  if (transaction.sampled === undefined) {
-    // we would have bailed at the beginning if neither `tracesSampler` nor `tracesSampleRate` were defined, so one of
-    // these should work; prefer the hook if so
-    const sampleRate =
-      typeof options.tracesSampler === 'function' ? options.tracesSampler(sampleContext) : options.tracesSampleRate;
-
-    // since this is coming from the user, who knows what we might get
-    if (!isValidSampleRate(sampleRate)) {
-      logger.warn(`[Tracing] Discarding trace because of invalid sample rate.`);
-      transaction.sampled = false;
-      return transaction;
-    }
+  // we would have bailed already if neither `tracesSampler` nor `tracesSampleRate` were defined, so one of these should
+  // work; prefer the hook if so
+  const sampleRate =
+    typeof options.tracesSampler === 'function'
+      ? options.tracesSampler(sampleContext)
+      : _inheritOrUseGivenRate(sampleContext.parentSampled, options.tracesSampleRate);
+
+  // 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(sampleRate)) {
+    logger.warn(`[Tracing] Discarding transaction because of invalid sample rate.`);
+    transaction.sampled = false;
+    return transaction;
+  }
 
-    // if the function returned 0, or if the sample rate is set to 0, it's a sign the transaction should be dropped
-    if (!sampleRate) {
-      logger.log(
-        `[Tracing] Discarding trace because ${
-          typeof options.tracesSampler === 'function' ? 'tracesSampler returned 0' : 'tracesSampleRate is set to 0'
-        }`,
-      );
-      transaction.sampled = false;
-      return transaction;
-    }
+  // if the function returned 0 (or false), or if `tracesSampleRate` is 0, it's a sign the transaction should be dropped
+  if (!sampleRate) {
+    logger.log(
+      `[Tracing] Discarding transaction because ${
+        typeof options.tracesSampler === 'function'
+          ? 'tracesSampler returned 0 or false'
+          : 'tracesSampleRate is set to 0'
+      }`,
+    );
+    transaction.sampled = false;
+    return transaction;
+  }
 
-    // now we roll the dice (Math.random is inclusive of 0, but not of 1, so strict < is safe here)
-    transaction.sampled = Math.random() < sampleRate;
+  // 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.
+  transaction.sampled = Math.random() < (sampleRate as number | boolean);
 
-    // if we're not going to keep it, we're done
-    if (!transaction.sampled) {
-      logger.log(
-        `[Tracing] Discarding trace because it's not included in the random sample (sampling rate = ${sampleRate})`,
-      );
-      return transaction;
-    }
+  // if we're not going to keep it, we're done
+  if (!transaction.sampled) {
+    logger.log(
+      `[Tracing] Discarding transaction because it's not included in the random sample (sampling rate = ${Number(
+        sampleRate,
+      )})`,
+    );
+    return transaction;
   }
 
   // at this point we know we're keeping the transaction, whether because of an inherited decision or because it got
@@ -148,17 +164,19 @@ function getDefaultSampleContext<T extends Transaction>(transaction: T): SampleC
 }
 
 /**
- * Checks the given sample rate to make sure it is valid (a number between 0 and 1).
+ * Checks the given sample rate to make sure it is valid type and value (a boolean, or a number between 0 and 1).
  */
 function isValidSampleRate(rate: unknown): boolean {
-  if (!(typeof rate === 'number')) {
+  if (!(typeof rate === 'number' || typeof rate === 'boolean')) {
     logger.warn(
-      `[Tracing] Given sample rate is invalid. Sample rate must be a number between 0 and 1. Got ${JSON.stringify(
+      `[Tracing] Given sample rate is invalid. 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) {
     logger.warn(`[Tracing] Given sample rate is invalid. Sample rate must be between 0 and 1. Got ${rate}.`);
     return false;
diff --git a/packages/tracing/test/hub.test.ts b/packages/tracing/test/hub.test.ts
@@ -41,7 +41,7 @@ describe('Hub', () => {
   });
 
   describe('transaction sampling', () => {
-    describe('options', () => {
+    describe('tracesSampleRate and tracesSampler options', () => {
       it("should call tracesSampler if it's defined", () => {
         const tracesSampler = jest.fn();
         const hub = new Hub(new BrowserClient({ tracesSampler }));
@@ -52,11 +52,20 @@ describe('Hub', () => {
 
       it('should prefer tracesSampler to tracesSampleRate', () => {
         const tracesSampler = jest.fn();
-        const hub = new Hub(new BrowserClient({ tracesSampleRate: 1, tracesSampler: tracesSampler }));
+        const hub = new Hub(new BrowserClient({ tracesSampleRate: 1, tracesSampler }));
         hub.startTransaction({ name: 'dogpark' });
 
         expect(tracesSampler).toHaveBeenCalled();
       });
+
+      it('tolerates tracesSampler returning a boolean', () => {
+        const tracesSampler = jest.fn().mockReturnValue(true);
+        const hub = new Hub(new BrowserClient({ tracesSampler }));
+        const transaction = hub.startTransaction({ name: 'dogpark' });
+
+        expect(tracesSampler).toHaveBeenCalled();
+        expect(transaction.sampled).toBe(true);
+      });
     });
 
     describe('default sample context', () => {
@@ -133,9 +142,25 @@ describe('Hub', () => {
 
         expect(tracesSampler).toHaveBeenCalledWith(expect.objectContaining({ location: dogParkLocation }));
       });
+
+      it("should add parent's sampling decision to default sample context", () => {
+        const tracesSampler = jest.fn();
+        const hub = new Hub(new BrowserClient({ tracesSampler }));
+        const parentSamplingDecsion = false;
+
+        hub.startTransaction({
+          name: 'dogpark',
+          parentSpanId: '12312012',
+          sampled: parentSamplingDecsion,
+        });
+
+        expect(tracesSampler).toHaveBeenLastCalledWith(
+          expect.objectContaining({ parentSampled: parentSamplingDecsion }),
+        );
+      });
     });
 
-    describe('while sampling', () => {
+    describe('sample()', () => {
       it('should not sample transactions when tracing is disabled', () => {
         // neither tracesSampleRate nor tracesSampler is defined -> tracing disabled
         const hub = new Hub(new BrowserClient({}));
@@ -157,36 +182,41 @@ describe('Hub', () => {
 
         expect(transaction.sampled).toBe(true);
       });
+    });
 
-      it("should reject tracesSampleRates which aren't numbers", () => {
+    describe('isValidSampleRate()', () => {
+      it("should reject tracesSampleRates which aren't numbers or booleans", () => {
         const hub = new Hub(new BrowserClient({ tracesSampleRate: 'dogs!' as any }));
         hub.startTransaction({ name: 'dogpark' });
 
-        expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be a number'));
+        expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be a boolean or a number'));
       });
 
+      // the rate might be a boolean, but for our purposes, false is equivalent to 0 and true is equivalent to 1
       it('should reject tracesSampleRates less than 0', () => {
         const hub = new Hub(new BrowserClient({ tracesSampleRate: -26 }));
         hub.startTransaction({ name: 'dogpark' });
 
         expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be between 0 and 1'));
       });
 
+      // the rate might be a boolean, but for our purposes, false is equivalent to 0 and true is equivalent to 1
       it('should reject tracesSampleRates greater than 1', () => {
         const hub = new Hub(new BrowserClient({ tracesSampleRate: 26 }));
         hub.startTransaction({ name: 'dogpark' });
 
         expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be between 0 and 1'));
       });
 
-      it("should reject tracesSampler return values which aren't numbers", () => {
+      it("should reject tracesSampler return values which aren't numbers or booleans", () => {
         const tracesSampler = jest.fn().mockReturnValue('dogs!');
         const hub = new Hub(new BrowserClient({ tracesSampler }));
         hub.startTransaction({ name: 'dogpark' });
 
-        expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be a number'));
+        expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be a boolean or a number'));
       });
 
+      // the rate might be a boolean, but for our purposes, false is equivalent to 0 and true is equivalent to 1
       it('should reject tracesSampler return values less than 0', () => {
         const tracesSampler = jest.fn().mockReturnValue(-12);
         const hub = new Hub(new BrowserClient({ tracesSampler }));
@@ -195,6 +225,7 @@ describe('Hub', () => {
         expect(logger.warn).toHaveBeenCalledWith(expect.stringContaining('Sample rate must be between 0 and 1'));
       });
 
+      // the rate might be a boolean, but for our purposes, false is equivalent to 0 and true is equivalent to 1
       it('should reject tracesSampler return values greater than 1', () => {
         const tracesSampler = jest.fn().mockReturnValue(31);
         const hub = new Hub(new BrowserClient({ tracesSampler }));
@@ -204,14 +235,6 @@ describe('Hub', () => {
       });
     });
 
-    it('should propagate sampling decision to child spans', () => {
-      const hub = new Hub(new BrowserClient({ tracesSampleRate: 0 }));
-      const transaction = hub.startTransaction({ name: 'dogpark' });
-      const child = transaction.startChild({ op: 'test' });
-
-      expect(child.sampled).toBe(false);
-    });
-
     it('should drop transactions with sampled = false', () => {
       const client = new BrowserClient({ tracesSampleRate: 0 });
       jest.spyOn(client, 'captureEvent');
@@ -226,5 +249,55 @@ describe('Hub', () => {
       expect(transaction.finish).toReturnWith(undefined);
       expect(client.captureEvent).not.toBeCalled();
     });
+
+    describe('sampling inheritance', () => {
+      it('should propagate sampling decision to child spans', () => {
+        const hub = new Hub(new BrowserClient({ tracesSampleRate: Math.random() }));
+        const transaction = hub.startTransaction({ name: 'dogpark' });
+        const child = transaction.startChild({ op: 'test' });
+
+        expect(child.sampled).toBe(transaction.sampled);
+      });
+
+      it('should propagate sampling decision to child transactions in XHR header', () => {
+        // TODO this doesn't currently happen, but it should
+      });
+
+      it('should propagate sampling decision to child transactions in fetch header', () => {
+        // TODO this doesn't currently happen, but it should
+      });
+
+      it("should inherit parent's sampling decision when creating a new transaction if tracesSampler is undefined", () => {
+        // tracesSampleRate = 1 means every transaction should end up with sampled = true, so make parent's decision the
+        // opposite to prove that inheritance takes precedence over tracesSampleRate
+        const hub = new Hub(new BrowserClient({ tracesSampleRate: 1 }));
+        const parentSamplingDecsion = false;
+
+        const transaction = hub.startTransaction({
+          name: 'dogpark',
+          parentSpanId: '12312012',
+          sampled: parentSamplingDecsion,
+        });
+
+        expect(transaction.sampled).toBe(parentSamplingDecsion);
+      });
+
+      it("should ignore parent's sampling decision when tracesSampler is defined", () => {
+        // this tracesSampler causes every transaction to end up with sampled = true, so make parent's decision the
+        // opposite to prove that tracesSampler takes precedence over inheritance
+        const tracesSampler = () => true;
+        const parentSamplingDecsion = false;
+
+        const hub = new Hub(new BrowserClient({ tracesSampler }));
+
+        const transaction = hub.startTransaction({
+          name: 'dogpark',
+          parentSpanId: '12312012',
+          sampled: parentSamplingDecsion,
+        });
+
+        expect(transaction.sampled).not.toBe(parentSamplingDecsion);
+      });
+    });
   });
 });
diff --git a/packages/types/src/options.ts b/packages/types/src/options.ts
@@ -132,9 +132,10 @@ export interface Options {
    * Will automatically be passed a context object of default and optional custom data. See
    * {@link Transaction.sampleContext} and {@link Hub.startTransaction}.
    *
-   * @returns A sample rate between 0 and 1 (0 drops the trace, 1 guarantees it will be sent).
+   * @returns A sample rate between 0 and 1 (0 drops the trace, 1 guarantees it will be sent). Returning `true` is
+   * equivalent to returning 1 and returning `false` is equivalent to returning 0.
    */
-  tracesSampler?(sampleContext: SampleContext): number;
+  tracesSampler?(sampleContext: SampleContext): number | boolean;
 
   /**
    * A callback invoked during event submission, allowing to optionally modify
