Fix priority sampling dropping traces

lib/ddtrace/sampler.rb

@@ -93,37 +93,49 @@ class PrioritySampler
     extend Forwardable
 
     def initialize(opts = {})
-      @base_sampler = opts[:base_sampler] || RateSampler.new
       @post_sampler = opts[:post_sampler] || RateByServiceSampler.new
     end
 
     def sample(span)
-      return perform_sampling(span) unless span.context
-      return sampled_by_upstream(span) if span.context.sampling_priority
-
-      perform_sampling(span).tap do |sampled|
-        span.context.sampling_priority = if sampled
-                                           Datadog::Ext::Priority::AUTO_KEEP
-                                         else
-                                           Datadog::Ext::Priority::AUTO_REJECT
-                                         end
+      # If we haven't sampled this trace yet, do so.
+      # Otherwise we want to keep whatever priority has already been assigned.
+      unless sampled_by_upstream(span)
+        # Use the underlying sampler to "roll the dice" and see how we assign it priority.
+        # This sampler derives rates from the agent, which updates the sampler's rates
+        # whenever traces are submitted to the agent.
+        perform_sampling(span).tap do |sampled|
+          value = sampled ? Datadog::Ext::Priority::AUTO_KEEP : Datadog::Ext::Priority::AUTO_REJECT
+
+          if span.context
+            span.context.sampling_priority = value
+          else
+            # Set the priority directly on the span instead, since otherwise
+            # it won't receive the appropriate tag.
+            span.set_metric(
+              Ext::DistributedTracing::SAMPLING_PRIORITY_KEY,
+              value
+            )
+          end
+        end
       end
+
+      # Priority sampling *always* marks spans as sampled, so we flush them to the agent.
+      # This will happen regardless of whether the trace is kept or ultimately rejected.
+      # Otherwise metrics for traces will not be accurate, since the agent will have an
+      # incomplete dataset.
+      span.sampled = true
     end
 
     def_delegators :@post_sampler, :update
 
     private
 
     def sampled_by_upstream(span)
-      span.sampled = priority_keep?(span.context.sampling_priority)
-    end
-
-    def priority_keep?(sampling_priority)
-      sampling_priority == Datadog::Ext::Priority::USER_KEEP || sampling_priority == Datadog::Ext::Priority::AUTO_KEEP
+      span.context && !span.context.sampling_priority.nil?
     end
 
     def perform_sampling(span)
-      @base_sampler.sample(span) && @post_sampler.sample(span)
+      @post_sampler.sample(span)
     end
   end
 end

lib/ddtrace/tracer.rb

@@ -139,7 +139,7 @@ def configure(options = {})
       # but neither are configured. Verify the sampler isn't already a
       # priority sampler too, so we don't wrap one with another.
       if priority_sampling != false && !@sampler.is_a?(PrioritySampler)
-        @sampler = PrioritySampler.new(base_sampler: @sampler)
+        @sampler = PrioritySampler.new
         @writer = Writer.new(priority_sampler: @sampler)
       elsif priority_sampling == false
         @sampler = sampler || Datadog::AllSampler.new if @sampler.is_a?(PrioritySampler)

spec/ddtrace/sampler_spec.rb

@@ -0,0 +1,177 @@
+require 'spec_helper'
+
+require 'ddtrace'
+require 'ddtrace/sampler'
+
+RSpec.describe Datadog::AllSampler do
+  subject(:sampler) { described_class.new }
+
+  before(:each) { Datadog::Tracer.log.level = Logger::FATAL }
+  after(:each) { Datadog::Tracer.log.level = Logger::WARN }
+
+  describe '#sample' do
+    let(:spans) do
+      [
+        Datadog::Span.new(nil, '', trace_id: 1),
+        Datadog::Span.new(nil, '', trace_id: 2),
+        Datadog::Span.new(nil, '', trace_id: 3)
+      ]
+    end
+
+    it 'samples all spans' do
+      spans.each do |span|
+        expect(sampler.sample(span)).to be true
+        expect(span.sampled).to be true
+      end
+    end
+  end
+end
+
+RSpec.describe Datadog::RateSampler do
+  subject(:sampler) { described_class.new(sample_rate) }
+
+  before(:each) { Datadog::Tracer.log.level = Logger::FATAL }
+  after(:each) { Datadog::Tracer.log.level = Logger::WARN }
+
+  describe '#initialize' do
+    context 'given a sample rate' do
+      context 'that is negative' do
+        let(:sample_rate) { -1.0 }
+        it { expect(sampler.sample_rate).to eq(1.0) }
+      end
+
+      context 'that is 0' do
+        let(:sample_rate) { 0.0 }
+        it { expect(sampler.sample_rate).to eq(1.0) }
+      end
+
+      context 'that is between 0 and 1.0' do
+        let(:sample_rate) { 0.5 }
+        it { expect(sampler.sample_rate).to eq(0.5) }
+      end
+
+      context 'that is greater than 1.0' do
+        let(:sample_rate) { 1.5 }
+        it { expect(sampler.sample_rate).to eq(1.0) }
+      end
+    end
+  end
+
+  describe '#sample' do
+    let(:spans) do
+      [
+        Datadog::Span.new(nil, '', trace_id: 1),
+        Datadog::Span.new(nil, '', trace_id: 2),
+        Datadog::Span.new(nil, '', trace_id: 3)
+      ]
+    end
+
+    shared_examples_for 'rate sampling' do
+      let(:span_count) { 1000 }
+      let(:rng) { Random.new(123) }
+
+      let(:spans) { Array.new(span_count) { Datadog::Span.new(nil, '', trace_id: rng.rand(Datadog::Span::MAX_ID)) } }
+      let(:expected_num_of_sampled_spans) { span_count * sample_rate }
+
+      it 'samples an appropriate proportion of spans' do
+        spans.each do |span|
+          sampled = sampler.sample(span)
+          expect(span.get_metric(Datadog::RateSampler::SAMPLE_RATE_METRIC_KEY)).to eq(sample_rate) if sampled
+        end
+
+        expect(spans.select(&:sampled).length).to be_within(expected_num_of_sampled_spans * 0.1)
+          .of(expected_num_of_sampled_spans)
+      end
+    end
+
+    it_behaves_like('rate sampling') { let(:sample_rate) { 0.1 } }
+    it_behaves_like('rate sampling') { let(:sample_rate) { 0.25 } }
+    it_behaves_like('rate sampling') { let(:sample_rate) { 0.5 } }
+    it_behaves_like('rate sampling') { let(:sample_rate) { 0.9 } }
+
+    context 'when a sample rate of 1.0 is set' do
+      let(:sample_rate) { 1.0 }
+
+      it 'samples all spans' do
+        spans.each do |span|
+          expect(sampler.sample(span)).to be true
+          expect(span.sampled).to be true
+          expect(span.get_metric(Datadog::RateSampler::SAMPLE_RATE_METRIC_KEY)).to eq(sample_rate)
+        end
+      end
+    end
+  end
+end
+
+RSpec.describe Datadog::PrioritySampler do
+  subject(:sampler) { described_class.new(options) }
+  let(:options) { {} }
+
+  before(:each) { Datadog::Tracer.log.level = Logger::FATAL }
+  after(:each) { Datadog::Tracer.log.level = Logger::WARN }
+
+  describe '#sample' do
+    subject(:sample) { sampler.sample(span) }
+
+    context 'given a span without a context' do
+      let(:span) { Datadog::Span.new(nil, '', trace_id: 1) }
+
+      it do
+        expect(sample).to be true
+        expect(span.sampled).to be(true)
+      end
+    end
+
+    context 'given a span with a context' do
+      let(:span) { Datadog::Span.new(nil, '', trace_id: 1, context: context) }
+      let(:context) { Datadog::Context.new }
+
+      context 'but no sampling priority' do
+        it do
+          expect(sample).to be true
+          expect(span.sampled).to be(true)
+        end
+      end
+
+      context 'and USER_KEEP sampling priority' do
+        before(:each) { context.sampling_priority = Datadog::Ext::Priority::USER_KEEP }
+
+        it do
+          expect(sample).to be true
+          expect(context.sampling_priority).to be(Datadog::Ext::Priority::USER_KEEP)
+          expect(span.sampled).to be(true)
+        end
+      end
+
+      context 'and AUTO_KEEP sampling priority' do
+        before(:each) { context.sampling_priority = Datadog::Ext::Priority::AUTO_KEEP }
+
+        it do
+          expect(sample).to be true
+          expect(context.sampling_priority).to be(Datadog::Ext::Priority::AUTO_KEEP)
+          expect(span.sampled).to be(true)
+        end
+      end
+
+      context 'and AUTO_REJECT sampling priority' do
+        before(:each) { context.sampling_priority = Datadog::Ext::Priority::AUTO_REJECT }
+
+        it do
+          expect(sample).to be true
+          expect(context.sampling_priority).to be(Datadog::Ext::Priority::AUTO_REJECT)
+          expect(span.sampled).to be(true) # Priority sampling always samples
+        end
+      end
+
+      context 'and USER_REJECT sampling priority' do
+        before(:each) { context.sampling_priority = Datadog::Ext::Priority::USER_REJECT }
+
+        it do
+          expect(sample).to be true
+          expect(context.sampling_priority).to be(Datadog::Ext::Priority::USER_REJECT)
+          expect(span.sampled).to be(true) # Priority sampling always samples
+        end
+      end
+    end
+  end
+end