internal: refactor DatadogSampler, subclass from RateByServiceSampler

There are implementation details of these two samplers which need to be shared
having DatadogSampler subclass from RateByServiceSampler simplifies coordination

Author: brettlangdon

ddtrace/sampler.py

@@ -138,15 +138,24 @@ def set_sample_rate(
         # type: (...) -> None
         self._by_service_samplers[self._key(service, env)] = RateSampler(sample_rate)
 
+    def _set_priority(self, span, priority):
+        # type: (Span, int) -> None
+        span.context.sampling_priority = priority
+        span.sampled = priority > 0  # Positive priorities mean it was kept
+
     def sample(self, span):
         # type: (Span) -> bool
         tags = span.tracer.tags if span.tracer else {}
         env = tags[ENV_KEY] if ENV_KEY in tags else None
         key = self._key(span.service, env)
 
         sampler = self._by_service_samplers.get(key, self._by_service_samplers[self._default_key])
+        sampled = sampler.sample(span)
+        priority = AUTO_KEEP if sampled else AUTO_REJECT
+
         span.set_metric(SAMPLING_AGENT_DECISION, sampler.sample_rate)
-        return sampler.sample(span)
+        self._set_priority(span, priority)
+        return sampled
 
     def update_rate_by_service_sample_rates(self, rate_by_service):
         # type: (Dict[str, float]) -> None
@@ -157,7 +166,7 @@ def update_rate_by_service_sample_rates(self, rate_by_service):
         self._by_service_samplers = new_by_service_samplers
 
 
-class DatadogSampler(BasePrioritySampler):
+class DatadogSampler(RateByServiceSampler):
     """
     Default sampler used by Tracer for determining if a trace should be kept or dropped.
 
@@ -188,7 +197,7 @@ class DatadogSampler(BasePrioritySampler):
     provided. It is not used when the agent supplied sample rates are used.
     """
 
-    __slots__ = ("_agent_sampler", "limiter", "rules")
+    __slots__ = ("limiter", "rules")
 
     NO_RATE_LIMIT = -1
     DEFAULT_RATE_LIMIT = 100
@@ -212,6 +221,9 @@ def __init__(
             applied to them, (default: ``100``)
         :type rate_limit: :obj:`int`
         """
+        # Use default sample rate of 1.0
+        super(DatadogSampler, self).__init__()
+
         if default_sample_rate is None:
             sample_rate = get_env("trace", "sample_rate")
             if sample_rate is not None:
@@ -238,9 +250,6 @@ def __init__(
         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)
 
@@ -250,11 +259,12 @@ def __init__(
     def default_sampler(self):
         if self.rules:
             return self.rules[-1]
-        return self._agent_sampler
+        return self
 
     def __str__(self):
-        return "{}(agent_sampler={!r}, limiter={!r}, rules={!r})".format(
-            self.__class__.__name__, self._agent_sampler, self.limiter, self.rules
+        rates = {key: sampler.sample_rate for key, sampler in self._by_service_samplers.items()}
+        return "{}(agent_rates={!r}, limiter={!r}, rules={!r})".format(
+            self.__class__.__name__, rates, self.limiter, self.rules
         )
 
     __repr__ = __str__
@@ -280,16 +290,6 @@ def _parse_rules_from_env_variable(self, rules):
                 sampling_rules.append(sampling_rule)
         return sampling_rules
 
-    def update_rate_by_service_sample_rates(self, sample_rates):
-        # type: (Dict[str, float]) -> None
-        # Pass through the call to our RateByServiceSampler
-        self._agent_sampler.update_rate_by_service_sample_rates(sample_rates)
-
-    def _set_priority(self, span, priority):
-        # type: (Span, int) -> None
-        span.context.sampling_priority = priority
-        span.sampled = priority > 0  # Positive priorities mean it was kept
-
     def sample(self, span):
         # type: (Span) -> bool
         """
@@ -313,12 +313,7 @@ def sample(self, span):
                 break
         else:
             # 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
+            return super(DatadogSampler, self).sample(span)
 
         # DEV: This should never happen, but since the type is Optional we have to check
         if not matching_rule:

tests/tracer/test_sampler.py

@@ -622,39 +622,33 @@ 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._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._agent_sampler, RateByServiceSampler)
 
     # With rate limit
     sampler = DatadogSampler(rate_limit=10)
     assert sampler.limiter.rate_limit == 10
-    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 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 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 sampler.rules == [SamplingRule(sample_rate=0)]
-        assert isinstance(sampler._agent_sampler, RateByServiceSampler)
 
     # Invalid env vars
     with override_env(dict(DD_TRACE_SAMPLE_RATE="asdf")):
@@ -683,7 +677,7 @@ def test_datadog_sampler_init():
     assert sampler.rules == [rule_1, rule_2, rule_3, SamplingRule(sample_rate=0.75)]
 
 
-@mock.patch("ddtrace.sampler.RateByServiceSampler.sample")
+@mock.patch("ddtrace.sampler.RateSampler.sample")
 def test_datadog_sampler_sample_no_rules(mock_sample, dummy_tracer):
     sampler = DatadogSampler()
     dummy_tracer.configure(sampler=sampler)
@@ -939,7 +933,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._agent_sampler._by_service_samplers):
+        for k, v in iteritems(sampler._by_service_samplers):
             rates[k] = v.sample_rate
         assert case == rates, "%s != %s" % (case, rates)
 
@@ -949,6 +943,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._agent_sampler._by_service_samplers):
+        for k, v in iteritems(sampler._by_service_samplers):
             rates[k] = v.sample_rate
         assert case == rates, "%s != %s" % (case, rates)