Changed RuleSampler initialization with ASM Standalone to Tracing::Co…

…mponent.build_sampler
DataDog · Oct 21, 2024 · e5f552f · e5f552f
1 parent 73f8af2
commit e5f552f
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 30 deletions.
diff --git a/lib/datadog/tracing/component.rb b/lib/datadog/tracing/component.rb
@@ -73,6 +73,24 @@ def build_sampler(settings)
           return sampler
         end
 
+        # AppSec events are sent to the backend using traces.
+        # Standalone ASM billing means that we don't want to charge clients for APM traces,
+        # so we want to send the minimum amount of traces possible (idealy only traces that contains security events),
+        # but for features such as API Security, we need to send at least one trace per minute,
+        # to keep the service alive on the backend side.
+        if settings.appsec.standalone.enabled
+          post_sampler = Tracing::Sampling::RuleSampler.new(
+            [Tracing::Sampling::SimpleRule.new(sample_rate: 1.0)],
+            rate_limit: nil,
+            rate_limiter: Datadog::Core::TokenBucket.new(1.0 / 60, 1.0),
+            default_sample_rate: nil,
+            # This is a hack, as we would like to set the default sampler to nil, but setting it to nil would
+            # create a RateByServiceSampler. 0 is not a RateByServiceSampler and does not respond to :update,
+            # so the result is the same as if @default_sampler was set to nil.
+            default_sampler: 0
+          )
+        end
+
         # Sampling rules are provided
         if (rules = settings.tracing.sampling.rules)
           post_sampler = Tracing::Sampling::RuleSampler.parse(

diff --git a/lib/datadog/tracing/sampling/rule_sampler.rb b/lib/datadog/tracing/sampling/rule_sampler.rb
@@ -29,34 +29,23 @@ def initialize(
           default_sample_rate: Datadog.configuration.tracing.sampling.default_rate,
           default_sampler: nil
         )
-          # AppSec events are sent to the backend using traces.
-          # Standalone ASM billing means that we don't want to charge clients for APM traces,
-          # so we want to send the minimum amount of traces possible (idealy only traces that contains security events),
-          # but for features such as API Security, we need to send at least one trace per minute,
-          # to keep the service alive on the backend side.
-          @rules = if Datadog.configuration.appsec.standalone.enabled
-                     [SimpleRule.new(sample_rate: 1.0)]
-                   elsif default_sample_rate && !default_sampler
+          @rules = if default_sample_rate && !default_sampler
                      # Add to the end of the rule list a rule always matches any trace
                      rules << SimpleRule.new(sample_rate: default_sample_rate)
                    else
                      rules
                    end
-          @rate_limiter = if Datadog.configuration.appsec.standalone.enabled
-                            # 1 trace per minute
-                            Core::TokenBucket.new(1.0 / 60, 1.0)
-                          elsif rate_limiter
+          @rate_limiter = if rate_limiter
                             rate_limiter
                           elsif rate_limit
                             Core::TokenBucket.new(rate_limit)
                           else
                             Core::UnlimitedLimiter.new
                           end
-          @default_sampler = if Datadog.configuration.appsec.standalone.enabled ||
-              (default_sample_rate && !default_sampler)
-                               nil
-                             elsif default_sampler
+          @default_sampler = if default_sampler
                                default_sampler
+                             elsif default_sample_rate
+                               nil
                              else
                                # TODO: Simplify .tags access, as `Tracer#tags` can't be arbitrarily changed anymore
                                RateByServiceSampler.new(1.0, env: -> { Tracing.send(:tracer).tags['env'] })

diff --git a/spec/datadog/tracing/sampling/rule_sampler_spec.rb b/spec/datadog/tracing/sampling/rule_sampler_spec.rb
@@ -42,20 +42,6 @@
     it_behaves_like 'a token bucket rate limiter', rate: 100
     it { expect(rule_sampler.default_sampler).to be_a(Datadog::Tracing::Sampling::RateByServiceSampler) }
 
-    context 'with appsec standalone enabled' do
-      before do
-        allow(Datadog.configuration.appsec.standalone).to receive(:enabled).and_return(true)
-      end
-
-      it_behaves_like 'a simple rule that matches all span operations', sample_rate: 1.0 do
-        let(:rule) { rule_sampler.rules.last }
-      end
-
-      it_behaves_like 'a token bucket rate limiter', rate: 1.0 / 60, max_tokens: 1.0
-
-      it { expect(rule_sampler.default_sampler).to be_nil }
-    end
-
     context 'with rate_limit ENV' do
       before do
         allow(Datadog.configuration.tracing.sampling).to receive(:rate_limit)