Line data    Source code

       1             : #include "trace_sampler.h"
       2             : 
       3             : #include <algorithm>
       4             : #include <cassert>
       5             : #include <cstdint>
       6             : #include <limits>
       7             : 
       8             : #include "collector_response.h"
       9             : #include "json.hpp"
      10             : #include "sampling_decision.h"
      11             : #include "sampling_priority.h"
      12             : #include "sampling_util.h"
      13             : #include "span_data.h"
      14             : 
      15             : namespace datadog {
      16             : namespace tracing {
      17             : 
      18         623 : TraceSampler::TraceSampler(const FinalizedTraceSamplerConfig& config,
      19         623 :                            const Clock& clock)
      20         623 :     : rules_(config.rules),
      21         623 :       limiter_(clock, config.max_per_second),
      22        1246 :       limiter_max_per_second_(config.max_per_second) {}
      23             : 
      24       82821 : SamplingDecision TraceSampler::decide(const SpanData& span) {
      25       82821 :   SamplingDecision decision;
      26       82821 :   decision.origin = SamplingDecision::Origin::LOCAL;
      27             : 
      28             :   // First check sampling rules.
      29             :   auto found_rule =
      30       82821 :       std::find_if(rules_.begin(), rules_.end(),
      31       52329 :                    [&](const auto& rule) { return rule.match(span); });
      32             : 
      33             :   // `mutex_` protects `limiter_`, `collector_sample_rates_`, and
      34             :   // `collector_default_sample_rate_`, so let's lock it here.
      35       82821 :   std::lock_guard lock(mutex_);
      36             : 
      37       82821 :   if (found_rule != rules_.end()) {
      38       52327 :     const auto& rule = *found_rule;
      39       52327 :     decision.mechanism = int(SamplingMechanism::RULE);
      40       52327 :     decision.limiter_max_per_second = limiter_max_per_second_;
      41       52327 :     decision.configured_rate = rule.sample_rate;
      42       52327 :     const std::uint64_t threshold = max_id_from_rate(rule.sample_rate);
      43       52327 :     if (knuth_hash(span.trace_id.low) < threshold) {
      44       25316 :       const auto result = limiter_.allow();
      45       25316 :       if (result.allowed) {
      46       25127 :         decision.priority = int(SamplingPriority::USER_KEEP);
      47             :       } else {
      48         189 :         decision.priority = int(SamplingPriority::USER_DROP);
      49             :       }
      50       25316 :       decision.limiter_effective_rate = result.effective_rate;
      51             :     } else {
      52       27011 :       decision.priority = int(SamplingPriority::USER_DROP);
      53             :     }
      54             : 
      55       52327 :     return decision;
      56             :   }
      57             : 
      58             :   // No sampling rule matched.  Find the appropriate collector-controlled
      59             :   // sample rate.
      60       30494 :   auto found_rate = collector_sample_rates_.find(
      61       60988 :       CollectorResponse::key(span.service, span.environment().value_or("")));
      62       30494 :   if (found_rate != collector_sample_rates_.end()) {
      63        9999 :     decision.configured_rate = found_rate->second;
      64        9999 :     decision.mechanism = int(SamplingMechanism::AGENT_RATE);
      65             :   } else {
      66       20495 :     if (collector_default_sample_rate_) {
      67       10000 :       decision.configured_rate = *collector_default_sample_rate_;
      68       10000 :       decision.mechanism = int(SamplingMechanism::AGENT_RATE);
      69             :     } else {
      70             :       // We have yet to receive a default rate from the collector.  This
      71             :       // corresponds to the `DEFAULT` sampling mechanism.
      72       10495 :       decision.configured_rate = Rate::one();
      73       10495 :       decision.mechanism = int(SamplingMechanism::DEFAULT);
      74             :     }
      75             :   }
      76             : 
      77       30494 :   const std::uint64_t threshold = max_id_from_rate(*decision.configured_rate);
      78       30494 :   if (knuth_hash(span.trace_id.low) < threshold) {
      79       20502 :     decision.priority = int(SamplingPriority::AUTO_KEEP);
      80             :   } else {
      81        9992 :     decision.priority = int(SamplingPriority::AUTO_DROP);
      82             :   }
      83             : 
      84       30494 :   return decision;
      85       82821 : }
      86             : 
      87       30006 : void TraceSampler::handle_collector_response(
      88             :     const CollectorResponse& response) {
      89             :   const auto found =
      90       30006 :       response.sample_rate_by_key.find(response.key_of_default_rate);
      91       30006 :   std::lock_guard<std::mutex> lock(mutex_);
      92             : 
      93       30006 :   if (found != response.sample_rate_by_key.end()) {
      94       10004 :     collector_default_sample_rate_ = found->second;
      95             :   }
      96             : 
      97       30006 :   collector_sample_rates_ = response.sample_rate_by_key;
      98       30006 : }
      99             : 
     100         503 : nlohmann::json TraceSampler::config_json() const {
     101         503 :   std::vector<nlohmann::json> rules;
     102         520 :   for (const auto& rule : rules_) {
     103          17 :     rules.push_back(to_json(rule));
     104             :   }
     105             : 
     106        3521 :   return nlohmann::json::object({
     107             :       {"rules", rules},
     108         503 :       {"max_per_second", limiter_max_per_second_},
     109        4527 :   });
     110         503 : }
     111             : 
     112             : }  // namespace tracing
     113             : }  // namespace datadog