datadog 2.0.0.beta1

data/lib/datadog/tracing/sampling/matcher.rb

@@ -0,0 +1,65 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Sampling
+      # Checks if a trace conforms to a matching criteria.
+      # @abstract
+      class Matcher
+        # Returns `true` if the trace should conforms to this rule, `false` otherwise
+        #
+        # @param [TraceOperation] trace
+        # @return [Boolean]
+        def match?(trace)
+          raise NotImplementedError
+        end
+      end
+
+      # A {Datadog::Sampling::Matcher} that supports matching a trace by
+      # trace name and/or service name.
+      class SimpleMatcher < Matcher
+        # Returns `true` for case equality (===) with any object
+        MATCH_ALL = Class.new do
+          # DEV: A class that implements `#===` is ~20% faster than
+          # DEV: a `Proc` that always returns `true`.
+          def ===(other)
+            true
+          end
+        end.new
+
+        attr_reader :name, :service
+
+        # @param name [String,Regexp,Proc] Matcher for case equality (===) with the trace name,
+        #             defaults to always match
+        # @param service [String,Regexp,Proc] Matcher for case equality (===) with the service name,
+        #                defaults to always match
+        def initialize(name: MATCH_ALL, service: MATCH_ALL)
+          super()
+          @name = name
+          @service = service
+        end
+
+        def match?(trace)
+          name === trace.name && service === trace.service
+        end
+      end
+
+      # A {Datadog::Tracing::Sampling::Matcher} that allows for arbitrary trace matching
+      # based on the return value of a provided block.
+      class ProcMatcher < Matcher
+        attr_reader :block
+
+        # @yield [name, service] Provides trace name and service to the block
+        # @yieldreturn [Boolean] Whether the trace conforms to this matcher
+        def initialize(&block)
+          super()
+          @block = block
+        end
+
+        def match?(trace)
+          block.call(trace.name, trace.service)
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/sampling/priority_sampler.rb

@@ -0,0 +1,160 @@
+# frozen_string_literal: true
+
+require_relative 'ext'
+require_relative 'all_sampler'
+require_relative 'rate_sampler'
+require_relative 'rate_by_service_sampler'
+
+module Datadog
+  module Tracing
+    module Sampling
+      # {Datadog::Tracing::Sampling::PrioritySampler}
+      class PrioritySampler
+        # NOTE: We do not advise using a pre-sampler. It can save resources,
+        # but pre-sampling at rates < 100% may result in partial traces, unless
+        # the pre-sampler knows exactly how to drop a span without dropping its ancestors.
+        #
+        # Additionally, as service metrics are calculated in the Datadog Agent,
+        # the service's throughput will be underestimated.
+        attr_reader :pre_sampler, :priority_sampler
+
+        def initialize(opts = {})
+          @pre_sampler = opts[:base_sampler] || AllSampler.new
+          @priority_sampler = opts[:post_sampler] || RateByServiceSampler.new(decision: Sampling::Ext::Decision::AGENT_RATE)
+        end
+
+        # DEV-2.0:We should get rid of this complicated interaction between @pre_sampler and @priority_sampler.
+        # DEV-2.0:If the user wants to configure a custom sampler, we should only allow them to provide a complete
+        # DEV-2.0:sampling suite, not having this convoluted support for mixing arbitrary provided samplers in
+        # DEV-2.0:the PrioritySampler. Ideally, the PrioritySampler is only used by Datadog.
+        # DEV-2.0:There are too many edge cases and combinations to work around currently in this class.
+        def sample!(trace)
+          # The priority that was set before the sampler ran.
+          # This comes from distributed tracing priority propagation.
+          distributed_sampling_priority = priority_assigned?(trace)
+
+          # If pre-sampling is configured, do it first. (By default, this will sample at 100%.)
+          # NOTE: Pre-sampling at rates < 100% may result in partial traces; not recommended.
+          trace.sampled = pre_sample?(trace) ? preserving_priority_sampling(trace) { @pre_sampler.sample!(trace) } : true
+
+          if trace.sampled?
+            # If priority sampling has already been applied upstream, use that value.
+            return true if priority_assigned?(trace)
+
+            # Check with post sampler how we set the priority.
+            sample = priority_sample!(trace)
+
+            # Check if post sampler has already assigned a priority.
+            return true if priority_assigned?(trace)
+
+            # If not, use agent priority values.
+            priority = sample ? Sampling::Ext::Priority::AUTO_KEEP : Sampling::Ext::Priority::AUTO_REJECT
+            assign_priority!(trace, priority)
+          else
+            # If discarded by pre-sampling, set "reject" priority, so other
+            # services for the same trace don't sample needlessly.
+            assign_priority!(trace, Sampling::Ext::Priority::AUTO_REJECT)
+          end
+
+          trace.sampled?
+        ensure
+          if trace.sampling_priority && trace.sampling_priority > 0
+            # Don't modify decision if priority was set upstream.
+            if !distributed_sampling_priority && !trace.has_tag?(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER)
+              # If no sampling priority being assigned at this point, a custom
+              # sampler implementation is configured: this means the user has
+              # full control over the sampling decision.
+              trace.set_tag(
+                Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER,
+                Sampling::Ext::Decision::MANUAL
+              )
+            end
+          else
+            # The sampler decided to not keep this span, removing sampling decision.
+            trace.clear_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER)
+          end
+        end
+
+        # (see Datadog::Tracing::Sampling::RateByServiceSampler#update)
+        def update(rate_by_service, decision: nil)
+          @priority_sampler.update(rate_by_service, decision: decision)
+        end
+
+        # Check if the Priority Sampling decision is to keep or drop the trace.
+        # Other factors can influence the sampling decision; this method is only
+        # responsible for interpreting the Sampling Priority decision.
+        #
+        # @param priority_sampling [Integer] priority sampling number
+        # @return [Boolean] true if trace is "kept" by priority sampling
+        # @return [Boolean] false if trace is "dropped" by priority sampling
+        def self.sampled?(priority_sampling)
+          priority_sampling >= Ext::Priority::AUTO_KEEP
+        end
+
+        private
+
+        def pre_sample?(trace)
+          case @pre_sampler
+          when RateSampler
+            @pre_sampler.sample_rate < 1.0
+          when RateByServiceSampler
+            @pre_sampler.sample_rate(trace) < 1.0
+          else
+            true
+          end
+        end
+
+        def priority_assigned?(trace)
+          !trace.sampling_priority.nil?
+        end
+
+        def priority_sample!(trace)
+          preserving_sampling(trace) do
+            @priority_sampler.sample!(trace)
+          end
+        end
+
+        # Ensures the trace's priority sampling decision is not changed by the @pre_sampler.
+        # The @pre_sampler should only change `trace.sampled`.
+        def preserving_priority_sampling(trace)
+          sampling_priority = trace.sampling_priority
+          sampling_decision = trace.get_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER)
+
+          yield.tap do
+            trace.sampling_priority = sampling_priority
+
+            if sampling_decision
+              trace.set_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER, sampling_decision)
+            else
+              trace.clear_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER)
+            end
+          end
+        end
+
+        # Ensures the trace is always propagated to the writer and that
+        # the sample rate metric represents the true client-side sampling.
+        def preserving_sampling(trace)
+          pre_sample_rate_metric = trace.sample_rate
+
+          yield.tap do
+            # NOTE: We'll want to leave `trace.sampled = true` here; all spans for priority sampling must
+            #       be sent to the agent. Otherwise metrics for traces will not be accurate, since the
+            #       agent will have an incomplete dataset.
+            #
+            #       We also ensure that the agent knows we that our `post_sampler` is not performing true sampling,
+            #       to avoid erroneous metric upscaling.
+            trace.sampled = true
+
+            # Restore true sampling metric, as only the @pre_sampler can reject traces.
+            # otherwise if @pre_sampler is not enabled, sending this metric would be misleading.
+            trace.sample_rate = pre_sample_rate_metric || nil
+          end
+        end
+
+        def assign_priority!(trace, priority)
+          trace.sampling_priority = priority
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/sampling/rate_by_key_sampler.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+require_relative 'sampler'
+require_relative 'rate_sampler'
+
+module Datadog
+  module Tracing
+    module Sampling
+      # Samples at different rates by key.
+      class RateByKeySampler < Sampler
+        attr_reader \
+          :default_key
+
+        def initialize(default_key, default_rate = 1.0, decision: nil, &block)
+          super()
+
+          raise ArgumentError, 'No resolver given!' unless block
+
+          @default_key = default_key
+          @resolver = block
+          @mutex = Mutex.new
+          @samplers = {}
+
+          set_rate(default_key, default_rate, decision)
+        end
+
+        def resolve(trace)
+          @resolver.call(trace)
+        end
+
+        def default_sampler
+          @samplers[default_key]
+        end
+
+        def sample!(trace)
+          key = resolve(trace)
+
+          @mutex.synchronize do
+            @samplers.fetch(key, default_sampler).sample!(trace)
+          end
+        end
+
+        def sample_rate(trace)
+          key = resolve(trace)
+
+          @mutex.synchronize do
+            @samplers.fetch(key, default_sampler).sample_rate
+          end
+        end
+
+        def update(key, rate, decision: nil)
+          @mutex.synchronize do
+            set_rate(key, rate, decision)
+          end
+        end
+
+        def update_all(rate_by_key, decision: nil)
+          @mutex.synchronize do
+            rate_by_key.each { |key, rate| set_rate(key, rate, decision) }
+          end
+        end
+
+        def delete(key)
+          @mutex.synchronize do
+            @samplers.delete(key)
+          end
+        end
+
+        def delete_if(&block)
+          @mutex.synchronize do
+            @samplers.delete_if(&block)
+          end
+        end
+
+        def length
+          @samplers.length
+        end
+
+        private
+
+        def set_rate(key, rate, decision)
+          @samplers[key] = RateSampler.new(rate, decision: decision)
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/sampling/rate_by_service_sampler.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require_relative 'rate_by_key_sampler'
+
+module Datadog
+  module Tracing
+    module Sampling
+      # {Datadog::Tracing::Sampling::RateByServiceSampler} samples different services at different rates
+      class RateByServiceSampler < RateByKeySampler
+        DEFAULT_KEY = 'service:,env:'
+
+        def initialize(default_rate = 1.0, env: nil, decision: Datadog::Tracing::Sampling::Ext::Decision::DEFAULT)
+          super(
+            DEFAULT_KEY,
+            default_rate,
+            decision: decision,
+            &method(:key_for)
+          )
+
+          @env = env
+        end
+
+        def update(rate_by_service, decision: nil)
+          # Remove any old services
+          delete_if { |key, _| key != DEFAULT_KEY && !rate_by_service.key?(key) }
+
+          # Update each service rate
+          update_all(rate_by_service, decision: decision)
+
+          # Emit metric for service cache size
+          Datadog.health_metrics.sampling_service_cache_length(length)
+        end
+
+        private
+
+        # DEV: Creating a string on every trace to perform a single Hash lookup is expensive.
+        #
+        # Using 2 nested hashes: 1 for env and 1 for service is the fastest option.
+        # This approach requires large API changes to `RateByKeySampler`.
+        #
+        # Reducing the interpolated string size, by using a 1 character separator,
+        # is also measurably faster than the current method. This approach does not
+        # require changes to `RateByKeySampler`.
+        #
+        # Keep in mind that these changes also require changes to `#update`.
+        #
+        # Comparison:
+        #  2 nested hashes: `service_hash.fetch(service, {}).fetch(env, default_rate)`
+        #                   7730045 i/s
+        # 1 char separator: `hash.fetch("#{service}\0#{env}", default_rate)`
+        #                   4302801 i/s - 1.80x slower
+        #          current: `hash.fetch("service:#{service},env:#{env}", default_rate)`
+        #                   2720459 i/s - 2.84x slower
+        def key_for(trace)
+          # Resolve env dynamically, if Proc is given.
+          env = @env.is_a?(Proc) ? @env.call : @env
+
+          "service:#{trace.service},env:#{env}"
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/sampling/rate_limiter.rb

@@ -0,0 +1,185 @@
+# frozen_string_literal: true
+
+require_relative '../../core/utils/time'
+
+module Datadog
+  module Tracing
+    module Sampling
+      # Checks for rate limiting on a resource.
+      class RateLimiter
+        # Checks if resource of specified size can be
+        # conforms with the current limit.
+        #
+        # Implementations of this method are not guaranteed
+        # to be side-effect free.
+        #
+        # @return [Boolean] whether a resource conforms with the current limit
+        def allow?(size); end
+
+        # The effective rate limiting ratio based on
+        # recent calls to `allow?`.
+        #
+        # @return [Float] recent allowance ratio
+        def effective_rate; end
+      end
+
+      # Implementation of the Token Bucket metering algorithm
+      # for rate limiting.
+      #
+      # @see https://en.wikipedia.org/wiki/Token_bucket Token bucket
+      class TokenBucket < RateLimiter
+        attr_reader :rate, :max_tokens
+
+        # @param rate [Numeric] Allowance rate, in units per second
+        #  if rate is negative, always allow
+        #  if rate is zero, never allow
+        # @param max_tokens [Numeric] Limit of available tokens
+        def initialize(rate, max_tokens = rate)
+          super()
+
+          raise ArgumentError, "rate must be a number: #{rate}" unless rate.is_a?(Numeric)
+          raise ArgumentError, "max_tokens must be a number: #{max_tokens}" unless max_tokens.is_a?(Numeric)
+
+          @rate = rate
+          @max_tokens = max_tokens
+
+          @tokens = max_tokens
+          @total_messages = 0
+          @conforming_messages = 0
+          @prev_conforming_messages = nil
+          @prev_total_messages = nil
+          @current_window = nil
+
+          @last_refill = Core::Utils::Time.get_time
+        end
+
+        # Checks if a message of provided +size+
+        # conforms with the current bucket limit.
+        #
+        # If it does, return +true+ and remove +size+
+        # tokens from the bucket.
+        # If it does not, return +false+ without affecting
+        # the tokens from the bucket.
+        #
+        # @return [Boolean] +true+ if message conforms with current bucket limit
+        def allow?(size)
+          allowed = should_allow?(size)
+          update_rate_counts(allowed)
+          allowed
+        end
+
+        # Ratio of 'conformance' per 'total messages' checked
+        # averaged for the past 2 buckets
+        #
+        # Returns +1.0+ when no messages have been checked yet.
+        #
+        # @return [Float] Conformance ratio, between +[0,1]+
+        def effective_rate
+          return 0.0 if @rate.zero?
+          return 1.0 if @rate < 0 || @total_messages.zero?
+
+          return current_window_rate if @prev_conforming_messages.nil? || @prev_total_messages.nil?
+
+          (@conforming_messages.to_f + @prev_conforming_messages.to_f) / (@total_messages + @prev_total_messages)
+        end
+
+        # Ratio of 'conformance' per 'total messages' checked
+        # on this bucket
+        #
+        # Returns +1.0+ when no messages have been checked yet.
+        #
+        # @return [Float] Conformance ratio, between +[0,1]+
+        def current_window_rate
+          return 1.0 if @total_messages.zero?
+
+          @conforming_messages.to_f / @total_messages
+        end
+
+        # @return [Numeric] number of tokens currently available
+        def available_tokens
+          @tokens
+        end
+
+        private
+
+        def refill_since_last_message
+          now = Core::Utils::Time.get_time
+          elapsed = now - @last_refill
+
+          # Update the number of available tokens, but ensure we do not exceed the max
+          # we return the min of tokens + rate*elapsed, or max tokens
+          refill_tokens(@rate * elapsed)
+
+          @last_refill = now
+        end
+
+        def refill_tokens(size)
+          @tokens += size
+          @tokens = @max_tokens if @tokens > @max_tokens
+        end
+
+        def increment_total_count
+          @total_messages += 1
+        end
+
+        def increment_conforming_count
+          @conforming_messages += 1
+        end
+
+        def should_allow?(size)
+          # rate limit of 0 blocks everything
+          return false if @rate.zero?
+
+          # negative rate limit disables rate limiting
+          return true if @rate < 0
+
+          refill_since_last_message
+
+          # if tokens < 1 we don't allow?
+          return false if @tokens < size
+
+          @tokens -= size
+
+          true
+        end
+
+        # Sets and Updates the past two 1 second windows for which
+        # the rate limiter must compute it's rate over and updates
+        # the total count, and conforming message count if +allowed+
+        def update_rate_counts(allowed)
+          now = Core::Utils::Time.get_time
+
+          # No tokens have been seen yet, start a new window
+          if @current_window.nil?
+            @current_window = now
+          # If more than 1 second has past since last window, reset
+          elsif now - @current_window >= 1
+            @prev_conforming_messages = @conforming_messages
+            @prev_total_messages = @total_messages
+            @conforming_messages = 0
+            @total_messages = 0
+            @current_window = now
+          end
+
+          increment_conforming_count if allowed
+
+          increment_total_count
+        end
+      end
+
+      # {Datadog::Tracing::Sampling::RateLimiter} that accepts all resources,
+      # with no limits.
+      class UnlimitedLimiter < RateLimiter
+        # @return [Boolean] always +true+
+        def allow?(_)
+          true
+        end
+
+        # @return [Float] always 100%
+        def effective_rate
+          1.0
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/sampling/rate_sampler.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+require_relative 'sampler'
+require_relative '../utils'
+
+module Datadog
+  module Tracing
+    module Sampling
+      # {Datadog::Tracing::Sampling::RateSampler} is based on a sample rate.
+      class RateSampler < Sampler
+        KNUTH_FACTOR = 1111111111111111111
+
+        # Initialize a {Datadog::Tracing::Sampling::RateSampler}.
+        # This sampler keeps a random subset of the traces. Its main purpose is to
+        # reduce the instrumentation footprint.
+        #
+        # @param sample_rate [Numeric] the sample rate between 0.0 and 1.0, inclusive.
+        #   0.0 means that no trace will be sampled; 1.0 means that all traces will be  sampled.
+        def initialize(sample_rate = 1.0, decision: nil)
+          super()
+
+          unless sample_rate >= 0.0 && sample_rate <= 1.0
+            Datadog.logger.error('sample rate is not between 0 and 1, falling back to 1')
+            sample_rate = 1.0
+          end
+
+          self.sample_rate = sample_rate
+
+          @decision = decision
+        end
+
+        def sample_rate(*_)
+          @sample_rate
+        end
+
+        def sample_rate=(sample_rate)
+          @sample_rate = sample_rate
+          @sampling_id_threshold = sample_rate * Tracing::Utils::EXTERNAL_MAX_ID
+        end
+
+        def sample?(trace)
+          ((trace.id * KNUTH_FACTOR) % Tracing::Utils::EXTERNAL_MAX_ID) <= @sampling_id_threshold
+        end
+
+        def sample!(trace)
+          sampled = sample?(trace)
+
+          return false unless sampled
+
+          trace.sample_rate = @sample_rate
+          trace.set_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER, @decision) if @decision
+
+          true
+        end
+      end
+    end
+  end
+end

data/lib/datadog/tracing/sampling/rule.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+require_relative 'matcher'
+require_relative 'rate_sampler'
+
+module Datadog
+  module Tracing
+    module Sampling
+      # Sampling rule that dictates if a trace matches
+      # a specific criteria and what sampling strategy to
+      # apply in case of a positive match.
+      class Rule
+        attr_reader :matcher, :sampler
+
+        # @param [Matcher] matcher A matcher to verify trace conformity against
+        # @param [Sampler] sampler A sampler to be consulted on a positive match
+        def initialize(matcher, sampler)
+          @matcher = matcher
+          @sampler = sampler
+        end
+
+        # Evaluates if the provided `trace` conforms to the `matcher`.
+        #
+        # @param [TraceOperation] trace
+        # @return [Boolean] whether this rules applies to the trace
+        # @return [NilClass] if the matcher fails errs during evaluation
+        def match?(trace)
+          @matcher.match?(trace)
+        rescue => e
+          Datadog.logger.error(
+            "Matcher failed. Cause: #{e.class.name} #{e.message} Source: #{Array(e.backtrace).first}"
+          )
+          nil
+        end
+
+        # (see Datadog::Tracing::Sampling::Sampler#sample!)
+        def sample!(trace)
+          @sampler.sample!(trace)
+        end
+
+        # (see Datadog::Tracing::Sampling::Sampler#sample_rate)
+        def sample_rate(trace)
+          @sampler.sample_rate(trace)
+        end
+      end
+
+      # A {Datadog::Tracing::Sampling::Rule} that matches a trace based on
+      # trace name and/or service name and
+      # applies a fixed sampling to matching spans.
+      class SimpleRule < Rule
+        # @param name [String,Regexp,Proc] Matcher for case equality (===) with the trace name, defaults to always match
+        # @param service [String,Regexp,Proc] Matcher for case equality (===) with the service name,
+        #                defaults to always match
+        # @param sample_rate [Float] Sampling rate between +[0,1]+
+        def initialize(name: SimpleMatcher::MATCH_ALL, service: SimpleMatcher::MATCH_ALL, sample_rate: 1.0)
+          super(SimpleMatcher.new(name: name, service: service), RateSampler.new(sample_rate))
+        end
+      end
+    end
+  end
+end