From 73496802eae9bbe44d491f9c02820dae021762b3 Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <156273877+p-datadog@users.noreply.github.com>
Date: Tue, 17 Dec 2024 08:42:12 -0500
Subject: [PATCH] DEBUG-3182 DI: remove hard dependency on tracing (#4223)

---
 lib/datadog/di/component.rb                   |  2 +
 lib/datadog/di/init.rb                        |  2 -
 lib/datadog/di/probe_notification_builder.rb  | 18 +++++-
 lib/datadog/di/transport.rb                   |  1 +
 sig/datadog/di/probe_notification_builder.rbs |  3 +
 spec/datadog/di/code_tracker_spec.rb          |  7 +--
 spec/datadog/di/instrumenter_spec.rb          | 57 ++++++++++++-------
 spec/datadog/di/probe_manager_spec.rb         |  2 +
 spec/datadog/di/probe_notifier_worker_spec.rb |  1 +
 spec/datadog/di/remote_spec.rb                |  2 +
 spec/datadog/di/spec_helper.rb                |  8 +++
 spec/datadog/di/transport_spec.rb             |  1 +
 spec/spec_helper.rb                           |  2 +-
 13 files changed, 76 insertions(+), 30 deletions(-)

diff --git a/lib/datadog/di/component.rb b/lib/datadog/di/component.rb
index c109bd4082e..8ef93e626d4 100644
--- a/lib/datadog/di/component.rb
+++ b/lib/datadog/di/component.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require_relative '../core'
+
 module Datadog
   module DI
     # Component for dynamic instrumentation.
diff --git a/lib/datadog/di/init.rb b/lib/datadog/di/init.rb
index a0eb8eb5cf3..0b6a9217c2d 100644
--- a/lib/datadog/di/init.rb
+++ b/lib/datadog/di/init.rb
@@ -4,8 +4,6 @@
 # enable dynamic instrumentation for third-party libraries used by the
 # application.
 
-require_relative '../tracing'
-require_relative '../tracing/contrib'
 require_relative '../di'
 
 # Code tracking is required for line probes to work; see the comments
diff --git a/lib/datadog/di/probe_notification_builder.rb b/lib/datadog/di/probe_notification_builder.rb
index f9be3bea746..f11bbdf8603 100644
--- a/lib/datadog/di/probe_notification_builder.rb
+++ b/lib/datadog/di/probe_notification_builder.rb
@@ -141,8 +141,8 @@ def build_snapshot(probe, rv: nil, snapshot: nil, path: nil,
             version: 2,
           },
           # TODO add tests that the trace/span id is correctly propagated
-          "dd.trace_id": Datadog::Tracing.active_trace&.id&.to_s,
-          "dd.span_id": Datadog::Tracing.active_span&.id&.to_s,
+          "dd.trace_id": active_trace&.id&.to_s,
+          "dd.span_id": active_span&.id&.to_s,
           ddsource: 'dd_debugger',
           message: probe.template && evaluate_template(probe.template,
             duration: duration ? duration * 1000 : 0),
@@ -150,6 +150,8 @@ def build_snapshot(probe, rv: nil, snapshot: nil, path: nil,
         }
       end
 
+      private
+
       def build_status(probe, message:, status:)
         {
           service: settings.service,
@@ -200,6 +202,18 @@ def get_local_variables(trace_point)
           map[name] = value
         end
       end
+
+      def active_trace
+        if defined?(Datadog::Tracing)
+          Datadog::Tracing.active_trace
+        end
+      end
+
+      def active_span
+        if defined?(Datadog::Tracing)
+          Datadog::Tracing.active_span
+        end
+      end
     end
   end
 end
diff --git a/lib/datadog/di/transport.rb b/lib/datadog/di/transport.rb
index a6a5522f645..aa8ebce9b75 100644
--- a/lib/datadog/di/transport.rb
+++ b/lib/datadog/di/transport.rb
@@ -2,6 +2,7 @@
 
 require 'ostruct'
 require_relative 'error'
+require_relative '../core/transport/http/adapters/net'
 
 module Datadog
   module DI
diff --git a/sig/datadog/di/probe_notification_builder.rbs b/sig/datadog/di/probe_notification_builder.rbs
index 3a46f3eef53..45891828166 100644
--- a/sig/datadog/di/probe_notification_builder.rbs
+++ b/sig/datadog/di/probe_notification_builder.rbs
@@ -27,6 +27,9 @@ module Datadog
       def timestamp_now: () -> Integer
 
       def get_local_variables: (TracePoint trace_point) -> Hash[Symbol,untyped]
+      
+      def active_trace: () -> Datadog::Tracing::TraceSegment?
+      def active_span: () -> Datadog::Tracing::SpanOperation?
     end
   end
 end
diff --git a/spec/datadog/di/code_tracker_spec.rb b/spec/datadog/di/code_tracker_spec.rb
index 5b9bf0d7fde..896eef6e6dc 100644
--- a/spec/datadog/di/code_tracker_spec.rb
+++ b/spec/datadog/di/code_tracker_spec.rb
@@ -1,12 +1,9 @@
+require 'datadog/di'
 require "datadog/di/spec_helper"
-require "datadog/di/code_tracker"
 
 RSpec.describe Datadog::DI::CodeTracker do
   di_test
-
-  before(:all) do
-    Datadog::DI.deactivate_tracking!
-  end
+  deactivate_code_tracking
 
   let(:tracker) do
     described_class.new
diff --git a/spec/datadog/di/instrumenter_spec.rb b/spec/datadog/di/instrumenter_spec.rb
index 695d8567386..93cc8aee457 100644
--- a/spec/datadog/di/instrumenter_spec.rb
+++ b/spec/datadog/di/instrumenter_spec.rb
@@ -1,8 +1,14 @@
 require "datadog/di/spec_helper"
-require 'datadog/di'
+require 'datadog/di/instrumenter'
+require 'datadog/di/code_tracker'
+require 'datadog/di/serializer'
+require 'datadog/di/probe'
 require_relative 'hook_line'
 require_relative 'hook_method'
+require 'logger'
 
+# The examples below use a local code tracker when they set line probes,
+# for better test encapsulation and to avoid having to clear/reset global state.
 RSpec.describe Datadog::DI::Instrumenter do
   di_test
 
@@ -55,6 +61,18 @@
     %i[caller_locations duration probe rv serialized_entry_args]
   end
 
+  shared_context 'with code tracking' do
+    let!(:code_tracker) do
+      Datadog::DI::CodeTracker.new.tap do |tracker|
+        tracker.start
+      end
+    end
+
+    after do
+      code_tracker.stop
+    end
+  end
+
   describe '.hook_method' do
     after do
       instrumenter.unhook(probe)
@@ -582,6 +600,12 @@
     end
 
     context 'when hooking two identical but different probes' do
+      include_context 'with code tracking'
+
+      before do
+        load File.join(File.dirname(__FILE__), 'hook_line_recursive.rb')
+      end
+
       let(:probe) do
         Datadog::DI::Probe.new(**base_probe_args.merge(
           type_name: 'HookTestClass', method_name: 'hook_test_method'
@@ -810,12 +834,10 @@
       end
 
       context 'with code tracking' do
-        let(:code_tracker) { Datadog::DI.code_tracker }
+        include_context 'with code tracking'
 
         before do
           expect(di_internal_settings).to receive(:untargeted_trace_points).and_return(false)
-          Datadog::DI.activate_tracking!
-          code_tracker.clear
         end
 
         let(:probe) do
@@ -906,7 +928,6 @@
 
     context 'when hooking same line twice with identical but different probes' do
       before(:all) do
-        Datadog::DI.activate_tracking!
         require_relative 'hook_line_basic'
       end
 
@@ -952,16 +973,14 @@
     end
 
     context 'when code tracking is available' do
-      before do
-        Datadog::DI.activate_tracking!
-        require_relative 'hook_line_targeted'
+      include_context 'with code tracking'
 
+      before do
         path = File.join(File.dirname(__FILE__), 'hook_line_targeted.rb')
-        expect(Datadog::DI.code_tracker.send(:registry)[path]).to be_a(RubyVM::InstructionSequence)
+        load path
+        expect(code_tracker.send(:registry)[path]).to be_a(RubyVM::InstructionSequence)
       end
 
-      let(:code_tracker) { Datadog::DI.code_tracker }
-
       let(:probe) do
         Datadog::DI::Probe.new(file: 'hook_line_targeted.rb', line_no: 3,
           id: 1, type: :log)
@@ -969,7 +988,7 @@
 
       it 'targets the trace point' do
         path = File.join(File.dirname(__FILE__), 'hook_line_targeted.rb')
-        target = Datadog::DI.code_tracker.send(:registry)[path]
+        target = code_tracker.send(:registry)[path]
         expect(target).to be_a(RubyVM::InstructionSequence)
 
         expect_any_instance_of(TracePoint).to receive(:enable).with(target: target, target_line: 3).and_call_original
@@ -1000,13 +1019,12 @@
     end
 
     context 'when method is recursive' do
-      before(:all) do
-        Datadog::DI.activate_tracking!
+      include_context 'with code tracking'
+
+      before do
         load File.join(File.dirname(__FILE__), 'hook_line_recursive.rb')
       end
 
-      let(:code_tracker) { Datadog::DI.code_tracker }
-
       context 'non-enriched probe' do
         let(:probe_args) do
           {file: 'hook_line_recursive.rb', line_no: 3}
@@ -1039,13 +1057,12 @@
     end
 
     context 'when method is infinitely recursive' do
-      before(:all) do
-        Datadog::DI.activate_tracking!
+      include_context 'with code tracking'
+
+      before do
         require_relative 'hook_line_recursive'
       end
 
-      let(:code_tracker) { Datadog::DI.code_tracker }
-
       # We need to use a rate limiter, otherwise the stack is exhausted
       # very slowly and this test burns 100% CPU for a long time performing
       # snapshot building etc.
diff --git a/spec/datadog/di/probe_manager_spec.rb b/spec/datadog/di/probe_manager_spec.rb
index f32cb71f577..a6f7a1c8d8a 100644
--- a/spec/datadog/di/probe_manager_spec.rb
+++ b/spec/datadog/di/probe_manager_spec.rb
@@ -1,5 +1,7 @@
 require "datadog/di/spec_helper"
 require 'datadog/di/probe_manager'
+require 'datadog/di/instrumenter'
+require 'logger'
 
 class ProbeManagerSpecTestClass; end
 
diff --git a/spec/datadog/di/probe_notifier_worker_spec.rb b/spec/datadog/di/probe_notifier_worker_spec.rb
index bbb5b2effee..f4ad4bdd30f 100644
--- a/spec/datadog/di/probe_notifier_worker_spec.rb
+++ b/spec/datadog/di/probe_notifier_worker_spec.rb
@@ -1,5 +1,6 @@
 require "datadog/di/spec_helper"
 require "datadog/di/probe_notifier_worker"
+require 'logger'
 
 RSpec.describe Datadog::DI::ProbeNotifierWorker do
   di_test
diff --git a/spec/datadog/di/remote_spec.rb b/spec/datadog/di/remote_spec.rb
index f3a7d386b9e..06368623d3a 100644
--- a/spec/datadog/di/remote_spec.rb
+++ b/spec/datadog/di/remote_spec.rb
@@ -1,5 +1,7 @@
 require "datadog/di/spec_helper"
+require 'datadog/di'
 require 'spec_helper'
+require 'logger'
 
 RSpec.describe Datadog::DI::Remote do
   di_test
diff --git a/spec/datadog/di/spec_helper.rb b/spec/datadog/di/spec_helper.rb
index 0e0258ddae9..09d8ffbd3c2 100644
--- a/spec/datadog/di/spec_helper.rb
+++ b/spec/datadog/di/spec_helper.rb
@@ -1,5 +1,13 @@
 module DIHelpers
   module ClassMethods
+    def deactivate_code_tracking
+      before(:all) do
+        if Datadog::DI.respond_to?(:deactivate_tracking!)
+          Datadog::DI.deactivate_tracking!
+        end
+      end
+    end
+
     def ruby_2_only
       if RUBY_VERSION >= '3'
         before(:all) do
diff --git a/spec/datadog/di/transport_spec.rb b/spec/datadog/di/transport_spec.rb
index 302cc3ecbe3..87b4af67f15 100644
--- a/spec/datadog/di/transport_spec.rb
+++ b/spec/datadog/di/transport_spec.rb
@@ -1,5 +1,6 @@
 require "datadog/di/spec_helper"
 require "datadog/di/transport"
+require 'datadog/core/configuration/agent_settings_resolver'
 
 RSpec.describe Datadog::DI::Transport do
   di_test
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 671a0b002c1..db9b84cd140 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -306,4 +306,4 @@ def initialize(*args, &block)
 # Code tracking calls out to the current DI component, which may reference
 # mock objects in the test suite. Disable it and tests that need code tracking
 # will enable it back for themselves.
-Datadog::DI.deactivate_tracking! if Datadog::DI.respond_to?(:deactivate_tracking!)
+Datadog::DI.deactivate_tracking! if defined?(Datadog::DI) && Datadog::DI.respond_to?(:deactivate_tracking!)