Add traces_sampler option (getsentry#863)

lobsterkatie · web-flow · commit 52830558bb53 · 2020-10-22T11:17:59.000-07:00
diff --git a/sentry_sdk/consts.py b/sentry_sdk/consts.py
@@ -69,7 +69,7 @@ def __init__(
         attach_stacktrace=False,  # type: bool
         ca_certs=None,  # type: Optional[str]
         propagate_traces=True,  # type: bool
-        traces_sample_rate=0.0,  # type: float
+        traces_sample_rate=None,  # type: Optional[float]
         traces_sampler=None,  # type: Optional[TracesSampler]
         auto_enabling_integrations=True,  # type: bool
         _experiments={},  # type: Experiments  # noqa: B006
diff --git a/sentry_sdk/hub.py b/sentry_sdk/hub.py
@@ -1,5 +1,4 @@
 import copy
-import random
 import sys
 
 from datetime import datetime
@@ -505,20 +504,28 @@ def start_transaction(
         When the transaction is finished, it will be sent to Sentry with all its
         finished child spans.
         """
+        custom_sampling_context = kwargs.pop("custom_sampling_context", {})
+
+        # if we haven't been given a transaction, make one
         if transaction is None:
             kwargs.setdefault("hub", self)
             transaction = Transaction(**kwargs)
 
-        client, scope = self._stack[-1]
-
-        if transaction.sampled is None:
-            sample_rate = client and client.options["traces_sample_rate"] or 0
-            transaction.sampled = random.random() < sample_rate
-
+        # use traces_sample_rate, traces_sampler, and/or inheritance to make a
+        # sampling decision
+        sampling_context = {
+            "transaction_context": transaction.to_json(),
+            "parent_sampled": transaction.parent_sampled,
+        }
+        sampling_context.update(custom_sampling_context)
+        transaction._set_initial_sampling_decision(sampling_context=sampling_context)
+
+        # we don't bother to keep spans if we already know we're not going to
+        # send the transaction
         if transaction.sampled:
             max_spans = (
-                client and client.options["_experiments"].get("max_spans") or 1000
-            )
+                self.client and self.client.options["_experiments"].get("max_spans")
+            ) or 1000
             transaction.init_span_recorder(maxlen=max_spans)
 
         return transaction
diff --git a/sentry_sdk/tracing.py b/sentry_sdk/tracing.py
@@ -2,14 +2,19 @@
 import uuid
 import contextlib
 import math
+import random
 import time
 
 from datetime import datetime, timedelta
 from numbers import Real
 
 import sentry_sdk
 
-from sentry_sdk.utils import capture_internal_exceptions, logger, to_string
+from sentry_sdk.utils import (
+    capture_internal_exceptions,
+    logger,
+    to_string,
+)
 from sentry_sdk._compat import PY2
 from sentry_sdk._types import MYPY
 
@@ -28,6 +33,8 @@
     from typing import List
     from typing import Tuple
 
+    from sentry_sdk._types import SamplingContext
+
 _traceparent_header_format_re = re.compile(
     "^[ \t]*"  # whitespace
     "([0-9a-f]{32})?"  # trace_id
@@ -337,7 +344,7 @@ def from_traceparent(
         return Transaction(
             trace_id=trace_id,
             parent_span_id=parent_span_id,
-            sampled=parent_sampled,
+            parent_sampled=parent_sampled,
             **kwargs
         )
 
@@ -555,6 +562,116 @@ def to_json(self):
 
         return rv
 
+    def _set_initial_sampling_decision(self, sampling_context):
+        # type: (SamplingContext) -> None
+        """
+        Sets the transaction's sampling decision, according to the following
+        precedence rules:
+
+        1. If a sampling decision is passed to `start_transaction`
+        (`start_transaction(name: "my transaction", sampled: True)`), that
+        decision will be used, regardlesss of anything else
+
+        2. If `traces_sampler` is defined, its decision will be used. It can
+        choose to keep or ignore any parent sampling decision, or use the
+        sampling context data to make its own decision or to choose a sample
+        rate for the transaction.
+
+        3. If `traces_sampler` is not defined, but there's a parent sampling
+        decision, the parent sampling decision will be used.
+
+        4. If `traces_sampler` is not defined and there's no parent sampling
+        decision, `traces_sample_rate` will be used.
+        """
+
+        hub = self.hub or sentry_sdk.Hub.current
+        client = hub.client
+        options = (client and client.options) or {}
+        transaction_description = "{op}transaction <{name}>".format(
+            op=("<" + self.op + "> " if self.op else ""), name=self.name
+        )
+
+        # nothing to do if there's no client or if tracing is disabled
+        if not client or not has_tracing_enabled(options):
+            self.sampled = False
+            return
+
+        # if the user has forced a sampling decision by passing a `sampled`
+        # value when starting the transaction, go with that
+        if self.sampled is not None:
+            return
+
+        # we would have bailed already if neither `traces_sampler` nor
+        # `traces_sample_rate` were defined, so one of these should work; prefer
+        # the hook if so
+        sample_rate = (
+            options["traces_sampler"](sampling_context)
+            if callable(options.get("traces_sampler"))
+            else (
+                # default inheritance behavior
+                sampling_context["parent_sampled"]
+                if sampling_context["parent_sampled"] is not None
+                else options["traces_sample_rate"]
+            )
+        )
+
+        # 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 not _is_valid_sample_rate(sample_rate):
+            logger.warning(
+                "[Tracing] Discarding {transaction_description} because of invalid sample rate.".format(
+                    transaction_description=transaction_description,
+                )
+            )
+            self.sampled = False
+            return
+
+        # if the function returned 0 (or false), or if `traces_sample_rate` is
+        # 0, it's a sign the transaction should be dropped
+        if not sample_rate:
+            logger.debug(
+                "[Tracing] Discarding {transaction_description} because {reason}".format(
+                    transaction_description=transaction_description,
+                    reason=(
+                        "traces_sampler returned 0 or False"
+                        if callable(options.get("traces_sampler"))
+                        else "traces_sample_rate is set to 0"
+                    ),
+                )
+            )
+            self.sampled = False
+            return
+
+        # Now we roll the dice. random.random is inclusive of 0, but not of 1,
+        # so strict < is safe here. In case sample_rate is a boolean, cast it
+        # to a float (True becomes 1.0 and False becomes 0.0)
+        self.sampled = random.random() < float(sample_rate)
+
+        if self.sampled:
+            logger.debug(
+                "[Tracing] Starting {transaction_description}".format(
+                    transaction_description=transaction_description,
+                )
+            )
+        else:
+            logger.debug(
+                "[Tracing] Discarding {transaction_description} because it's not included in the random sample (sampling rate = {sample_rate})".format(
+                    transaction_description=transaction_description,
+                    sample_rate=float(sample_rate),
+                )
+            )
+
+
+def has_tracing_enabled(options):
+    # type: (Dict[str, Any]) -> bool
+    """
+    Returns True if either traces_sample_rate or traces_sampler is
+    non-zero/defined, False otherwise.
+    """
+
+    return bool(options.get("traces_sample_rate") or options.get("traces_sampler"))
+
 
 def _is_valid_sample_rate(rate):
     # type: (Any) -> bool
diff --git a/sentry_sdk/utils.py b/sentry_sdk/utils.py
@@ -968,13 +968,3 @@ def run(self):
                 integer_configured_timeout
             )
         )
-
-
-def has_tracing_enabled(options):
-    # type: (Dict[str, Any]) -> bool
-    """
-    Returns True if either traces_sample_rate or traces_sampler is
-    non-zero/defined, False otherwise.
-    """
-
-    return bool(options.get("traces_sample_rate") or options.get("traces_sampler"))
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -1,6 +1,5 @@
 import os
 import json
-from types import FunctionType
 
 import pytest
 import jsonschema
@@ -37,11 +36,6 @@ def benchmark():
 else:
     del pytest_benchmark
 
-try:
-    from unittest import mock  # python 3.3 and above
-except ImportError:
-    import mock  # python < 3.3
-
 
 @pytest.fixture(autouse=True)
 def internal_exceptions(request, monkeypatch):
@@ -400,18 +394,3 @@ def __eq__(self, test_dict):
             return all(test_dict.get(key) == self.subdict[key] for key in self.subdict)
 
     return DictionaryContaining
-
-
-@pytest.fixture(name="FunctionMock")
-def function_mock():
-    """
-    Just like a mock.Mock object, but one which always passes an isfunction
-    test.
-    """
-
-    class FunctionMock(mock.Mock):
-        def __init__(self, *args, **kwargs):
-            super(FunctionMock, self).__init__(*args, **kwargs)
-            self.__class__ = FunctionType
-
-    return FunctionMock
diff --git a/tests/integrations/sqlalchemy/test_sqlalchemy.py b/tests/integrations/sqlalchemy/test_sqlalchemy.py
@@ -76,7 +76,9 @@ class Address(Base):
 def test_transactions(sentry_init, capture_events, render_span_tree):
 
     sentry_init(
-        integrations=[SqlalchemyIntegration()], _experiments={"record_sql_params": True}
+        integrations=[SqlalchemyIntegration()],
+        _experiments={"record_sql_params": True},
+        traces_sample_rate=1.0,
     )
     events = capture_events()
 
diff --git a/tests/tracing/test_integration_tests.py b/tests/tracing/test_integration_tests.py
@@ -70,7 +70,7 @@ def test_continue_from_headers(sentry_init, capture_events, sampled):
     # correctly
     transaction = Transaction.continue_from_headers(headers, name="WRONG")
     assert transaction is not None
-    assert transaction.sampled == sampled
+    assert transaction.parent_sampled == sampled
     assert transaction.trace_id == old_span.trace_id
     assert transaction.same_process_as_parent is False
     assert transaction.parent_span_id == old_span.span_id
diff --git a/tests/tracing/test_sampling.py b/tests/tracing/test_sampling.py

Original file line number	Diff line number	Diff line change
`@@ -968,13 +968,3 @@ def run(self):`
`968`	`968`	`integer_configured_timeout`
`969`	`969`	`)`
`970`	`970`	`)`
`971`		`-`
`972`		`-`
`973`		`-def has_tracing_enabled(options):`
`974`		`- # type: (Dict[str, Any]) -> bool`
`975`		`- """`
`976`		`- Returns True if either traces_sample_rate or traces_sampler is`
`977`		`- non-zero/defined, False otherwise.`
`978`		`- """`
`979`		`-`
`980`		`- return bool(options.get("traces_sample_rate") or options.get("traces_sampler"))`
Original file line number	Diff line number	Diff line change
`@@ -76,7 +76,9 @@ class Address(Base):`
`76`	`76`	`def test_transactions(sentry_init, capture_events, render_span_tree):`
`77`	`77`
`78`	`78`	`sentry_init(`
`79`		`- integrations=[SqlalchemyIntegration()], _experiments={"record_sql_params": True}`
	`79`	`+ integrations=[SqlalchemyIntegration()],`
	`80`	`+ _experiments={"record_sql_params": True},`
	`81`	`+ traces_sample_rate=1.0,`
`80`	`82`	`)`
`81`	`83`	`events = capture_events()`
`82`	`84`