From daae013e54b6dcae26956feae4c128f43f94de1f Mon Sep 17 00:00:00 2001
From: David Elner <david@davidelner.com>
Date: Thu, 29 Oct 2020 14:03:02 -0400
Subject: [PATCH] Changed: Lazily clone trace context after a fork.

---
 lib/ddtrace/context_provider.rb       | 13 ++++++++++-
 spec/ddtrace/context_provider_spec.rb | 33 ++++++++++++++++++++++++---
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/lib/ddtrace/context_provider.rb b/lib/ddtrace/context_provider.rb
index 8bdd2554e0e..c466294b3a4 100644
--- a/lib/ddtrace/context_provider.rb
+++ b/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
 
diff --git a/spec/ddtrace/context_provider_spec.rb b/spec/ddtrace/context_provider_spec.rb
index a8b1f182e03..0248cb0be39 100644
--- a/spec/ddtrace/context_provider_spec.rb
+++ b/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