From d603a46c7b965577798d3639b97f8359400db469 Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Tue, 4 Feb 2025 17:37:38 -0500
Subject: [PATCH 1/9] Run flaky test over and over and over...

---
 run_test.rb | 4 ++++
 1 file changed, 4 insertions(+)
 create mode 100644 run_test.rb

diff --git a/run_test.rb b/run_test.rb
new file mode 100644
index 00000000000..e095ccd5622
--- /dev/null
+++ b/run_test.rb
@@ -0,0 +1,4 @@
+1000.times do |i|
+  tests_passed = system('BUNDLE_GEMFILE=/app/gemfiles/jruby_9.2_contrib.gemfile bundle exec rspec ./spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb')
+  break unless tests_passed
+end

From 3a9cc9b1f9fc9fbb08e4c469509faea511ad09f9 Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Tue, 4 Feb 2025 17:38:39 -0500
Subject: [PATCH 2/9] Add mutex to tracer_helpers.

---
 .../tracing/contrib/support/tracer_helpers.rb | 29 +++++++++++--------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/spec/datadog/tracing/contrib/support/tracer_helpers.rb b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
index 9cb6316e294..7d2568fa280 100644
--- a/spec/datadog/tracing/contrib/support/tracer_helpers.rb
+++ b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
@@ -14,6 +14,7 @@ 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)+).
@@ -35,20 +36,24 @@ 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)
-      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
+      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
             else
-              a.start_time <=> b.start_time
+              a.resource <=> b.resource
             end
           else
-            a.resource <=> b.resource
+            a.name <=> b.name
           end
-        else
-          a.name <=> b.name
         end
       end
     end
@@ -57,6 +62,7 @@ def fetch_spans(tracer = self.tracer)
     # 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
@@ -73,10 +79,9 @@ 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
-              write_lock.synchronize do
+              tracer.instance_variable_get(:@write_lock).synchronize do
                 @traces ||= []
                 @traces << trace
               end

From f9cfe16a78d3e237bc8f8eafdcc49bcadeaf7ef5 Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Wed, 5 Feb 2025 17:01:14 -0500
Subject: [PATCH 3/9] Isolate sucker_punch fetch_spans calls.

---
 .../contrib/sucker_punch/patcher_spec.rb      | 23 +++++++-----
 .../tracing/contrib/support/tracer_helpers.rb | 35 ++++++++++---------
 2 files changed, 33 insertions(+), 25 deletions(-)

diff --git a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
index 62bbc90c00c..57427252863 100644
--- a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
+++ b/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')
diff --git a/spec/datadog/tracing/contrib/support/tracer_helpers.rb b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
index 7d2568fa280..738ae6cc65d 100644
--- a/spec/datadog/tracing/contrib/support/tracer_helpers.rb
+++ b/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

From e485c127ca24d0987ca148928eb5b974f8258ebf Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Thu, 6 Feb 2025 14:05:15 -0500
Subject: [PATCH 4/9] Separate out spans' call to fetch_spans.

---
 .../contrib/sucker_punch/patcher_spec.rb      |  2 +-
 .../tracing/contrib/support/tracer_helpers.rb | 45 ++++++++++++-------
 2 files changed, 29 insertions(+), 18 deletions(-)

diff --git a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
index 57427252863..2adc4d33f0d 100644
--- a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
+++ b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
@@ -53,7 +53,7 @@ def perform(action = :none, **_)
 
   def check_span_count_eq(count)
     mutex.synchronize do
-      try_wait_until { fetch_spans_without_sorting.length == count }
+      try_wait_until { fetch_spans.length == count }
     end
   end
 
diff --git a/spec/datadog/tracing/contrib/support/tracer_helpers.rb b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
index 738ae6cc65d..bd693f4e945 100644
--- a/spec/datadog/tracing/contrib/support/tracer_helpers.rb
+++ b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
@@ -11,6 +11,10 @@ module Contrib
   # For contrib, we only allow one tracer to be active:
   # the global tracer in +Datadog::Tracing+.
   module TracerHelpers
+    def mutex
+      @mutex ||= Mutex.new
+    end
+
     # Returns the current tracer instance
     def tracer
       Datadog::Tracing.send(:tracer)
@@ -23,7 +27,9 @@ def traces
 
     # Returns spans and caches it (similar to +let(:spans)+).
     def spans
-      @spans ||= fetch_spans
+      mutex.synchronize do
+        @spans ||= fetch_spans_without_sorting
+      end
     end
 
     # Retrieves all traces in the current tracer instance.
@@ -35,39 +41,44 @@ 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)
-      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
+      mutex.synchronize do
+        traces = fetch_traces(tracer)
+        spans = traces.collect(&:spans)
+        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
             else
-              a.start_time <=> b.start_time
+              a.resource <=> b.resource
             end
           else
-            a.resource <=> b.resource
+            a.name <=> b.name
           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 = traces.map { |trace| trace.instance_variable_get(:@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, [])
+      mutex.synchronize do
+        tracer.instance_variable_set(:@traces, [])
 
-      @traces = nil
-      @trace = nil
-      @spans = nil
-      @span = nil
+        @traces = nil
+        @trace = nil
+        @spans = nil
+        @span = nil
+      end
     end
 
     RSpec.configure do |config|

From 5b972d5eeb9fe73d1b436e502cd17a6827210378 Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Thu, 6 Feb 2025 15:33:17 -0500
Subject: [PATCH 5/9] Create separate method to get spans_count.

---
 run_test.rb                                   |  4 --
 .../contrib/sucker_punch/patcher_spec.rb      | 20 +++----
 .../tracing/contrib/support/tracer_helpers.rb | 52 ++++++++-----------
 3 files changed, 32 insertions(+), 44 deletions(-)
 delete mode 100644 run_test.rb

diff --git a/run_test.rb b/run_test.rb
deleted file mode 100644
index e095ccd5622..00000000000
--- a/run_test.rb
+++ /dev/null
@@ -1,4 +0,0 @@
-1000.times do |i|
-  tests_passed = system('BUNDLE_GEMFILE=/app/gemfiles/jruby_9.2_contrib.gemfile bundle exec rspec ./spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb')
-  break unless tests_passed
-end
diff --git a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
index 2adc4d33f0d..3955ead8982 100644
--- a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
+++ b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
@@ -51,9 +51,9 @@ def perform(action = :none, **_)
     end
   end
 
-  def check_span_count_eq(count)
+  def check_spans_count_eq(count)
     mutex.synchronize do
-      try_wait_until { fetch_spans.length == count }
+      try_wait_until { get_spans_count == count }
     end
   end
 
@@ -67,19 +67,19 @@ def check_span_count_eq(count)
     it_behaves_like 'measured span for integration', true do
       before do
         dummy_worker_success
-        check_span_count_eq(2)
+        check_spans_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
-      check_span_count_eq(2)
+      check_spans_count_eq(2)
       expect(spans.length).to eq(2)
     end
 
     it 'instruments successful job' do
       is_expected.to be true
-      check_span_count_eq(2)
+      check_spans_count_eq(2)
 
       expect(job_span.service).to eq(tracer.default_service)
       expect(job_span.name).to eq('sucker_punch.perform')
@@ -92,7 +92,7 @@ def check_span_count_eq(count)
 
     it 'instruments successful enqueuing' do
       is_expected.to be true
-      check_span_count_eq(2)
+      check_spans_count_eq(2)
 
       expect(enqueue_span.service).to eq(tracer.default_service)
       expect(enqueue_span.name).to eq('sucker_punch.perform_async')
@@ -113,13 +113,13 @@ def check_span_count_eq(count)
     it_behaves_like 'measured span for integration', true do
       before do
         dummy_worker_fail
-        check_span_count_eq(2)
+        check_spans_count_eq(2)
       end
     end
 
     it 'instruments a failed job' do
       is_expected.to be true
-      check_span_count_eq(2)
+      check_spans_count_eq(2)
 
       expect(job_span.service).to eq(tracer.default_service)
       expect(job_span.name).to eq('sucker_punch.perform')
@@ -140,13 +140,13 @@ def check_span_count_eq(count)
     it_behaves_like 'measured span for integration', true do
       before do
         dummy_worker_delay
-        check_span_count_eq(2)
+        check_spans_count_eq(2)
       end
     end
 
     it 'instruments enqueuing for a delayed job' do
       dummy_worker_delay
-      check_span_count_eq(2)
+      check_spans_count_eq(2)
 
       expect(enqueue_span.service).to eq(tracer.default_service)
       expect(enqueue_span.name).to eq('sucker_punch.perform_in')
diff --git a/spec/datadog/tracing/contrib/support/tracer_helpers.rb b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
index bd693f4e945..ef059b84ab8 100644
--- a/spec/datadog/tracing/contrib/support/tracer_helpers.rb
+++ b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
@@ -11,10 +11,6 @@ module Contrib
   # For contrib, we only allow one tracer to be active:
   # the global tracer in +Datadog::Tracing+.
   module TracerHelpers
-    def mutex
-      @mutex ||= Mutex.new
-    end
-
     # Returns the current tracer instance
     def tracer
       Datadog::Tracing.send(:tracer)
@@ -27,9 +23,7 @@ def traces
 
     # Returns spans and caches it (similar to +let(:spans)+).
     def spans
-      mutex.synchronize do
-        @spans ||= fetch_spans_without_sorting
-      end
+      @spans ||= fetch_spans
     end
 
     # Retrieves all traces in the current tracer instance.
@@ -41,44 +35,42 @@ 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)
-      mutex.synchronize do
-        traces = fetch_traces(tracer)
-        spans = traces.collect(&:spans)
-        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)
+    def get_spans_count(tracer = self.tracer)
+      count = 0
       traces = fetch_traces(tracer)
-      spans = traces.map { |trace| trace.instance_variable_get(:@spans) || [] }
-      spans.flatten # gets spans for every trace in the tracer instance
+      traces.each do |trace|
+        count += trace.instance_variable_get(:@spans).length
+      end
+      count
     end
 
     # Remove all traces from the current tracer instance and
     # busts cache of +#spans+ and +#span+.
     def clear_traces!
-      mutex.synchronize do
-        tracer.instance_variable_set(:@traces, [])
+      tracer.instance_variable_set(:@traces, [])
 
-        @traces = nil
-        @trace = nil
-        @spans = nil
-        @span = nil
-      end
+      @traces = nil
+      @trace = nil
+      @spans = nil
+      @span = nil
     end
 
     RSpec.configure do |config|

From cd45c0647b2afa8cd34e05e7c5b7a265d7c9dbe1 Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Thu, 6 Feb 2025 16:18:45 -0500
Subject: [PATCH 6/9] Remove unnecessary sucker_punch mutexes.

---
 .../contrib/sucker_punch/patcher_spec.rb      | 23 +++++++------------
 1 file changed, 8 insertions(+), 15 deletions(-)

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

From e02ad490498ae4ab6e736210280388c45f60f6a3 Mon Sep 17 00:00:00 2001
From: Sarah Chen <sarah.chen@datadoghq.com>
Date: Fri, 7 Feb 2025 11:05:36 -0500
Subject: [PATCH 7/9] Keep everything to sucker_punch test file.

---
 .../datadog/tracing/contrib/sucker_punch/patcher_spec.rb | 4 ++++
 spec/datadog/tracing/contrib/support/tracer_helpers.rb   | 9 ---------
 2 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
index bcf18bb1b57..0b8298db83a 100644
--- a/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
+++ b/spec/datadog/tracing/contrib/sucker_punch/patcher_spec.rb
@@ -50,6 +50,10 @@ def perform(action = :none, **_)
     end
   end
 
+  def get_spans_count(tracer = self.tracer)
+    fetch_traces(tracer).sum { |trace| trace.spans.length }
+  end
+
   context 'successful job' do
     subject(:dummy_worker_success) { worker_class.perform_async }
 
diff --git a/spec/datadog/tracing/contrib/support/tracer_helpers.rb b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
index ef059b84ab8..9cb6316e294 100644
--- a/spec/datadog/tracing/contrib/support/tracer_helpers.rb
+++ b/spec/datadog/tracing/contrib/support/tracer_helpers.rb
@@ -53,15 +53,6 @@ def fetch_spans(tracer = self.tracer)
       end
     end
 
-    def get_spans_count(tracer = self.tracer)
-      count = 0
-      traces = fetch_traces(tracer)
-      traces.each do |trace|
-        count += trace.instance_variable_get(:@spans).length
-      end
-      count
-    end
-
     # Remove all traces from the current tracer instance and
     # busts cache of +#spans+ and +#span+.
     def clear_traces!

From 971040f015704e50dcd1798b173e2a8ab962b053 Mon Sep 17 00:00:00 2001
From: Marco Costa <marco.costa@datadoghq.com>
Date: Mon, 10 Feb 2025 15:37:50 -0500
Subject: [PATCH 8/9] Serialize span events via a dedicated field (#4279)

* Increase type checking coverage

* Serialize span events via a dedicated field

* Revert unrelated test changes

* Skip test running against incorrect combinations

* Use kwarg
---
 .github/forced-tests-list.json                |   9 +-
 Steepfile                                     |   2 -
 lib/datadog/core/configuration/components.rb  |   8 +-
 lib/datadog/core/encoding.rb                  |  16 +++
 lib/datadog/core/environment/agent_info.rb    |  77 +++++++++++++
 .../core/remote/transport/http/negotiation.rb |   1 +
 .../core/remote/transport/negotiation.rb      |  14 ++-
 .../tracing/transport/serializable_trace.rb   |  12 ++-
 lib/datadog/tracing/transport/traces.rb       |  33 ++++--
 sig/datadog/core/configuration/components.rbs |   2 +
 sig/datadog/core/encoding.rbs                 |  30 +++---
 sig/datadog/core/environment/agent_info.rbs   |  13 +++
 .../core/remote/transport/negotiation.rbs     |  12 ++-
 sig/datadog/core/transport/response.rbs       |   2 +-
 .../tracing/transport/serializable_trace.rbs  |   6 +-
 sig/datadog/tracing/transport/traces.rbs      |  10 +-
 .../core/configuration/components_spec.rb     |   2 +
 spec/datadog/tracing/integration_spec.rb      |  85 +++++++++++++++
 .../transport/serializable_trace_spec.rb      |   2 +-
 spec/datadog/tracing/transport/traces_spec.rb | 101 ++++++++++++++++--
 spec/support/spy_transport.rb                 |   8 +-
 21 files changed, 394 insertions(+), 51 deletions(-)
 create mode 100644 lib/datadog/core/environment/agent_info.rb
 create mode 100644 sig/datadog/core/environment/agent_info.rbs

diff --git a/.github/forced-tests-list.json b/.github/forced-tests-list.json
index 077404aaa41..fa3cc348b11 100644
--- a/.github/forced-tests-list.json
+++ b/.github/forced-tests-list.json
@@ -1,3 +1,10 @@
 {
-  
+  "AGENT_NOT_SUPPORTING_SPAN_EVENTS":
+  [
+    "tests/test_span_events.py"
+  ],
+  "PARAMETRIC":
+  [
+    "tests/parametric/test_span_events.py"
+  ]
 }
\ No newline at end of file
diff --git a/Steepfile b/Steepfile
index 7acf3d2057b..26bf7f7fc24 100644
--- a/Steepfile
+++ b/Steepfile
@@ -137,10 +137,8 @@ target :datadog do
   ignore 'lib/datadog/tracing/transport/http/traces.rb'
   ignore 'lib/datadog/tracing/transport/io/client.rb'
   ignore 'lib/datadog/tracing/transport/io/traces.rb'
-  ignore 'lib/datadog/tracing/transport/serializable_trace.rb'
   ignore 'lib/datadog/tracing/transport/statistics.rb'
   ignore 'lib/datadog/tracing/transport/trace_formatter.rb'
-  ignore 'lib/datadog/tracing/transport/traces.rb'
   ignore 'lib/datadog/tracing/workers.rb'
   ignore 'lib/datadog/tracing/workers/trace_writer.rb'
   ignore 'lib/datadog/tracing/writer.rb'
diff --git a/lib/datadog/core/configuration/components.rb b/lib/datadog/core/configuration/components.rb
index 25afac8bff2..1bf96a53f71 100644
--- a/lib/datadog/core/configuration/components.rb
+++ b/lib/datadog/core/configuration/components.rb
@@ -16,6 +16,8 @@
 require_relative '../../di/component'
 require_relative '../crashtracking/component'
 
+require_relative '../environment/agent_info'
+
 module Datadog
   module Core
     module Configuration
@@ -85,7 +87,8 @@ def build_crashtracker(settings, agent_settings, logger:)
           :tracer,
           :crashtracker,
           :dynamic_instrumentation,
-          :appsec
+          :appsec,
+          :agent_info
 
         def initialize(settings)
           @logger = self.class.build_logger(settings)
@@ -96,6 +99,9 @@ def initialize(settings)
           # the Core resolver from within your product/component's namespace.
           agent_settings = AgentSettingsResolver.call(settings, logger: @logger)
 
+          # Exposes agent capability information for detection by any components
+          @agent_info = Core::Environment::AgentInfo.new(agent_settings)
+
           @telemetry = self.class.build_telemetry(settings, agent_settings, @logger)
 
           @remote = Remote::Component.build(settings, agent_settings, telemetry: telemetry)
diff --git a/lib/datadog/core/encoding.rb b/lib/datadog/core/encoding.rb
index 951d221d4e2..8061d5a898b 100644
--- a/lib/datadog/core/encoding.rb
+++ b/lib/datadog/core/encoding.rb
@@ -10,6 +10,7 @@ module Encoding
       # Encoder interface that provides the logic to encode traces and service
       # @abstract
       module Encoder
+        # :nocov:
         def content_type
           raise NotImplementedError
         end
@@ -23,6 +24,13 @@ def join(encoded_elements)
         def encode(_)
           raise NotImplementedError
         end
+
+        # Deserializes a value serialized with {#encode}.
+        # This method is used for debugging purposes.
+        def decode(_)
+          raise NotImplementedError
+        end
+        # :nocov:
       end
 
       # Encoder for the JSON format
@@ -41,6 +49,10 @@ def encode(obj)
           JSON.dump(obj)
         end
 
+        def decode(obj)
+          JSON.parse(obj)
+        end
+
         def join(encoded_data)
           "[#{encoded_data.join(',')}]"
         end
@@ -62,6 +74,10 @@ def encode(obj)
           MessagePack.pack(obj)
         end
 
+        def decode(obj)
+          MessagePack.unpack(obj)
+        end
+
         def join(encoded_data)
           packer = MessagePack::Packer.new
           packer.write_array_header(encoded_data.size)
diff --git a/lib/datadog/core/environment/agent_info.rb b/lib/datadog/core/environment/agent_info.rb
new file mode 100644
index 00000000000..249b6af9b74
--- /dev/null
+++ b/lib/datadog/core/environment/agent_info.rb
@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Core
+    module Environment
+      # Retrieves the agent's `/info` endpoint data.
+      # This data can be used to determine the capabilities of the local Datadog agent.
+      #
+      # @example Example response payload
+      #   {
+      #     "version" : "7.57.2",
+      #     "git_commit" : "38ba0c7",
+      #     "endpoints" : [ "/v0.4/traces", "/v0.4/services", "/v0.7/traces", "/v0.7/config" ],
+      #     "client_drop_p0s" : true,
+      #     "span_meta_structs" : true,
+      #     "long_running_spans" : true,
+      #     "evp_proxy_allowed_headers" : [ "Content-Type", "Accept-Encoding", "Content-Encoding", "User-Agent" ],
+      #     "config" : {
+      #       "default_env" : "none",
+      #       "target_tps" : 10,
+      #       "max_eps" : 200,
+      #       "receiver_port" : 8126,
+      #       "receiver_socket" : "/var/run/datadog/apm.socket",
+      #       "connection_limit" : 0,
+      #       "receiver_timeout" : 0,
+      #       "max_request_bytes" : 26214400,
+      #       "statsd_port" : 8125,
+      #       "analyzed_spans_by_service" : { },
+      #       "obfuscation" : {
+      #         "elastic_search" : true,
+      #         "mongo" : true,
+      #         "sql_exec_plan" : false,
+      #         "sql_exec_plan_normalize" : false,
+      #         "http" : {
+      #           "remove_query_string" : false,
+      #           "remove_path_digits" : false
+      #         },
+      #         "remove_stack_traces" : false,
+      #         "redis" : {
+      #           "Enabled" : true,
+      #           "RemoveAllArgs" : false
+      #         },
+      #         "memcached" : {
+      #           "Enabled" : true,
+      #           "KeepCommand" : false
+      #         }
+      #       }
+      #     },
+      #     "peer_tags" : null
+      #   }
+      #
+      # @see https://github.com/DataDog/datadog-agent/blob/f07df0a3c1fca0c83b5a15f553bd994091b0c8ac/pkg/trace/api/info.go#L20
+      class AgentInfo
+        attr_reader :agent_settings
+
+        def initialize(agent_settings)
+          @agent_settings = agent_settings
+          @client = Remote::Transport::HTTP.root(agent_settings: agent_settings)
+        end
+
+        # Fetches the information from the agent.
+        # @return [Datadog::Core::Remote::Transport::HTTP::Negotiation::Response] the response from the agent
+        # @return [nil] if an error occurred while fetching the information
+        def fetch
+          res = @client.send_info
+          return unless res.ok?
+
+          res
+        end
+
+        def ==(other)
+          other.is_a?(self.class) && other.agent_settings == agent_settings
+        end
+      end
+    end
+  end
+end
diff --git a/lib/datadog/core/remote/transport/http/negotiation.rb b/lib/datadog/core/remote/transport/http/negotiation.rb
index 5cc00c8a8f0..4a939222ef6 100644
--- a/lib/datadog/core/remote/transport/http/negotiation.rb
+++ b/lib/datadog/core/remote/transport/http/negotiation.rb
@@ -43,6 +43,7 @@ def initialize(http_response, options = {})
                 @version = options[:version]
                 @endpoints = options[:endpoints]
                 @config = options[:config]
+                @span_events = options[:span_events]
               end
             end
 
diff --git a/lib/datadog/core/remote/transport/negotiation.rb b/lib/datadog/core/remote/transport/negotiation.rb
index e9209087092..b6e7f6c8833 100644
--- a/lib/datadog/core/remote/transport/negotiation.rb
+++ b/lib/datadog/core/remote/transport/negotiation.rb
@@ -32,7 +32,19 @@ class Request < Datadog::Core::Transport::Request
 
           # Negotiation response
           module Response
-            attr_reader :version, :endpoints, :config
+            # @!attribute [r] version
+            #   The version of the agent.
+            #   @return [String]
+            # @!attribute [r] endpoints
+            #   The HTTP endpoints the agent supports.
+            #   @return [Array<String>]
+            # @!attribute [r] config
+            #   The agent configuration. These are configured by the user when starting the agent, as well as any defaults.
+            #   @return [Hash]
+            # @!attribute [r] span_events
+            #   Whether the agent supports the top-level span events field in flushed spans.
+            #   @return [Boolean,nil]
+            attr_reader :version, :endpoints, :config, :span_events
           end
 
           # Negotiation transport
diff --git a/lib/datadog/tracing/transport/serializable_trace.rb b/lib/datadog/tracing/transport/serializable_trace.rb
index 380f8623954..b3786903312 100644
--- a/lib/datadog/tracing/transport/serializable_trace.rb
+++ b/lib/datadog/tracing/transport/serializable_trace.rb
@@ -14,7 +14,7 @@ class SerializableTrace
 
         # @param trace [Datadog::Trace] the trace to serialize
         # @param native_events_supported [Boolean] whether the agent supports span events as a top-level field
-        def initialize(trace, native_events_supported = false)
+        def initialize(trace, native_events_supported:)
           @trace = trace
           @native_events_supported = native_events_supported
         end
@@ -29,13 +29,17 @@ def initialize(trace, native_events_supported = false)
         # @param packer [MessagePack::Packer] serialization buffer, can be +nil+ with JRuby
         def to_msgpack(packer = nil)
           # As of 1.3.3, JRuby implementation doesn't pass an existing packer
-          trace.spans.map { |s| SerializableSpan.new(s, @native_events_supported) }.to_msgpack(packer)
+          trace.spans.map do |s|
+            SerializableSpan.new(s, native_events_supported: @native_events_supported)
+          end.to_msgpack(packer)
         end
 
         # JSON serializer interface.
         # Used by older version of the transport.
         def to_json(*args)
-          trace.spans.map { |s| SerializableSpan.new(s, @native_events_supported).to_hash }.to_json(*args)
+          trace.spans.map do |s|
+            SerializableSpan.new(s, native_events_supported: @native_events_supported).to_hash
+          end.to_json(*args)
         end
       end
 
@@ -46,7 +50,7 @@ class SerializableSpan
 
         # @param span [Datadog::Span] the span to serialize
         # @param native_events_supported [Boolean] whether the agent supports span events as a top-level field
-        def initialize(span, native_events_supported)
+        def initialize(span, native_events_supported:)
           @span = span
           @trace_id = Tracing::Utils::TraceId.to_low_order(span.trace_id)
           @native_events_supported = native_events_supported
diff --git a/lib/datadog/tracing/transport/traces.rb b/lib/datadog/tracing/transport/traces.rb
index 644f80fe94c..441477ac7ca 100644
--- a/lib/datadog/tracing/transport/traces.rb
+++ b/lib/datadog/tracing/transport/traces.rb
@@ -50,8 +50,9 @@ class Chunker
           #
           # @param encoder [Datadog::Core::Encoding::Encoder]
           # @param max_size [String] maximum acceptable payload size
-          def initialize(encoder, max_size: DEFAULT_MAX_PAYLOAD_SIZE)
+          def initialize(encoder, native_events_supported:, max_size: DEFAULT_MAX_PAYLOAD_SIZE)
             @encoder = encoder
+            @native_events_supported = native_events_supported
             @max_size = max_size
           end
 
@@ -77,7 +78,7 @@ def encode_in_chunks(traces)
           private
 
           def encode_one(trace)
-            encoded = Encoder.encode_trace(encoder, trace)
+            encoded = Encoder.encode_trace(encoder, trace, native_events_supported: @native_events_supported)
 
             if encoded.size > max_size
               # This single trace is too large, we can't flush it
@@ -95,17 +96,18 @@ def encode_one(trace)
         module Encoder
           module_function
 
-          def encode_trace(encoder, trace)
+          def encode_trace(encoder, trace, native_events_supported:)
             # Format the trace for transport
             TraceFormatter.format!(trace)
 
             # Make the trace serializable
-            serializable_trace = SerializableTrace.new(trace)
-
-            Datadog.logger.debug { "Flushing trace: #{JSON.dump(serializable_trace)}" }
+            serializable_trace = SerializableTrace.new(trace, native_events_supported: native_events_supported)
 
             # Encode the trace
-            encoder.encode(serializable_trace)
+            encoder.encode(serializable_trace).tap do |encoded|
+              # Print the actual serialized trace, since the encoder can change make non-trivial changes
+              Datadog.logger.debug { "Flushing trace: #{encoder.decode(encoded)}" }
+            end
           end
         end
 
@@ -126,7 +128,10 @@ def initialize(apis, default_api)
 
           def send_traces(traces)
             encoder = current_api.encoder
-            chunker = Datadog::Tracing::Transport::Traces::Chunker.new(encoder)
+            chunker = Datadog::Tracing::Transport::Traces::Chunker.new(
+              encoder,
+              native_events_supported: native_events_supported?
+            )
 
             responses = chunker.encode_in_chunks(traces.lazy).map do |encoded_traces, trace_count|
               request = Request.new(EncodedParcel.new(encoded_traces, trace_count))
@@ -188,6 +193,18 @@ def change_api!(api_id)
             @client = HTTP::Client.new(current_api)
           end
 
+          # Queries the agent for native span events serialization support.
+          # This changes how the serialization of span events performed.
+          def native_events_supported?
+            return @native_events_supported if defined?(@native_events_supported)
+
+            if (res = Datadog.send(:components).agent_info.fetch)
+              @native_events_supported = res.span_events == true
+            else
+              false
+            end
+          end
+
           # Raised when configured with an unknown API version
           class UnknownApiVersionError < StandardError
             attr_reader :version
diff --git a/sig/datadog/core/configuration/components.rbs b/sig/datadog/core/configuration/components.rbs
index b4e2ad4aa92..b7d96f7b728 100644
--- a/sig/datadog/core/configuration/components.rbs
+++ b/sig/datadog/core/configuration/components.rbs
@@ -36,6 +36,8 @@ module Datadog
 
         attr_reader remote: Datadog::Core::Remote::Component
 
+        attr_reader agent_info: Datadog::Core::Environment::AgentInfo
+
         def initialize: (untyped settings) -> untyped
 
         def startup!: (untyped settings) -> untyped
diff --git a/sig/datadog/core/encoding.rbs b/sig/datadog/core/encoding.rbs
index f3fee9b1aed..9320f0fb615 100644
--- a/sig/datadog/core/encoding.rbs
+++ b/sig/datadog/core/encoding.rbs
@@ -5,39 +5,43 @@ module Datadog
       # Encoder interface that provides the logic to encode traces and service
       # @abstract
       module Encoder
-        def content_type: () -> untyped
+        def content_type: () -> String
 
-        # Concatenates a list of elements previously encoded by +#encode+.
-        def join: (untyped encoded_elements) -> untyped
+        def encode: (untyped obj) -> String
 
-        # Serializes a single trace into a String suitable for network transmission.
-        def encode: (untyped _) -> untyped
+        def join: (Array[untyped] encoded_data) -> String
+
+        def decode: (String obj)-> untyped
       end
 
       # Encoder for the JSON format
       module JSONEncoder
         extend Encoder
 
-        CONTENT_TYPE: "application/json"
+        CONTENT_TYPE: String
+
+        def self?.content_type: () -> String
 
-        def self?.content_type: () -> untyped
+        def self?.encode: (untyped obj) -> String
 
-        def self?.encode: (untyped obj) -> untyped
+        def self?.join: (Array[untyped] encoded_data) -> String
 
-        def self?.join: (untyped encoded_data) -> ::String
+        def self?.decode: (String obj)-> untyped
       end
 
       # Encoder for the Msgpack format
       module MsgpackEncoder
         extend Encoder
 
-        CONTENT_TYPE: "application/msgpack"
+        CONTENT_TYPE: String
+
+        def self?.content_type: () -> String
 
-        def self?.content_type: () -> untyped
+        def self?.encode: (untyped obj) -> String
 
-        def self?.encode: (untyped obj) -> untyped
+        def self?.join: (Array[untyped] encoded_data) -> String
 
-        def self?.join: (untyped encoded_data) -> untyped
+        def self?.decode: (String obj)-> untyped
       end
     end
   end
diff --git a/sig/datadog/core/environment/agent_info.rbs b/sig/datadog/core/environment/agent_info.rbs
new file mode 100644
index 00000000000..ddf013809b9
--- /dev/null
+++ b/sig/datadog/core/environment/agent_info.rbs
@@ -0,0 +1,13 @@
+module Datadog
+  module Core
+    module Environment
+      class AgentInfo
+        attr_reader agent_settings: Configuration::AgentSettingsResolver::AgentSettings
+
+        def initialize: (Configuration::AgentSettingsResolver::AgentSettings agent_settings) -> void
+
+        def fetch: -> Remote::Transport::HTTP::Negotiation::Response?
+      end
+    end
+  end
+end
diff --git a/sig/datadog/core/remote/transport/negotiation.rbs b/sig/datadog/core/remote/transport/negotiation.rbs
index 0f90aa77b57..82d21ebe05e 100644
--- a/sig/datadog/core/remote/transport/negotiation.rbs
+++ b/sig/datadog/core/remote/transport/negotiation.rbs
@@ -7,11 +7,13 @@ module Datadog
           end
 
           module Response
-            attr_reader version: untyped
+            attr_reader version: String
 
-            attr_reader endpoints: untyped
+            attr_reader endpoints: Array[String]
 
-            attr_reader config: untyped
+            attr_reader config: Hash[String,untyped]
+
+            attr_reader span_events: bool
           end
 
           class Transport
@@ -25,7 +27,9 @@ module Datadog
 
             def initialize: (untyped apis, untyped default_api) -> void
 
-            def send_info: () -> untyped
+            type send_info_return = HTTP::Negotiation::Response & Core::Transport::InternalErrorResponse
+
+            def send_info: () -> send_info_return
 
             def current_api: () -> untyped
           end
diff --git a/sig/datadog/core/transport/response.rbs b/sig/datadog/core/transport/response.rbs
index 888e67568e1..11558cfe48b 100644
--- a/sig/datadog/core/transport/response.rbs
+++ b/sig/datadog/core/transport/response.rbs
@@ -22,7 +22,7 @@ module Datadog
       class InternalErrorResponse
         include Response
 
-        attr_reader error: untyped
+        attr_reader error: Exception
 
         def initialize: (untyped error) -> void
 
diff --git a/sig/datadog/tracing/transport/serializable_trace.rbs b/sig/datadog/tracing/transport/serializable_trace.rbs
index ab84fcfc23d..de27a49c2cd 100644
--- a/sig/datadog/tracing/transport/serializable_trace.rbs
+++ b/sig/datadog/tracing/transport/serializable_trace.rbs
@@ -4,9 +4,9 @@ module Datadog
       class SerializableTrace
         @native_events_supported: bool
 
-        attr_reader trace: Span
+        attr_reader trace: TraceSegment
 
-        def initialize: (untyped trace, bool native_events_supported) -> void
+        def initialize: (untyped trace, native_events_supported: bool) -> void
 
         def to_msgpack: (?untyped? packer) -> untyped
 
@@ -19,7 +19,7 @@ module Datadog
 
         attr_reader span: Span
 
-        def initialize: (untyped span, bool native_events_supported) -> void
+        def initialize: (untyped span, native_events_supported: bool) -> void
 
         def to_msgpack: (?untyped? packer) -> untyped
 
diff --git a/sig/datadog/tracing/transport/traces.rbs b/sig/datadog/tracing/transport/traces.rbs
index 00e3eabc06d..25f22332c17 100644
--- a/sig/datadog/tracing/transport/traces.rbs
+++ b/sig/datadog/tracing/transport/traces.rbs
@@ -28,7 +28,7 @@ module Datadog
 
           attr_reader max_size: untyped
 
-          def initialize: (untyped encoder, ?max_size: untyped) -> void
+          def initialize: (untyped encoder, native_events_supported: bool, ?max_size: untyped) -> void
 
           def encode_in_chunks: (untyped traces) -> untyped
 
@@ -38,10 +38,12 @@ module Datadog
         end
 
         module Encoder
-          def self?.encode_trace: (untyped encoder, untyped trace) -> untyped
+          def self?.encode_trace: (untyped encoder, untyped trace, native_events_supported: bool) -> untyped
         end
 
         class Transport
+          @native_events_supported: bool
+
           attr_reader client: untyped
 
           attr_reader apis: untyped
@@ -52,7 +54,7 @@ module Datadog
 
           def initialize: (untyped apis, untyped default_api) -> void
 
-          def send_traces: (untyped traces) -> untyped
+          def send_traces: (Array[Tracing::TraceOperation] traces) -> untyped
 
           def stats: () -> untyped
 
@@ -81,6 +83,8 @@ module Datadog
 
             def message: () -> ::String
           end
+
+          def native_events_supported?: -> bool
         end
       end
     end
diff --git a/spec/datadog/core/configuration/components_spec.rb b/spec/datadog/core/configuration/components_spec.rb
index 8073c653517..efda72cc6da 100644
--- a/spec/datadog/core/configuration/components_spec.rb
+++ b/spec/datadog/core/configuration/components_spec.rb
@@ -31,6 +31,7 @@
   let(:logger) { instance_double(Datadog::Core::Logger) }
   let(:settings) { Datadog::Core::Configuration::Settings.new }
   let(:agent_settings) { Datadog::Core::Configuration::AgentSettingsResolver.call(settings, logger: nil) }
+  let(:agent_info) { Datadog::Core::Environment::AgentInfo.new(agent_settings) }
 
   let(:profiler_setup_task) { Datadog::Profiling.supported? ? instance_double(Datadog::Profiling::Tasks::Setup) : nil }
   let(:remote) { instance_double(Datadog::Core::Remote::Component, start: nil, shutdown!: nil) }
@@ -95,6 +96,7 @@
       expect(components.profiler).to be profiler
       expect(components.runtime_metrics).to be runtime_metrics
       expect(components.health_metrics).to be health_metrics
+      expect(components.agent_info).to eq agent_info
     end
 
     describe '@environment_logger_extra' do
diff --git a/spec/datadog/tracing/integration_spec.rb b/spec/datadog/tracing/integration_spec.rb
index 6824a123612..649a2f17505 100644
--- a/spec/datadog/tracing/integration_spec.rb
+++ b/spec/datadog/tracing/integration_spec.rb
@@ -867,6 +867,83 @@ def agent_receives_span_step3(previous_success)
     end
   end
 
+  shared_examples 'flushes traces with span events' do |native_span_events_support: true|
+    context 'a trace with span events' do
+      subject(:trace_with_event) do
+        tracer.trace('parent_span') do |span|
+          span.span_events << Datadog::Tracing::SpanEvent.new(
+            'event_name',
+            time_unix_nano: 123,
+            attributes: { 'key' => 'value' }
+          )
+        end
+
+        try_wait_until(seconds: 2) { tracer.writer.stats[:traces_flushed] >= 1 }
+      end
+
+      before do
+        allow_any_instance_of(Datadog::Core::Remote::Transport::HTTP::Negotiation::Response)
+          .to receive(:span_events).and_return(span_events_support)
+
+        allow(encoder).to receive(:encode).and_wrap_original do |m, *args|
+          encoded = m.call(*args)
+          traces = encoder.decode(encoded)
+
+          traces = traces['traces'].first if traces.is_a?(Hash) # For Transport::IO
+
+          expect(traces).to have(1).item
+          @flushed_trace = traces.first
+
+          encoded
+        end
+      end
+
+      let(:flushed_trace) { @flushed_trace }
+
+      context 'with agent supporting native span events' do
+        before do
+          skip 'Environment does not support native span events' unless native_span_events_support
+        end
+
+        let(:span_events_support) { true }
+
+        it 'flushes events using the span_events field' do
+          trace_with_event
+
+          expect(flushed_trace['meta']).to_not have_key('events')
+          expect(flushed_trace['span_events']).to eq(
+            [
+              { 'name' => 'event_name',
+                'time_unix_nano' => 123,
+                'attributes' => { 'key' => {
+                  'string_value' => 'value', 'type' => 0
+                } }, }
+            ]
+          )
+        end
+      end
+
+      context 'with agent not supporting native span events' do
+        let(:span_events_support) { false }
+
+        it 'flushes events using the span_events field' do
+          trace_with_event
+
+          expect(flushed_trace['meta']['events']).to eq(
+            JSON.dump(
+              [
+                { 'name' => 'event_name',
+                  'time_unix_nano' => 123,
+                  'attributes' => { 'key' => 'value' } }
+              ]
+            )
+          )
+          expect(flushed_trace).to_not have_key('span_events')
+        end
+      end
+    end
+  end
+
   describe 'Transport::IO' do
     include_context 'agent-based test'
 
@@ -921,6 +998,10 @@ def agent_receives_span_step3(previous_success)
 
       expect(out).to have_received(:puts)
     end
+
+    it_behaves_like 'flushes traces with span events', native_span_events_support: false do
+      let(:encoder) { Datadog.send(:components).tracer.writer.transport.encoder }
+    end
   end
 
   describe 'Transport::HTTP' do
@@ -965,6 +1046,10 @@ def agent_receives_span_step3(previous_success)
         expect(stats[:transport].internal_error).to eq(0)
       end
     end
+
+    it_behaves_like 'flushes traces with span events' do
+      let(:encoder) { Datadog.send(:components).tracer.writer.transport.current_api.encoder }
+    end
   end
 
   describe 'fiber-local context' do
diff --git a/spec/datadog/tracing/transport/serializable_trace_spec.rb b/spec/datadog/tracing/transport/serializable_trace_spec.rb
index de88a519225..afef82a68ef 100644
--- a/spec/datadog/tracing/transport/serializable_trace_spec.rb
+++ b/spec/datadog/tracing/transport/serializable_trace_spec.rb
@@ -8,7 +8,7 @@
 require 'datadog/tracing/transport/serializable_trace'
 
 RSpec.describe Datadog::Tracing::Transport::SerializableTrace do
-  subject(:serializable_trace) { described_class.new(trace, native_events_supported) }
+  subject(:serializable_trace) { described_class.new(trace, native_events_supported: native_events_supported) }
 
   let(:trace) { Datadog::Tracing::TraceSegment.new(spans) }
   let(:native_events_supported) { false }
diff --git a/spec/datadog/tracing/transport/traces_spec.rb b/spec/datadog/tracing/transport/traces_spec.rb
index ab78c7ac8b0..5931bc17326 100644
--- a/spec/datadog/tracing/transport/traces_spec.rb
+++ b/spec/datadog/tracing/transport/traces_spec.rb
@@ -60,8 +60,9 @@
 end
 
 RSpec.describe Datadog::Tracing::Transport::Traces::Chunker do
-  let(:chunker) { described_class.new(encoder, max_size: max_size) }
+  let(:chunker) { described_class.new(encoder, native_events_supported: native_events_supported, max_size: max_size) }
   let(:encoder) { instance_double(Datadog::Core::Encoding::Encoder) }
+  let(:native_events_supported) { double }
   let(:trace_encoder) { Datadog::Tracing::Transport::Traces::Encoder }
   let(:max_size) { 10 }
 
@@ -72,9 +73,21 @@
       let(:traces) { get_test_traces(3) }
 
       before do
-        allow(trace_encoder).to receive(:encode_trace).with(encoder, traces[0]).and_return('1')
-        allow(trace_encoder).to receive(:encode_trace).with(encoder, traces[1]).and_return('22')
-        allow(trace_encoder).to receive(:encode_trace).with(encoder, traces[2]).and_return('333')
+        allow(trace_encoder).to receive(:encode_trace).with(
+          encoder,
+          traces[0],
+          native_events_supported: native_events_supported
+        ).and_return('1')
+        allow(trace_encoder).to receive(:encode_trace).with(
+          encoder,
+          traces[1],
+          native_events_supported: native_events_supported
+        ).and_return('22')
+        allow(trace_encoder).to receive(:encode_trace).with(
+          encoder,
+          traces[2],
+          native_events_supported: native_events_supported
+        ).and_return('333')
         allow(encoder).to receive(:join) { |arr| arr.join(',') }
       end
 
@@ -138,8 +151,8 @@
 
     let(:api_v1) { instance_double(Datadog::Tracing::Transport::HTTP::API::Instance, 'v1', encoder: encoder_v1) }
     let(:api_v2) { instance_double(Datadog::Tracing::Transport::HTTP::API::Instance, 'v2', encoder: encoder_v2) }
-    let(:encoder_v1) { instance_double(Datadog::Core::Encoding::Encoder, content_type: 'text/plain') }
-    let(:encoder_v2) { instance_double(Datadog::Core::Encoding::Encoder, content_type: 'text/csv') }
+    let(:encoder_v1) { instance_double(Datadog::Core::Encoding::Encoder, 'v1', content_type: 'text/plain') }
+    let(:encoder_v2) { instance_double(Datadog::Core::Encoding::Encoder, 'v2', content_type: 'text/csv') }
   end
 
   describe '#initialize' do
@@ -157,6 +170,7 @@
     subject(:send_traces) { transport.send_traces(traces) }
 
     let(:traces) { [] }
+
     let(:response) { Class.new { include Datadog::Core::Transport::Response }.new }
     let(:responses) { [response] }
 
@@ -170,10 +184,20 @@
     let(:client_v1) { instance_double(Datadog::Tracing::Transport::HTTP::Client) }
 
     let(:chunker) { instance_double(Datadog::Tracing::Transport::Traces::Chunker, max_size: 1) }
+    let(:native_events_supported) { nil }
+    let(:agent_info_response) do
+      instance_double(Datadog::Core::Remote::Transport::HTTP::Negotiation::Response, span_events: native_events_supported)
+    end
 
     before do
-      allow(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(encoder_v1).and_return(chunker)
-      allow(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(encoder_v2).and_return(chunker)
+      allow(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(
+        encoder_v1,
+        native_events_supported: false
+      ).and_return(chunker)
+      allow(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(
+        encoder_v2,
+        native_events_supported: false
+      ).and_return(chunker)
 
       allow(chunker).to receive(:encode_in_chunks).and_return(lazy_chunks)
 
@@ -183,6 +207,9 @@
       allow(client_v2).to receive(:send_traces_payload).with(request).and_return(response)
 
       allow(Datadog::Tracing::Transport::Traces::Request).to receive(:new).and_return(request)
+
+      allow_any_instance_of(Datadog::Core::Environment::AgentInfo).to receive(:fetch)
+        .and_return(agent_info_response)
     end
 
     context 'which returns an OK response' do
@@ -254,6 +281,64 @@
         expect(health_metrics).to have_received(:transport_chunked).with(1)
       end
     end
+
+    context 'for native span event support by the agent' do
+      context 'on a successful agent info call' do
+        context 'with support not advertised' do
+          let(:native_events_supported) { nil }
+
+          it 'does not encode native span events' do
+            expect(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(
+              encoder_v2,
+              native_events_supported: false
+            ).and_return(chunker)
+            send_traces
+          end
+        end
+
+        context 'with support advertised as supported' do
+          let(:native_events_supported) { true }
+
+          it 'encodes native span events' do
+            expect(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(
+              encoder_v2,
+              native_events_supported: true
+            ).and_return(chunker)
+            send_traces
+          end
+        end
+
+        context 'with support advertised as unsupported' do
+          let(:native_events_supported) { false }
+
+          it 'encodes native span events' do
+            expect(Datadog::Tracing::Transport::Traces::Chunker).to receive(:new).with(
+              encoder_v2,
+              native_events_supported: false
+            ).and_return(chunker)
+            send_traces
+          end
+        end
+
+        it 'caches the agent result' do
+          transport.send_traces(traces)
+          transport.send_traces(traces)
+
+          expect(Datadog.send(:components).agent_info).to have_received(:fetch).once
+        end
+      end
+
+      context 'on an unsuccessful agent info call' do
+        let(:agent_info_response) { nil }
+
+        it 'does not cache the agent result' do
+          transport.send_traces(traces)
+          transport.send_traces(traces)
+
+          expect(Datadog.send(:components).agent_info).to have_received(:fetch).twice
+        end
+      end
+    end
   end
 
   describe '#downgrade?' do
diff --git a/spec/support/spy_transport.rb b/spec/support/spy_transport.rb
index f36439e408f..45f28226a7a 100644
--- a/spec/support/spy_transport.rb
+++ b/spec/support/spy_transport.rb
@@ -43,7 +43,13 @@ def initialize(*)
 
   def send_traces(data)
     encoded_data = data.map do |trace|
-      @helper_encoder.join([Datadog::Tracing::Transport::Traces::Encoder.encode_trace(@helper_encoder, trace)])
+      @helper_encoder.join(
+        [Datadog::Tracing::Transport::Traces::Encoder.encode_trace(
+          @helper_encoder,
+          trace,
+          native_events_supported: true
+        )]
+      )
     end
 
     @helper_mutex.synchronize do

From 5319afcdba36777acf428670d65023bccee23f6b Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <156273877+p-datadog@users.noreply.github.com>
Date: Mon, 10 Feb 2025 15:51:19 -0500
Subject: [PATCH 9/9] DEBUG-3472 send snapshot and status events together
 (#4360)

Co-authored-by: Oleg Pudeyev <code@olegp.name>
---
 lib/datadog/di/probe_notifier_worker.rb | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/datadog/di/probe_notifier_worker.rb b/lib/datadog/di/probe_notifier_worker.rb
index 03efaece9e0..0e497c870ec 100644
--- a/lib/datadog/di/probe_notifier_worker.rb
+++ b/lib/datadog/di/probe_notifier_worker.rb
@@ -263,7 +263,7 @@ def set_sleep_remaining
 
       def maybe_send
         rv = maybe_send_status
-        rv || maybe_send_snapshot
+        maybe_send_snapshot || rv
       end
     end
   end