Isolate sucker_punch fetch_spans calls.

Author: sarahchen6

spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb

@@ -39,6 +39,7 @@
   end
 
   let(:expect_thread?) { true }
+  let(:mutex) { Mutex.new }
 
   let(:worker_class) do
     Class.new do
@@ -50,6 +51,12 @@ def perform(action = :none, **_)
     end
   end
 
+  def check_span_count_eq(count)
+    mutex.synchronize do
+      try_wait_until { fetch_spans_without_sorting.length == count }
+    end
+  end
+
   context 'successful job' do
     subject(:dummy_worker_success) { worker_class.perform_async }
 
@@ -60,19 +67,19 @@ def perform(action = :none, **_)
     it_behaves_like 'measured span for integration', true do
       before do
         dummy_worker_success
-        try_wait_until { fetch_spans.length == 2 }
+        check_span_count_eq(2)
       end
     end
 
     it 'generates two spans, one for pushing to enqueue and one for the job itself' do
       is_expected.to be true
-      try_wait_until { fetch_spans.length == 2 }
+      check_span_count_eq(2)
       expect(spans.length).to eq(2)
     end
 
     it 'instruments successful job' do
       is_expected.to be true
-      try_wait_until { fetch_spans.length == 2 }
+      check_span_count_eq(2)
 
       expect(job_span.service).to eq(tracer.default_service)
       expect(job_span.name).to eq('sucker_punch.perform')
@@ -85,7 +92,7 @@ def perform(action = :none, **_)
 
     it 'instruments successful enqueuing' do
       is_expected.to be true
-      try_wait_until { fetch_spans.length == 2 }
+      check_span_count_eq(2)
 
       expect(enqueue_span.service).to eq(tracer.default_service)
       expect(enqueue_span.name).to eq('sucker_punch.perform_async')
@@ -106,13 +113,13 @@ def perform(action = :none, **_)
     it_behaves_like 'measured span for integration', true do
       before do
         dummy_worker_fail
-        try_wait_until { fetch_spans.length == 2 }
+        check_span_count_eq(2)
       end
     end
 
     it 'instruments a failed job' do
       is_expected.to be true
-      try_wait_until { fetch_spans.length == 2 }
+      check_span_count_eq(2)
 
       expect(job_span.service).to eq(tracer.default_service)
       expect(job_span.name).to eq('sucker_punch.perform')
@@ -133,13 +140,13 @@ def perform(action = :none, **_)
     it_behaves_like 'measured span for integration', true do
       before do
         dummy_worker_delay
-        try_wait_until { fetch_spans.length == 2 }
+        check_span_count_eq(2)
       end
     end
 
     it 'instruments enqueuing for a delayed job' do
       dummy_worker_delay
-      try_wait_until { fetch_spans.length == 2 }
+      check_span_count_eq(2)
 
       expect(enqueue_span.service).to eq(tracer.default_service)
       expect(enqueue_span.name).to eq('sucker_punch.perform_in')

spec/datadog/tracing/contrib/support/tracer_helpers.rb

@@ -14,7 +14,6 @@ module TracerHelpers
     # Returns the current tracer instance
     def tracer
       Datadog::Tracing.send(:tracer)
-      @write_lock = Mutex.new
     end
 
     # Returns spans and caches it (similar to +let(:spans)+).
@@ -36,33 +35,34 @@ def fetch_traces(tracer = self.tracer)
     # Retrieves and sorts all spans in the current tracer instance.
     # This method does not cache its results.
     def fetch_spans(tracer = self.tracer)
-      lock = tracer.instance_variable_get(:@write_lock)
-      return [] if lock.nil?
-      lock.synchronize do
-        traces = fetch_traces(tracer)
-        traces.collect(&:spans).flatten.sort! do |a, b|
-          if a.name == b.name
-            if a.resource == b.resource
-              if a.start_time == b.start_time
-                a.end_time <=> b.end_time
-              else
-                a.start_time <=> b.start_time
-              end
+      traces = fetch_traces(tracer)
+      traces.collect(&:spans).flatten.sort! do |a, b|
+        if a.name == b.name
+          if a.resource == b.resource
+            if a.start_time == b.start_time
+              a.end_time <=> b.end_time
             else
-              a.resource <=> b.resource
+              a.start_time <=> b.start_time
             end
           else
-            a.name <=> b.name
+            a.resource <=> b.resource
           end
+        else
+          a.name <=> b.name
         end
       end
     end
 
+    def fetch_spans_without_sorting(tracer = self.tracer)
+      traces = fetch_traces(tracer)
+      spans = traces.collect(&:spans)
+      spans.flatten # gets spans for every trace in the tracer instance
+    end
+
     # Remove all traces from the current tracer instance and
     # busts cache of +#spans+ and +#span+.
     def clear_traces!
       tracer.instance_variable_set(:@traces, [])
-      tracer.instance_variable_set(:@write_lock, Mutex.new)
 
       @traces = nil
       @trace = nil
@@ -79,9 +79,10 @@ def clear_traces!
           instance = method.call(**args, &block)
 
           # The mutex must be eagerly initialized to prevent race conditions on lazy initialization
+          write_lock = Mutex.new
           allow(instance).to receive(:write) do |trace|
             instance.instance_exec do
-              tracer.instance_variable_get(:@write_lock).synchronize do
+              write_lock.synchronize do
                 @traces ||= []
                 @traces << trace
               end