internal: refactor DatadogSampler default sampling rule and agent sampler

Right now DatadogSampler.default_sampler has the responsibility of either
being a default SamplingRule or the older RateByServiceSampler.

Instead of having a single attribute which tries to represent both of
these different paths, separating them into two makes juggling types
and reasoning about the flow of the sampler better.

Author: brettlangdon

ddtrace/sampler.py

@@ -27,6 +27,7 @@
 from .internal.rate_limiter import RateLimiter
 from .internal.utils.cache import cachedmethod
 from .internal.utils.formats import get_env
+from .vendor.debtcollector.removals import removed_property
 
 
 try:
@@ -47,6 +48,10 @@
 KNUTH_FACTOR = 1111111111111111111
 
 
+class SamplingError(Exception):
+    pass
+
+
 class BaseSampler(six.with_metaclass(abc.ABCMeta)):
     @abc.abstractmethod
     def sample(self, span):
@@ -153,7 +158,37 @@ def update_rate_by_service_sample_rates(self, rate_by_service):
 
 
 class DatadogSampler(BasePrioritySampler):
-    __slots__ = ("default_sampler", "limiter", "rules")
+    """
+    Default sampler used by Tracer for determining if a trace should be kept or dropped.
+
+    By default this sampler will rely on dynamic sample rates provided by the trace agent
+    to determine which traces are kept or dropped.
+
+    You can also configure a static sample rate via ``default_sample_rate`` to use for sampling.
+    When a ``default_sample_rate`` is configured, that is the only sample rate used, the agent
+    provided rates are ignored.
+
+    You may also supply a list of ``SamplingRule`` to determine sample rates for specific
+    services or operation names.
+
+    Example rules::
+
+        DatadogSampler(rules=[
+            SamplingRule(sample_rate=1.0, service="my-svc"),
+            SamplingRule(sample_rate=0.0, service="less-important"),
+        ])
+
+    Rules are evaluated in the order they are provided, and the first rule that matches is used.
+    If no rule matches, then the agent sample rates are used.
+
+
+    Lastly, this sampler can be configured with a rate limit. This will ensure the max number of
+    sampled traces per second does not exceed the supplied limit. The default is 100 traces kept
+    per second. This rate limiter is only used when ``default_sample_rate`` or ``rules`` are
+    provided. It is not used when the agent supplied sample rates are used.
+    """
+
+    __slots__ = ("_agent_sampler", "limiter", "rules")
 
     NO_RATE_LIMIT = -1
     DEFAULT_RATE_LIMIT = 100
@@ -186,6 +221,7 @@ def __init__(
             rate_limit = int(get_env("trace", "rate_limit", default=self.DEFAULT_RATE_LIMIT))  # type: ignore[arg-type]
 
         # Ensure rules is a list
+        self.rules = []  # type: List[SamplingRule]
         if rules is None:
             env_sampling_rules = get_env("trace", "sampling_rules")
             if env_sampling_rules:
@@ -198,21 +234,30 @@ def __init__(
             if not isinstance(rule, SamplingRule):
                 raise TypeError("Rule {!r} must be a sub-class of type ddtrace.sampler.SamplingRules".format(rule))
         self.rules = rules
+        # DEV: Default sampling rule must come last
+        if default_sample_rate is not None:
+            self.rules.append(SamplingRule(sample_rate=default_sample_rate))
+
+        # If no rules are configured (or match), fallback to agent based sampling
+        self._agent_sampler = RateByServiceSampler()
 
         # Configure rate limiter
         self.limiter = RateLimiter(rate_limit)
 
-        if default_sample_rate is None:
-            log.debug("initialized DatadogSampler, limit %r traces per second", rate_limit)
-            # Default to previous default behavior of RateByServiceSampler
-            self.default_sampler = RateByServiceSampler()  # type: BaseSampler
-        else:
-            log.debug(
-                "initialized DatadogSampler, sample %s%% traces, limit %r traces per second",
-                100 * default_sample_rate,
-                rate_limit,
-            )
-            self.default_sampler = SamplingRule(sample_rate=default_sample_rate)
+        log.debug("initialized %r", self)
+
+    @removed_property(removal_version="1.0.0")
+    def default_sampler(self):
+        if self.rules:
+            return self.rules[-1]
+        return self._agent_sampler
+
+    def __str__(self):
+        return "{}(agent_sampler={!r}, limiter={!r}, rules={!r})".format(
+            self.__class__.__name__, self._agent_sampler, self.limiter, self.rules
+        )
+
+    __repr__ = __str__
 
     def _parse_rules_from_env_variable(self, rules):
         sampling_rules = []
@@ -238,8 +283,7 @@ def _parse_rules_from_env_variable(self, rules):
     def update_rate_by_service_sample_rates(self, sample_rates):
         # type: (Dict[str, float]) -> None
         # Pass through the call to our RateByServiceSampler
-        if isinstance(self.default_sampler, RateByServiceSampler):
-            self.default_sampler.update_rate_by_service_sample_rates(sample_rates)
+        self._agent_sampler.update_rate_by_service_sample_rates(sample_rates)
 
     def _set_priority(self, span, priority):
         # type: (Span, int) -> None
@@ -259,29 +303,29 @@ def sample(self, span):
         :rtype: :obj:`bool`
         """
         # If there are rules defined, then iterate through them and find one that wants to sample
-        matching_rule = None  # type: Optional[BaseSampler]
+        matching_rule = None  # type: Optional[SamplingRule]
+
         # Go through all rules and grab the first one that matched
         # DEV: This means rules should be ordered by the user from most specific to least specific
         for rule in self.rules:
             if rule.matches(span):
                 matching_rule = rule
                 break
         else:
-            # If this is the old sampler, sample and return
-            if isinstance(self.default_sampler, RateByServiceSampler):
-                if self.default_sampler.sample(span):
-                    self._set_priority(span, AUTO_KEEP)
-                    return True
-                else:
-                    self._set_priority(span, AUTO_REJECT)
-                    return False
-
-            # If no rules match, use our default sampler
-            matching_rule = self.default_sampler
+            # No rules matches so use agent based sampling
+            if self._agent_sampler.sample(span):
+                self._set_priority(span, AUTO_KEEP)
+                return True
+            else:
+                self._set_priority(span, AUTO_REJECT)
+                return False
+
+        # DEV: This should never happen, but since the type is Optional we have to check
+        if not matching_rule:
+            raise SamplingError("No sampling rule found for span {!r} from {!r}".format(span, self))
 
         # Sample with the matching sampling rule
-        if isinstance(matching_rule, (RateSampler, SamplingRule)):
-            span.set_metric(SAMPLING_RULE_DECISION, matching_rule.sample_rate)
+        span.set_metric(SAMPLING_RULE_DECISION, matching_rule.sample_rate)
         if not matching_rule.sample(span):
             self._set_priority(span, USER_REJECT)
             return False
@@ -451,3 +495,10 @@ def __repr__(self):
         )
 
     __str__ = __repr__
+
+    def __eq__(self, other):
+        # type: (Any) -> bool
+        if not isinstance(other, SamplingRule):
+            raise TypeError("Cannot compare SamplingRule to {}".format(type(other)))
+
+        return self.sample_rate == other.sample_rate and self.service == other.service and self.name == other.name

releasenotes/notes/deprecate-datadogsampler-default_sampler-7f4a8f9923cd3c7b.yaml

@@ -0,0 +1,4 @@
+---
+deprecations:
+  - |
+    ``ddtrace.sampler.DatadogSampler.default_sampler`` property is deprecated and will be removed in 1.0.

tests/tracer/test_sampler.py

@@ -263,6 +263,94 @@ def test_sampling_rule_init():
     assert rule.name == name_regex
 
 
+@pytest.mark.parametrize(
+    "rule_1,rule_2,expected",
+    [
+        # Sample rate only
+        (SamplingRule(sample_rate=1.0), SamplingRule(sample_rate=1.0), True),
+        (SamplingRule(sample_rate=0.5), SamplingRule(sample_rate=0.5), True),
+        (SamplingRule(sample_rate=0.0), SamplingRule(sample_rate=0.0), True),
+        (SamplingRule(sample_rate=0.5), SamplingRule(sample_rate=1.0), False),
+        # Sample rate, and service name
+        (SamplingRule(sample_rate=1.0, service="my-svc"), SamplingRule(sample_rate=1.0, service="my-svc"), True),
+        (
+            SamplingRule(sample_rate=1.0, service=re.compile("my-svc")),
+            SamplingRule(sample_rate=1.0, service=re.compile("my-svc")),
+            True,
+        ),
+        (SamplingRule(sample_rate=1.0, service="my-svc"), SamplingRule(sample_rate=1.0, service="other-svc"), False),
+        (SamplingRule(sample_rate=1.0, service="my-svc"), SamplingRule(sample_rate=0.5, service="my-svc"), False),
+        (
+            SamplingRule(sample_rate=1.0, service=re.compile("my-svc")),
+            SamplingRule(sample_rate=0.5, service=re.compile("my-svc")),
+            False,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, service=re.compile("my-svc")),
+            SamplingRule(sample_rate=1.0, service=re.compile("other")),
+            False,
+        ),
+        # Sample rate, and operation name
+        (
+            SamplingRule(sample_rate=1.0, name="span.name"),
+            SamplingRule(sample_rate=1.0, name="span.name"),
+            True,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, name=re.compile("span.name")),
+            SamplingRule(sample_rate=1.0, name=re.compile("span.name")),
+            True,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, name=re.compile("span.name")),
+            SamplingRule(sample_rate=1.0, name=re.compile("span.other")),
+            False,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, name="span.name"),
+            SamplingRule(sample_rate=0.5, name="span.name"),
+            False,
+        ),
+        (SamplingRule(sample_rate=1.0, name="span.name"), SamplingRule(sample_rate=1.0, name="span.other"), False),
+        (SamplingRule(sample_rate=1.0, name="span.name"), SamplingRule(sample_rate=0.5, name="span.name"), False),
+        # Sample rate, service, and operation name
+        (
+            SamplingRule(sample_rate=1.0, service="my-svc", name="span.name"),
+            SamplingRule(sample_rate=1.0, service="my-svc", name="span.name"),
+            True,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, service="my-svc", name=re.compile("span.name")),
+            SamplingRule(sample_rate=1.0, service="my-svc", name=re.compile("span.name")),
+            True,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, service=re.compile("my-svc"), name=re.compile("span.name")),
+            SamplingRule(sample_rate=1.0, service=re.compile("my-svc"), name=re.compile("span.name")),
+            True,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, service="my-svc", name="span.name"),
+            SamplingRule(sample_rate=0.5, service="my-svc", name="span.name"),
+            False,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, service="my-svc", name="span.name"),
+            SamplingRule(sample_rate=1.0, service="other", name="span.name"),
+            False,
+        ),
+        (
+            SamplingRule(sample_rate=1.0, service="my-svc", name="span.name"),
+            SamplingRule(sample_rate=1.0, service="my-svc", name="span.other"),
+            False,
+        ),
+    ],
+)
+def test_sampling_rule_eq(rule_1, rule_2, expected):
+    result = rule_1 == rule_2
+    assert result == expected
+
+
 def test_sampling_rule_init_via_env():
     # Testing single sampling rule
     with override_env(dict(DD_TRACE_SAMPLING_RULES='[{"sample_rate":1.0,"service":"xyz","name":"abc"}]')):
@@ -534,39 +622,39 @@ def test_datadog_sampler_init():
     assert sampler.rules == []
     assert isinstance(sampler.limiter, RateLimiter)
     assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
-    assert isinstance(sampler.default_sampler, RateByServiceSampler)
+    assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # With rules
     rule = SamplingRule(sample_rate=1)
     sampler = DatadogSampler(rules=[rule])
     assert sampler.rules == [rule]
     assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
-    assert isinstance(sampler.default_sampler, RateByServiceSampler)
+    assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # With rate limit
     sampler = DatadogSampler(rate_limit=10)
     assert sampler.limiter.rate_limit == 10
-    assert isinstance(sampler.default_sampler, RateByServiceSampler)
+    assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # With default_sample_rate
     sampler = DatadogSampler(default_sample_rate=0.5)
     assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
-    assert isinstance(sampler.default_sampler, SamplingRule)
-    assert sampler.default_sampler.sample_rate == 0.5
+    assert sampler.rules == [SamplingRule(sample_rate=0.5)]
+    assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # From env variables
     with override_env(dict(DD_TRACE_SAMPLE_RATE="0.5", DD_TRACE_RATE_LIMIT="10")):
         sampler = DatadogSampler()
         assert sampler.limiter.rate_limit == 10
-        assert isinstance(sampler.default_sampler, SamplingRule)
-        assert sampler.default_sampler.sample_rate == 0.5
+        assert sampler.rules == [SamplingRule(sample_rate=0.5)]
+        assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # DD_TRACE_SAMPLE_RATE=0
     with override_env(dict(DD_TRACE_SAMPLE_RATE="0")):
         sampler = DatadogSampler()
         assert sampler.limiter.rate_limit == DatadogSampler.DEFAULT_RATE_LIMIT
-        assert isinstance(sampler.default_sampler, SamplingRule)
-        assert sampler.default_sampler.sample_rate == 0
+        assert sampler.rules == [SamplingRule(sample_rate=0)]
+        assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # Invalid env vars
     with override_env(dict(DD_TRACE_SAMPLE_RATE="asdf")):
@@ -590,6 +678,10 @@ def test_datadog_sampler_init():
     sampler = DatadogSampler(rules=[rule_1, rule_2, rule_3])
     assert sampler.rules == [rule_1, rule_2, rule_3]
 
+    # Ensure default rule is appended
+    sampler = DatadogSampler(rules=[rule_1, rule_2, rule_3], default_sample_rate=0.75)
+    assert sampler.rules == [rule_1, rule_2, rule_3, SamplingRule(sample_rate=0.75)]
+
 
 @mock.patch("ddtrace.sampler.RateByServiceSampler.sample")
 def test_datadog_sampler_sample_no_rules(mock_sample, dummy_tracer):
@@ -838,7 +930,7 @@ def test_datadog_sampler_update_rate_by_service_sample_rates(dummy_tracer):
     for case in cases:
         sampler.update_rate_by_service_sample_rates(case)
         rates = {}
-        for k, v in iteritems(sampler.default_sampler._by_service_samplers):
+        for k, v in iteritems(sampler._agent_sampler._by_service_samplers):
             rates[k] = v.sample_rate
         assert case == rates, "%s != %s" % (case, rates)
 
@@ -848,6 +940,6 @@ def test_datadog_sampler_update_rate_by_service_sample_rates(dummy_tracer):
     for case in cases:
         sampler.update_rate_by_service_sample_rates(case)
         rates = {}
-        for k, v in iteritems(sampler.default_sampler._by_service_samplers):
+        for k, v in iteritems(sampler._agent_sampler._by_service_samplers):
             rates[k] = v.sample_rate
         assert case == rates, "%s != %s" % (case, rates)