Add dynamically configurable sampling rules

Signed-off-by: Marco Costa <[email protected]>

lib/datadog/tracing/configuration/dynamic.rb

@@ -53,8 +53,46 @@ def configuration_object
           end
         end
 
+        # Dynamic configuration for `DD_TRACE_SAMPLING_RULES`.
+        class TracingSamplingRules < SimpleOption
+          def initialize
+            super('tracing_sampling_rules', 'DD_TRACE_SAMPLING_RULES', :rules)
+          end
+
+          # Ensures sampler is rebuilt and new configuration is applied
+          def call(tracing_sampling_rules)
+            # Modify the remote configuration value that it matches the
+            # local environment variable it configures.
+            if tracing_sampling_rules
+              tracing_sampling_rules.each do |rule|
+                next unless (tags = rule['tags'])
+
+                # Tag maps come in as arrays of 'key' and `value_glob`.
+                # We need to convert them into a hash for local use.
+                tag_array = tags.map! do |tag|
+                  [tag['key'], tag['value_glob']]
+                end
+
+                rule['tags'] = tag_array.to_h
+              end
+
+              # The configuration is stored as JSON, so we need to convert it back
+              tracing_sampling_rules = tracing_sampling_rules.to_json
+            end
+
+            super(tracing_sampling_rules)
+            Datadog.send(:components).reconfigure_live_sampler
+          end
+
+          protected
+
+          def configuration_object
+            Datadog.configuration.tracing.sampling
+          end
+        end
+
         # List of all tracing dynamic configuration options supported.
-        OPTIONS = [LogInjectionEnabled, TracingHeaderTags, TracingSamplingRate].map do |option_class|
+        OPTIONS = [LogInjectionEnabled, TracingHeaderTags, TracingSamplingRate, TracingSamplingRules].map do |option_class|
           option = option_class.new
           [option.name, option.env_var, option]
         end

lib/datadog/tracing/configuration/settings.rb

@@ -327,6 +327,7 @@ def self.extended(base)
                 # @return [String,nil]
                 # @public_api
                 option :rules do |o|
+                  o.type :string, nilable: true
                   o.default { ENV.fetch(Configuration::Ext::Sampling::ENV_RULES, nil) }
                 end
 

lib/datadog/tracing/remote.rb

@@ -12,12 +12,16 @@ class ReadError < StandardError; end
       class << self
         PRODUCT = 'APM_TRACING'
 
+        CAPABILITIES = [
+          1 << 29 # APM_TRACING_SAMPLE_RULES: Dynamic trace sampling rules configuration
+        ].freeze
+
         def products
           [PRODUCT]
         end
 
         def capabilities
-          [] # No capabilities advertised
+          CAPABILITIES
         end
 
         def process_config(config, content)

lib/datadog/tracing/sampling/ext.rb

@@ -40,7 +40,7 @@ module Decision
           DEFAULT = '-0'
           # The sampling rate received in the agent's http response.
           AGENT_RATE = '-1'
-          # Sampling rule or sampling rate based on tracer config.
+          # Locally configured rule.
           TRACE_SAMPLING_RULE = '-3'
           # User directly sets sampling priority via {Tracing.reject!} or {Tracing.keep!},
           # or by a custom sampler implementation.
@@ -49,6 +49,10 @@ module Decision
           ASM = '-5'
           # Single Span Sampled.
           SPAN_SAMPLING_RATE = '-8'
+          # Dynamically configured rule, explicitly created by the user.
+          REMOTE_USER_RULE = '-10'
+          # Dynamically configured  rule, automatically generated by Datadog.
+          REMOTE_DYNAMIC_RULE = '-11'
         end
       end
     end

lib/datadog/tracing/sampling/rule.rb

@@ -9,13 +9,18 @@ module Sampling
       # apply in case of a positive match.
       # @public_api
       class Rule
-        attr_reader :matcher, :sampler
+        PROVENANCE_LOCAL = :local
+        PROVENANCE_REMOTE_USER = :customer
+        PROVENANCE_REMOTE_DYNAMIC = :dynamic
+
+        attr_reader :matcher, :sampler, :provenance
 
         # @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)
+        def initialize(matcher, sampler, provenance)
           @matcher = matcher
           @sampler = sampler
+          @provenance = provenance
         end
 
         # Evaluates if the provided `trace` conforms to the `matcher`.
@@ -56,7 +61,9 @@ class SimpleRule < Rule
         # @param sample_rate [Float] Sampling rate between +[0,1]+
         def initialize(
           name: SimpleMatcher::MATCH_ALL, service: SimpleMatcher::MATCH_ALL,
-          resource: SimpleMatcher::MATCH_ALL, tags: {}, sample_rate: 1.0
+          resource: SimpleMatcher::MATCH_ALL, tags: {},
+          provenance: Rule::PROVENANCE_LOCAL,
+          sample_rate: 1.0
         )
           # We want to allow 0.0 to drop all traces, but {Datadog::Tracing::Sampling::RateSampler}
           # considers 0.0 an invalid rate and falls back to 100% sampling.
@@ -69,7 +76,7 @@ def initialize(
           sampler = RateSampler.new
           sampler.sample_rate = sample_rate
 
-          super(SimpleMatcher.new(name: name, service: service, resource: resource, tags: tags), sampler)
+          super(SimpleMatcher.new(name: name, service: service, resource: resource, tags: tags), sampler, provenance)
         end
       end
     end

lib/datadog/tracing/sampling/rule_sampler.rb

@@ -64,6 +64,12 @@ def self.parse(rules, rate_limit, default_sample_rate)
               resource: rule['resource'],
               tags: rule['tags'],
               sample_rate: sample_rate,
+              provenance: if (provenance = rule['provenance'])
+                            # `Rule::PROVENANCE_*` values are symbols, so convert strings to match
+                            provenance.to_sym
+                          else
+                            Rule::PROVENANCE_LOCAL
+                          end,
             }
 
             Core::BackportFrom24.hash_compact!(kwargs)
@@ -127,7 +133,17 @@ def sample_trace(trace)
           rate_limiter.allow?(1).tap do |allowed|
             set_priority(trace, allowed)
             set_limiter_metrics(trace, rate_limiter.effective_rate)
-            trace.set_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER, Ext::Decision::TRACE_SAMPLING_RULE)
+
+            provenance = case rule.provenance
+                         when Rule::PROVENANCE_REMOTE_USER
+                           Ext::Decision::REMOTE_USER_RULE
+                         when Rule::PROVENANCE_REMOTE_DYNAMIC
+                           Ext::Decision::REMOTE_DYNAMIC_RULE
+                         else
+                           Ext::Decision::TRACE_SAMPLING_RULE
+                         end
+
+            trace.set_tag(Tracing::Metadata::Ext::Distributed::TAG_DECISION_MAKER, provenance)
           end
         rescue StandardError => e
           Datadog.logger.error(

sig/datadog/core/configuration/options.rbs

@@ -19,14 +19,16 @@ module Datadog
         module InstanceMethods
           def options: () -> untyped
 
-          def set_option: (untyped name, untyped value) -> untyped
+          def set_option: (untyped name, untyped value, precedence: Option::Precedence::Value) -> untyped
 
           def get_option: (untyped name) -> untyped
 
           def reset_option: (untyped name) -> untyped
 
           def option_defined?: (untyped name) -> untyped
 
+          def unset_option: (untyped name, precedence: Option::Precedence::Value) -> void
+
           def using_default?: (Symbol option) -> bool
 
           def options_hash: () -> untyped

sig/datadog/tracing/configuration/dynamic.rbs

@@ -14,7 +14,13 @@ module Datadog
           def initialize: () -> void
           def call: (untyped tracing_sampling_rate) -> untyped
 
-          def configuration_object: () -> untyped
+          def configuration_object: () -> Core::Configuration::Options::InstanceMethods
+        end
+
+        class TracingSamplingRules < SimpleOption
+          def initialize: () -> void
+          def call: (Array[Hash[String, untyped]] tracing_sampling_rules) -> void
+          def configuration_object: () -> Core::Configuration::Options::InstanceMethods
         end
 
         # Correct type is `OPTIONS: Array[[String, String, Option]]`

sig/datadog/tracing/configuration/dynamic/option.rbs

@@ -14,7 +14,7 @@ module Datadog
 
           def initialize: (String name, String env_var, Symbol setting_key) -> void
           def call: (Object value) -> void
-          def configuration_object: () -> untyped
+          def configuration_object: () -> Core::Configuration::Options::InstanceMethods
         end
       end
     end

sig/datadog/tracing/remote.rbs

@@ -5,6 +5,7 @@ module Datadog
       end
 
       PRODUCT: "APM_TRACING"
+      CAPABILITIES: Array[Integer]
 
       def self.products: () -> ::Array[String]
 

sig/datadog/tracing/sampling/ext.rbs

@@ -16,6 +16,8 @@ module Datadog
         module Decision
           DEFAULT: ::String
           AGENT_RATE: ::String
+          REMOTE_DYNAMIC_RULE: ::String
+          REMOTE_USER_RULE: ::String
           TRACE_SAMPLING_RULE: ::String
           MANUAL: ::String
           ASM: ::String

sig/datadog/tracing/sampling/rule_sampler.rbs

@@ -13,6 +13,8 @@ module Datadog
         def sample!: (untyped trace) -> untyped
         def update: (*untyped args, **untyped kwargs) -> (false | untyped)
 
+        def self.parse: (Array[Hash[Symbol, untyped]] rules, Float rate_limit, Float default_sample_rate) -> RuleSampler?
+
         private
 
         def sample_trace: (untyped trace) { (untyped) -> untyped } -> untyped

spec/datadog/core/remote/client/capabilities_spec.rb

@@ -33,8 +33,8 @@
       end
 
       describe '#base64_capabilities' do
-        it 'returns an empty string' do
-          expect(capabilities.base64_capabilities).to eq('')
+        it 'matches tracing capabilities only' do
+          expect(capabilities.base64_capabilities).to eq('IAAAAA==')
         end
       end
     end

spec/datadog/tracing/configuration/dynamic_spec.rb

@@ -79,3 +79,36 @@
     option.call(new_value)
   end
 end
+
+RSpec.describe Datadog::Tracing::Configuration::Dynamic::TracingSamplingRules do
+  let(:old_value) { nil }
+
+  include_examples 'tracing dynamic simple option',
+    name: 'tracing_sampling_rules',
+    env_var: 'DD_TRACE_SAMPLING_RULES',
+    config_key: :rules,
+    value: RSpec::Matchers::BuiltIn::Match.new(->(rules) { rules == '[{"sample_rate":1}]' }),
+    config_object: Datadog.configuration.tracing.sampling do
+      let(:new_value) { [{ sample_rate: 1 }] }
+    end
+
+  context 'with tags' do
+    include_examples 'tracing dynamic simple option',
+      name: 'tracing_sampling_rules',
+      env_var: 'DD_TRACE_SAMPLING_RULES',
+      config_key: :rules,
+      value: RSpec::Matchers::BuiltIn::Match.new(
+        lambda do |rules|
+          rules == '[{"sample_rate":1,"tags":[{"key":"k","value_glob":"v"}]}]'
+        end
+      ),
+      config_object: Datadog.configuration.tracing.sampling do
+        let(:new_value) { [{ sample_rate: 1, tags: [{ key: 'k', value_glob: 'v' }] }] }
+      end
+  end
+
+  it 'reconfigures the live sampler' do
+    expect(Datadog.send(:components)).to receive(:reconfigure_live_sampler)
+    option.call(new_value)
+  end
+end

spec/datadog/tracing/remote_spec.rb

@@ -8,8 +8,8 @@
     expect(remote.products).to contain_exactly('APM_TRACING')
   end
 
-  it 'declares no capabilities' do
-    expect(remote.capabilities).to be_empty
+  it 'declares rule sampling capabilities' do
+    expect(remote.capabilities).to eq([1 << 29])
   end
 
   it 'declares matches that match APM_TRACING' do
@@ -46,7 +46,8 @@
             .with(contain_exactly(
               ['DD_LOGS_INJECTION', nil],
               ['DD_TRACE_HEADER_TAGS', nil],
-              ['DD_TRACE_SAMPLE_RATE', nil]
+              ['DD_TRACE_SAMPLE_RATE', nil],
+              ['DD_TRACE_SAMPLING_RULES', nil],
             ))
 
           process_config
@@ -64,7 +65,8 @@
             .with(contain_exactly(
               ['DD_LOGS_INJECTION', false],
               ['DD_TRACE_HEADER_TAGS', nil],
-              ['DD_TRACE_SAMPLE_RATE', nil]
+              ['DD_TRACE_SAMPLE_RATE', nil],
+              ['DD_TRACE_SAMPLING_RULES', nil],
             ))
 
           process_config

spec/datadog/tracing/sampling/rule_sampler_spec.rb

@@ -100,6 +100,9 @@
       it 'parses as a match any' do
         expect(actual_rule.matcher.name).to eq(Datadog::Tracing::Sampling::SimpleMatcher::MATCH_ALL)
         expect(actual_rule.matcher.service).to eq(Datadog::Tracing::Sampling::SimpleMatcher::MATCH_ALL)
+        expect(actual_rule.matcher.resource).to eq(Datadog::Tracing::Sampling::SimpleMatcher::MATCH_ALL)
+        expect(actual_rule.matcher.tags).to eq({})
+        expect(actual_rule.provenance).to eq(Datadog::Tracing::Sampling::Rule::PROVENANCE_LOCAL)
         expect(actual_rule.sampler.sample_rate).to eq(0.1)
       end
 
@@ -149,6 +152,23 @@
           expect(actual_rules[1].sampler.sample_rate).to eq(0.2)
         end
       end
+
+      context 'with provenance' do
+        context 'from customer' do
+          let(:rule) { { sample_rate: 1, provenance: 'customer' } }
+          it 'parses the provenance' do
+            expect(actual_rule.provenance).to eq(Datadog::Tracing::Sampling::Rule::PROVENANCE_REMOTE_USER)
+          end
+        end
+
+        context 'from dynamic configuration' do
+          let(:rule) { { sample_rate: 1, provenance: 'dynamic' } }
+
+          it 'parses the provenance' do
+            expect(actual_rule.provenance).to eq(Datadog::Tracing::Sampling::Rule::PROVENANCE_REMOTE_DYNAMIC)
+          end
+        end
+      end
     end
 
     context 'with a non-float sample_rate' do
@@ -192,11 +212,13 @@
     let(:rules) { [rule] }
     let(:rule) { instance_double(Datadog::Tracing::Sampling::Rule) }
     let(:sample_rate) { 0.8 }
+    let(:provenance) { :local }
 
     before do
       allow(rule).to receive(:match?).with(trace).and_return(true)
       allow(rule).to receive(:sample?).with(trace).and_return(sampled)
       allow(rule).to receive(:sample_rate).with(trace).and_return(sample_rate)
+      allow(rule).to receive(:provenance).and_return(provenance)
     end
   end
 
@@ -244,6 +266,26 @@
             let(:expected_sampled) { true }
             let(:sampling_priority) { 2 }
           end
+
+          context 'when the rule is from a remote user' do
+            let(:provenance) { :customer }
+
+            it_behaves_like 'a sampled! trace' do
+              let(:expected_sampled) { true }
+              let(:sampling_priority) { 2 }
+              let(:sampling_decision) { '-10' }
+            end
+          end
+
+          context 'when the rule is dynamically configured' do
+            let(:provenance) { :dynamic }
+
+            it_behaves_like 'a sampled! trace' do
+              let(:expected_sampled) { true }
+              let(:sampling_priority) { 2 }
+              let(:sampling_decision) { '-11' }
+            end
+          end
         end
 
         context 'and rate limited' do