Merge pull request #1172 from DataDog/brettlangdon/fallback.sampling

core: Update default sampler to new DatadogSampler

Author: brettlangdon

ddtrace/internal/writer.py

@@ -6,6 +6,7 @@
 from .. import api
 from .. import _worker
 from ..internal.logger import get_logger
+from ..sampler import BasePrioritySampler
 from ..settings import config
 from ..vendor import monotonic
 from ddtrace.vendor.six.moves.queue import Queue, Full, Empty
@@ -25,13 +26,14 @@ class AgentWriter(_worker.PeriodicWorkerThread):
 
     def __init__(self, hostname='localhost', port=8126, uds_path=None, https=False,
                  shutdown_timeout=DEFAULT_TIMEOUT,
-                 filters=None, priority_sampler=None,
+                 filters=None, sampler=None, priority_sampler=None,
                  dogstatsd=None):
         super(AgentWriter, self).__init__(interval=self.QUEUE_PROCESSING_INTERVAL,
                                           exit_timeout=shutdown_timeout,
                                           name=self.__class__.__name__)
         self._trace_queue = Q(maxsize=MAX_TRACES)
         self._filters = filters
+        self._sampler = sampler
         self._priority_sampler = priority_sampler
         self._last_error_ts = 0
         self.dogstatsd = dogstatsd
@@ -94,10 +96,17 @@ def flush_queue(self):
         for response in traces_responses:
             if isinstance(response, Exception) or response.status >= 400:
                 self._log_error_status(response)
-            elif self._priority_sampler:
+            elif self._priority_sampler or isinstance(self._sampler, BasePrioritySampler):
                 result_traces_json = response.get_json()
                 if result_traces_json and 'rate_by_service' in result_traces_json:
-                    self._priority_sampler.set_sample_rate_by_service(result_traces_json['rate_by_service'])
+                    if self._priority_sampler:
+                        self._priority_sampler.update_rate_by_service_sample_rates(
+                            result_traces_json['rate_by_service'],
+                        )
+                    if isinstance(self._sampler, BasePrioritySampler):
+                        self._sampler.update_rate_by_service_sample_rates(
+                            result_traces_json['rate_by_service'],
+                        )
 
         # Dump statistics
         # NOTE: Do not use the buffering of dogstatsd as it's not thread-safe

ddtrace/sampler.py

@@ -27,6 +27,12 @@ def sample(self, span):
         pass
 
 
+class BasePrioritySampler(six.with_metaclass(abc.ABCMeta)):
+    @abc.abstractmethod
+    def update_rate_by_service_sample_rates(self, sample_rates):
+        pass
+
+
 class AllSampler(BaseSampler):
     """Sampler sampling all the traces"""
 
@@ -60,7 +66,7 @@ def sample(self, span):
         return ((span.trace_id * KNUTH_FACTOR) % MAX_TRACE_ID) <= self.sampling_id_threshold
 
 
-class RateByServiceSampler(BaseSampler):
+class RateByServiceSampler(BaseSampler, BasePrioritySampler):
     """Sampler based on a rate, by service
 
     Keep (100 * `sample_rate`)% of the traces.
@@ -97,7 +103,7 @@ def sample(self, span):
         span.set_metric(SAMPLING_AGENT_DECISION, sampler.sample_rate)
         return sampler.sample(span)
 
-    def set_sample_rate_by_service(self, rate_by_service):
+    def update_rate_by_service_sample_rates(self, rate_by_service):
         new_by_service_samplers = self._get_new_by_service_sampler()
         for key, sample_rate in iteritems(rate_by_service):
             new_by_service_samplers[key] = RateSampler(sample_rate)
@@ -109,30 +115,41 @@ def set_sample_rate_by_service(self, rate_by_service):
 RateByServiceSampler._default_key = RateByServiceSampler._key()
 
 
-class DatadogSampler(BaseSampler):
+class DatadogSampler(BaseSampler, BasePrioritySampler):
     """
     This sampler is currently in ALPHA and it's API may change at any time, use at your own risk.
     """
     __slots__ = ('default_sampler', 'limiter', 'rules')
 
     NO_RATE_LIMIT = -1
     DEFAULT_RATE_LIMIT = 100
-    DEFAULT_SAMPLE_RATE = 1.0
+    DEFAULT_SAMPLE_RATE = None
 
     def __init__(self, rules=None, default_sample_rate=None, rate_limit=None):
         """
         Constructor for DatadogSampler sampler
 
         :param rules: List of :class:`SamplingRule` rules to apply to the root span of every trace, default no rules
         :type rules: :obj:`list` of :class:`SamplingRule`
-        :param default_sample_rate: The default sample rate to apply if no rules matched (default: 1.0)
+        :param default_sample_rate: The default sample rate to apply if no rules matched (default: ``None`` /
+            Use :class:`RateByServiceSampler` only)
         :type default_sample_rate: float 0 <= X <= 1.0
         :param rate_limit: Global rate limit (traces per second) to apply to all traces regardless of the rules
             applied to them, (default: ``100``)
         :type rate_limit: :obj:`int`
         """
         if default_sample_rate is None:
-            default_sample_rate = float(get_env('trace', 'sample_rate', default=self.DEFAULT_SAMPLE_RATE))
+            # If no sample rate was provided explicitly in code, try to load from environment variable
+            sample_rate = get_env('trace', 'sample_rate', default=self.DEFAULT_SAMPLE_RATE)
+
+            # If no env variable was found, just use the default
+            if sample_rate is None:
+                default_sample_rate = self.DEFAULT_SAMPLE_RATE
+
+            # Otherwise, try to convert it to a float
+            else:
+                default_sample_rate = float(sample_rate)
+
         if rate_limit is None:
             rate_limit = int(get_env('trace', 'rate_limit', default=self.DEFAULT_RATE_LIMIT))
 
@@ -148,7 +165,16 @@ def __init__(self, rules=None, default_sample_rate=None, rate_limit=None):
 
         # Configure rate limiter
         self.limiter = RateLimiter(rate_limit)
-        self.default_sampler = SamplingRule(sample_rate=default_sample_rate)
+
+        # Default to previous default behavior of RateByServiceSampler
+        self.default_sampler = RateByServiceSampler()
+        if default_sample_rate is not None:
+            self.default_sampler = SamplingRule(sample_rate=default_sample_rate)
+
+    def update_rate_by_service_sample_rates(self, sample_rates):
+        # Pass through the call to our RateByServiceSampler
+        if isinstance(self.default_sampler, RateByServiceSampler):
+            self.default_sampler.update_rate_by_service_sample_rates(sample_rates)
 
     def _set_priority(self, span, priority):
         if span._context:
@@ -175,7 +201,16 @@ def sample(self, span):
                 matching_rule = rule
                 break
         else:
-            # No rule matches, use the default sampler
+            # 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 defualt sampler
             matching_rule = self.default_sampler
 
         # Sample with the matching sampling rule
@@ -202,7 +237,7 @@ def sample(self, span):
         return True
 
 
-class SamplingRule(object):
+class SamplingRule(BaseSampler):
     """
     Definition of a sampling rule used by :class:`DatadogSampler` for applying a sample rate on a span
     """

ddtrace/tracer.py

@@ -12,7 +12,7 @@
 from .internal.writer import AgentWriter
 from .provider import DefaultContextProvider
 from .context import Context
-from .sampler import AllSampler, DatadogSampler, RateSampler, RateByServiceSampler
+from .sampler import DatadogSampler, RateSampler, RateByServiceSampler
 from .span import Span
 from .utils.formats import get_env
 from .utils.deprecation import deprecated, RemovedInDDTrace10Warning
@@ -114,7 +114,7 @@ def __init__(self, url=DEFAULT_AGENT_URL, dogstatsd_url=DEFAULT_DOGSTATSD_URL):
             port=port,
             https=https,
             uds_path=uds_path,
-            sampler=AllSampler(),
+            sampler=DatadogSampler(),
             context_provider=DefaultContextProvider(),
             dogstatsd_url=dogstatsd_url,
         )
@@ -229,7 +229,7 @@ def configure(self, enabled=None, hostname=None, port=None, uds_path=None, https
             self._dogstatsd_client = DogStatsd(**dogstatsd_kwargs)
 
         if hostname is not None or port is not None or uds_path is not None or https is not None or \
-                filters is not None or priority_sampling is not None:
+                filters is not None or priority_sampling is not None or sampler is not None:
             # Preserve hostname and port when overriding filters or priority sampling
             # This is clumsy and a good reason to get rid of this configure() API
             if hasattr(self, 'writer') and hasattr(self.writer, 'api'):
@@ -247,6 +247,7 @@ def configure(self, enabled=None, hostname=None, port=None, uds_path=None, https
                 uds_path=uds_path,
                 https=https,
                 filters=filters,
+                sampler=self.sampler,
                 priority_sampler=self.priority_sampler,
                 dogstatsd=self._dogstatsd_client,
             )
@@ -370,6 +371,8 @@ def start_span(self, name, child_of=None, service=None, resource=None, span_type
                         context.sampling_priority = AUTO_REJECT
             else:
                 context.sampling_priority = AUTO_KEEP if span.sampled else AUTO_REJECT
+                # We must always mark the span as sampled so it is forwarded to the agent
+                span.sampled = True
 
             # add tags to root span to correlate trace with runtime metrics
             # only applied to spans with types that are internal to applications