Reset trace context after fork

lib/ddtrace/context.rb

@@ -3,6 +3,7 @@
 
 require 'ddtrace/context_flush'
 require 'ddtrace/context_provider'
+require 'ddtrace/utils/forking'
 
 module Datadog
   # \Context is used to keep track of a hierarchy of spans for the current
@@ -19,6 +20,8 @@ module Datadog
   # This data structure is thread-safe.
   # rubocop:disable Metrics/ClassLength
   class Context
+    include Datadog::Utils::Forking
+
     # 100k spans is about a 100Mb footprint
     DEFAULT_MAX_LENGTH = 100_000
 
@@ -232,6 +235,21 @@ def to_s
       end
     end
 
+    # Generates equivalent context for forked processes.
+    #
+    # When Context from parent process is forked, child process
+    # should have a Context belonging to the same trace but not
+    # have the parent process spans.
+    def fork_clone
+      self.class.new(
+        trace_id: trace_id,
+        span_id: span_id,
+        sampled: sampled?,
+        sampling_priority: sampling_priority,
+        origin: origin
+      )
+    end
+
     private
 
     def reset(options = {})

lib/ddtrace/context_provider.rb

@@ -15,7 +15,18 @@ def context=(ctx)
 
     # Return the local context.
     def context(key = nil)
-      key.nil? ? @context.local : @context.local(key)
+      current_context = key.nil? ? @context.local : @context.local(key)
+
+      # Rebuild/reset context after a fork
+      #
+      # We don't want forked processes to copy and retransmit spans
+      # that were generated from the parent process. Reset it such
+      # that it acts like a distributed trace.
+      current_context.after_fork! do
+        current_context = self.context = current_context.fork_clone
+      end
+
+      current_context
     end
   end
 

spec/ddtrace/context_provider_spec.rb

@@ -5,6 +5,7 @@
 RSpec.describe Datadog::DefaultContextProvider do
   let(:provider) { described_class.new }
   let(:local_context) { instance_double(Datadog::ThreadLocalContext) }
+  let(:trace_context) { Datadog::Context.new }
 
   context '#context=' do
     subject(:context=) { provider.context = ctx }
@@ -25,7 +26,10 @@
 
     context 'when given no arguments' do
       it do
-        expect(local_context).to receive(:local)
+        expect(local_context)
+          .to receive(:local)
+          .and_return(trace_context)
+
         subject
       end
     end
@@ -38,22 +42,45 @@
         expect(local_context)
           .to receive(:local)
           .with(key)
+          .and_return(trace_context)
 
         subject
       end
     end
   end
 
+  context 'when fork occurs' do
+    before { skip 'Java not supported' if RUBY_PLATFORM == 'java' }
+
+    it 'clones the context and returns the clone' do
+      # Initialize a context for the current process
+      parent_context = provider.context
+      expect(parent_context.forked?).to be false
+
+      # Fork the process, clone context.
+      expect_in_fork do
+        expect(parent_context).to receive(:fork_clone).and_call_original
+        child_context = provider.context
+
+        # Check context changed
+        expect(child_context).to_not be parent_context
+
+        # Check context doesn't change again
+        expect(provider.context).to be(child_context)
+      end
+    end
+  end
+
   context 'with multiple instances' do
     it 'holds independent values for each instance' do
       provider1 = described_class.new
       provider2 = described_class.new
 
-      ctx1 = provider1.context = double
+      ctx1 = provider1.context = Datadog::Context.new
       expect(provider1.context).to be(ctx1)
       expect(provider2.context).to_not be(ctx1)
 
-      ctx2 = provider2.context = double
+      ctx2 = provider2.context = Datadog::Context.new
       expect(provider1.context).to be(ctx1)
       expect(provider2.context).to be(ctx2)
     end

spec/ddtrace/context_spec.rb

@@ -408,6 +408,31 @@ def new_unfinished_span(name = nil)
     end
   end
 
+  describe '#fork_clone' do
+    subject(:fork_clone) { context.fork_clone }
+
+    let(:options) do
+      {
+        trace_id: SecureRandom.uuid,
+        span_id: SecureRandom.uuid,
+        sampled: true,
+        sampling_priority: Datadog::Ext::Priority::AUTO_KEEP,
+        origin: 'synthetics'
+      }
+    end
+
+    it do
+      is_expected.to be_a_kind_of(described_class)
+      is_expected.to have_attributes(
+        trace_id: context.trace_id,
+        span_id: context.span_id,
+        sampled?: context.sampled?,
+        sampling_priority: context.sampling_priority,
+        origin: context.origin
+      )
+    end
+  end
+
   describe '#length' do
     subject(:ctx) { context }
     let(:span) { new_span }