[APMAPI-1258] Adjust trace sampling formula (#4616)

Author: genesor

lib/datadog/tracing/distributed/b3_multi.rb

@@ -55,7 +55,7 @@ def extract(data)
           span_id = Helpers.parse_hex_id(fetcher[@span_id_key])
 
           # Return early if this propagation is not valid
-          return if span_id.nil? || span_id <= 0 || span_id >= Tracing::Utils::EXTERNAL_MAX_ID
+          return if span_id.nil? || span_id <= 0 || span_id > Tracing::Utils::EXTERNAL_MAX_ID
 
           # We don't need to try and convert sampled since B3 supports 0/1 (AUTO_REJECT/AUTO_KEEP)
           sampling_priority = Helpers.parse_decimal_id(fetcher[@sampled_key])

lib/datadog/tracing/distributed/b3_single.rb

@@ -54,7 +54,7 @@ def extract(env)
 
           span_id = Helpers.parse_hex_id(parts[1]) if parts.length > 1
           # Return early if this propagation is not valid
-          return if span_id.nil? || span_id <= 0 || span_id >= Tracing::Utils::EXTERNAL_MAX_ID
+          return if span_id.nil? || span_id <= 0 || span_id > Tracing::Utils::EXTERNAL_MAX_ID
 
           sampling_priority = Helpers.parse_decimal_id(parts[2]) if parts.length > 2
 

lib/datadog/tracing/distributed/datadog.rb

@@ -89,7 +89,7 @@ def parse_trace_id(fetcher_object)
           trace_id = Helpers.parse_decimal_id(fetcher_object[@trace_id_key])
 
           return unless trace_id
-          return if trace_id <= 0 || trace_id >= Tracing::Utils::EXTERNAL_MAX_ID
+          return if trace_id <= 0 || trace_id > Tracing::Utils::EXTERNAL_MAX_ID
 
           trace_id
         end
@@ -98,7 +98,7 @@ def parse_parent_id(fetcher_object)
           parent_id = Helpers.parse_decimal_id(fetcher_object[@parent_id_key])
 
           return unless parent_id
-          return if parent_id <= 0 || parent_id >= Tracing::Utils::EXTERNAL_MAX_ID
+          return if parent_id <= 0 || parent_id > Tracing::Utils::EXTERNAL_MAX_ID
 
           parent_id
         end

lib/datadog/tracing/sampling/rate_sampler.rb

@@ -9,6 +9,7 @@ module Sampling
       # {Datadog::Tracing::Sampling::RateSampler} is based on a sample rate.
       class RateSampler < Sampler
         KNUTH_FACTOR = 1111111111111111111
+        UINT64_MODULO = (1 << 64)
 
         # Initialize a {Datadog::Tracing::Sampling::RateSampler}.
         # This sampler keeps a random subset of the traces. Its main purpose is to
@@ -39,7 +40,7 @@ def sample_rate=(sample_rate)
         end
 
         def sample?(trace)
-          ((trace.id * KNUTH_FACTOR) % Tracing::Utils::EXTERNAL_MAX_ID) <= @sampling_id_threshold
+          ((trace.id * KNUTH_FACTOR) % UINT64_MODULO) <= @sampling_id_threshold
         end
 
         def sample!(trace)

lib/datadog/tracing/utils.rb

@@ -24,7 +24,7 @@ module Utils
 
       # While we only generate 63-bit integers due to limitations in other languages, we support
       # parsing 64-bit integers for distributed tracing since an upstream system may generate one
-      EXTERNAL_MAX_ID = 1 << 64
+      EXTERNAL_MAX_ID = (1 << 64) - 1
 
       # We use a custom random number generator because we want no interference
       # with the default one. Using the default prng, we could break code that

sig/datadog/tracing/sampling/rate_sampler.rbs

@@ -2,7 +2,8 @@ module Datadog
   module Tracing
     module Sampling
       class RateSampler < Sampler
-        KNUTH_FACTOR: 1111111111111111111
+        KNUTH_FACTOR: Integer
+        UINT64_MODULO: Integer
         def initialize: (?::Float sample_rate, ?decision: untyped?) -> void
 
         def sample_rate: (*untyped _) -> untyped

spec/datadog/tracing/sampling/rate_sampler_spec.rb

@@ -122,5 +122,30 @@
         expect(trace.get_tag('_dd.p.dm')).to be_nil
       end
     end
+
+    context 'with specific trace IDs and 0.5 sample rate' do
+      let(:sample_rate) { 0.5 }
+      let(:trace_expectations) do
+        {
+          12078589664685934330 => false,
+          13794769880582338323 => true,
+          1882305164521835798 => true,
+          5198373796167680436 => false,
+          6272545487220484606 => true,
+          8696342848850656916 => true,
+          18444899399302180860 => false,
+          18444899399302180862 => true,
+          9223372036854775808 => true, # Lands exactly on the sampling threshold 0.5 * MAX_UINT64
+        }
+      end
+
+      it 'produces deterministic sampling results' do
+        trace_expectations.each do |id, expected|
+          trace = Datadog::Tracing::TraceOperation.new(id: id)
+          expect(sampler.sample!(trace)).to eq(expected),
+            "Expected trace ID #{id} to be #{expected ? 'sampled' : 'not sampled'}"
+        end
+      end
+    end
   end
 end