Clear thread local variables when forked

lib/elastic_apm/agent.rb

@@ -288,6 +288,8 @@ def detect_forking!
       instrumenter.handle_forking!
       metrics.handle_forking!
 
+      Spies::MongoSpy::Subscriber.handle_forking!
+
       @pid = Process.pid
     end
   end

lib/elastic_apm/instrumenter.rb

@@ -38,6 +38,11 @@ def initialize
         self.spans = []
       end
 
+      def clear!
+        Thread.current[TRANSACTION_KEY] = nil
+        Thread.current[SPAN_KEY] = nil
+      end
+
       def transaction
         Thread.current[TRANSACTION_KEY]
       end
@@ -84,6 +89,7 @@ def stop
     end
 
     def handle_forking!
+      @current.clear!
       stop
       start
     end

lib/elastic_apm/spies/mongo.rb

@@ -37,6 +37,10 @@ class Subscriber
 
         EVENT_KEY = :__elastic_instrumenter_mongo_events_key
 
+        def self.handle_forking!
+          Thread.current[EVENT_KEY] = []
+        end
+
         def events
           Thread.current[EVENT_KEY] ||= []
         end

spec/elastic_apm/agent_spec.rb

@@ -192,15 +192,25 @@ class AgentTestError < StandardError; end
         subject.stop
       end
 
-      it 'calls handle_forking! on all associated objects' do
-        allow(Process).to receive(:pid).and_return(1)
-
-        expect(subject.central_config).to receive(:handle_forking!)
-        expect(subject.transport).to receive(:handle_forking!)
-        expect(subject.instrumenter).to receive(:handle_forking!)
-        expect(subject.metrics).to receive(:handle_forking!)
-
-        subject.report_message('Everything went boom')
+      it 'handles forking on all associated objects', :intercept do
+        with_agent do
+          ElasticAPM.with_transaction do
+            expect(ElasticAPM.current_transaction).not_to be_nil
+
+            expect(ElasticAPM.agent.central_config).to receive(:handle_forking!)
+            expect(ElasticAPM.agent.transport).to receive(:handle_forking!)
+
+            # The thread local variables cached by the instrumenter need to be clear, so call original
+            expect(ElasticAPM.agent.instrumenter).to receive(:handle_forking!).and_call_original
+            expect(ElasticAPM.agent.metrics).to receive(:handle_forking!)
+            expect(ElasticAPM::Spies::MongoSpy::Subscriber).to receive(:handle_forking!)
+
+            # Simulate a new fork
+            allow(Process).to receive(:pid).and_return(1)
+            ElasticAPM.report_message('Everything went boom')
+            expect(ElasticAPM.current_transaction).to be_nil
+          end
+        end
       end
     end
   end

spec/elastic_apm/spies/mongo_spec.rb

@@ -133,5 +133,34 @@ module ElasticAPM
         expect(@intercepted.spans.length).to be(thread_count)
       end
     end
+
+    context 'forking', :intercept do
+      let(:event) do
+        double('event',
+          command: { 'find' => 'testing',
+                     'filter' => { 'a' => 'bc' } },
+          command_name: 'find',
+          database_name: 'elastic-apm-test',
+          operation_id: 456)
+      end
+      let(:subscriber) { Spies::MongoSpy::Subscriber.new }
+
+      it 'clears the thread local variables' do
+        with_agent do
+          ElasticAPM.with_transaction do
+            subscriber.started(event)
+            expect(Thread.current[Spies::MongoSpy::Subscriber::EVENT_KEY][0]).to be_a(Span)
+
+            # Simulate the following happening in a new fork
+            allow(Process).to receive(:pid).and_return(1)
+            ElasticAPM.agent.detect_forking!
+            expect(Thread.current[Spies::MongoSpy::Subscriber::EVENT_KEY]).to be_empty
+            subscriber.succeeded(event)
+
+            expect(@intercepted.spans.length).to be 0
+          end
+        end
+      end
+    end
   end
 end

spec/support/intercept.rb

@@ -50,6 +50,8 @@ def start; end
 
     def stop; end
 
+    def handle_forking!; end
+
     def validate_span!(span)
       type, subtype = [span.type, span.subtype]