DEBUG-3182 Rework DI loading

lib/datadog/di.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require_relative 'di/base'
 require_relative 'di/error'
 require_relative 'di/code_tracker'
 require_relative 'di/component'
@@ -46,67 +47,7 @@ def enabled?
     # Expose DI to global shared objects
     Extensions.activate!
 
-    LOCK = Mutex.new
-
     class << self
-      attr_reader :code_tracker
-
-      # Activates code tracking. Normally this method should be called
-      # when the application starts. If instrumenting third-party code,
-      # code tracking needs to be enabled before the third-party libraries
-      # are loaded. If you definitely will not be instrumenting
-      # third-party libraries, activating tracking after third-party libraries
-      # have been loaded may improve lookup performance.
-      #
-      # TODO test that activating tracker multiple times preserves
-      # existing mappings in the registry
-      def activate_tracking!
-        (@code_tracker ||= CodeTracker.new).start
-      end
-
-      # Activates code tracking if possible.
-      #
-      # This method does nothing if invoked in an environment that does not
-      # implement required trace points for code tracking (MRI Ruby < 2.6,
-      # JRuby) and rescues any exceptions that may be raised by downstream
-      # DI code.
-      def activate_tracking
-        # :script_compiled trace point was added in Ruby 2.6.
-        return unless RUBY_VERSION >= '2.6'
-
-        begin
-          # Activate code tracking by default because line trace points will not work
-          # without it.
-          Datadog::DI.activate_tracking!
-        rescue => exc
-          if defined?(Datadog.logger)
-            Datadog.logger.warn("Failed to activate code tracking for DI: #{exc.class}: #{exc}")
-          else
-            # We do not have Datadog logger potentially because DI code tracker is
-            # being loaded early in application boot process and the rest of datadog
-            # wasn't loaded yet. Output to standard error.
-            warn("Failed to activate code tracking for DI: #{exc.class}: #{exc}")
-          end
-        end
-      end
-
-      # Deactivates code tracking. In normal usage of DI this method should
-      # never be called, however it is used by DI's test suite to reset
-      # state for individual tests.
-      #
-      # Note that deactivating tracking clears out the registry, losing
-      # the ability to look up files that have been loaded into the process
-      # already.
-      def deactivate_tracking!
-        code_tracker&.stop
-      end
-
-      # Returns whether code tracking is available.
-      # This method should be used instead of querying #code_tracker
-      # because the latter one may be nil.
-      def code_tracking_active?
-        code_tracker&.active? || false
-      end
 
       # This method is called from DI Remote handler to issue DI operations
       # to the probe manager (add or remove probes).
@@ -120,39 +61,6 @@ def code_tracking_active?
       def component
         Datadog.send(:components).dynamic_instrumentation
       end
-
-      # DI code tracker is instantiated globally before the regular set of
-      # components is created, but the code tracker needs to call out to the
-      # "current" DI component to perform instrumentation when application
-      # code is loaded. Because this call may happen prior to Datadog
-      # components having been initialized, we maintain the "current component"
-      # which contains a reference to the most recently instantiated
-      # DI::Component. This way, if a DI component hasn't been instantiated,
-      # we do not try to reference Datadog.components.
-      def current_component
-        LOCK.synchronize do
-          @current_components&.last
-        end
-      end
-
-      # To avoid potential races with DI::Component being added and removed,
-      # we maintain a list of the components. Normally the list should contain
-      # either zero or one component depending on whether DI is enabled in
-      # Datadog configuration. However, if a new instance of DI::Component
-      # is created while the previous instance is still running, we are
-      # guaranteed to not end up with no component when one is running.
-      def add_current_component(component)
-        LOCK.synchronize do
-          @current_components ||= []
-          @current_components << component
-        end
-      end
-
-      def remove_current_component(component)
-        LOCK.synchronize do
-          @current_components&.delete(component)
-        end
-      end
     end
   end
 end

lib/datadog/di/base.rb

@@ -0,0 +1,116 @@
+# frozen_string_literal: true
+
+# This file is loaded by datadog/di/init.rb.
+# It contains just the global DI reference to the (normally one and only)
+# code tracker for the current process.
+# This file should not require the rest of DI, specifically none of the
+# contrib code that is meant to be loaded after third-party libraries
+# are loaded, and also none of the rest of datadog library which also
+# has contrib code in other products.
+
+require_relative 'code_tracker'
+
+module Datadog
+  # Namespace for Datadog dynamic instrumentation.
+  #
+  # @api private
+  module DI
+    LOCK = Mutex.new
+
+    class << self
+      attr_reader :code_tracker
+
+      # Activates code tracking. Normally this method should be called
+      # when the application starts. If instrumenting third-party code,
+      # code tracking needs to be enabled before the third-party libraries
+      # are loaded. If you definitely will not be instrumenting
+      # third-party libraries, activating tracking after third-party libraries
+      # have been loaded may improve lookup performance.
+      #
+      # TODO test that activating tracker multiple times preserves
+      # existing mappings in the registry
+      def activate_tracking!
+        (@code_tracker ||= CodeTracker.new).start
+      end
+
+      # Activates code tracking if possible.
+      #
+      # This method does nothing if invoked in an environment that does not
+      # implement required trace points for code tracking (MRI Ruby < 2.6,
+      # JRuby) and rescues any exceptions that may be raised by downstream
+      # DI code.
+      def activate_tracking
+        # :script_compiled trace point was added in Ruby 2.6.
+        return unless RUBY_VERSION >= '2.6'
+
+        begin
+          # Activate code tracking by default because line trace points will not work
+          # without it.
+          Datadog::DI.activate_tracking!
+        rescue => exc
+          if defined?(Datadog.logger)
+            Datadog.logger.warn("Failed to activate code tracking for DI: #{exc.class}: #{exc}")
+          else
+            # We do not have Datadog logger potentially because DI code tracker is
+            # being loaded early in application boot process and the rest of datadog
+            # wasn't loaded yet. Output to standard error.
+            warn("Failed to activate code tracking for DI: #{exc.class}: #{exc}")
+          end
+        end
+      end
+
+      # Deactivates code tracking. In normal usage of DI this method should
+      # never be called, however it is used by DI's test suite to reset
+      # state for individual tests.
+      #
+      # Note that deactivating tracking clears out the registry, losing
+      # the ability to look up files that have been loaded into the process
+      # already.
+      def deactivate_tracking!
+        code_tracker&.stop
+      end
+
+      # Returns whether code tracking is available.
+      # This method should be used instead of querying #code_tracker
+      # because the latter one may be nil.
+      def code_tracking_active?
+        code_tracker&.active? || false
+      end
+
+      # DI code tracker is instantiated globally before the regular set of
+      # components is created, but the code tracker needs to call out to the
+      # "current" DI component to perform instrumentation when application
+      # code is loaded. Because this call may happen prior to Datadog
+      # components having been initialized, we maintain the "current component"
+      # which contains a reference to the most recently instantiated
+      # DI::Component. This way, if a DI component hasn't been instantiated,
+      # we do not try to reference Datadog.components.
+      # In other words, this method exists so that we never attempt to call
+      # Datadog.components from the code tracker.
+      def current_component
+        LOCK.synchronize do
+          @current_components&.last
+        end
+      end
+
+      # To avoid potential races with DI::Component being added and removed,
+      # we maintain a list of the components. Normally the list should contain
+      # either zero or one component depending on whether DI is enabled in
+      # Datadog configuration. However, if a new instance of DI::Component
+      # is created while the previous instance is still running, we are
+      # guaranteed to not end up with no component when one is running.
+      def add_current_component(component)
+        LOCK.synchronize do
+          @current_components ||= []
+          @current_components << component
+        end
+      end
+
+      def remove_current_component(component)
+        LOCK.synchronize do
+          @current_components&.delete(component)
+        end
+      end
+    end
+  end
+end

lib/datadog/di/code_tracker.rb

@@ -2,6 +2,8 @@
 
 # rubocop:disable Lint/AssignmentInCondition
 
+require_relative 'error'
+
 module Datadog
   module DI
     # Tracks loaded Ruby code by source file and maintains a map from
@@ -87,9 +89,10 @@ def start
           # rescue any exceptions that might not be handled to not break said
           # customer applications.
           rescue => exc
-            # TODO we do not have DI.component defined yet, remove steep:ignore
-            # before release.
-            if component = DI.current_component # steep:ignore
+            # Code tracker may be loaded without the rest of DI,
+            # in which case DI.component will not yet be defined,
+            # but we will have DI.current_component (set to nil).
+            if component = DI.current_component
               raise if component.settings.dynamic_instrumentation.internal.propagate_all_exceptions
               component.logger.warn("Unhandled exception in script_compiled trace point: #{exc.class}: #{exc}")
               component.telemetry&.report(exc, description: "Unhandled exception in script_compiled trace point")

lib/datadog/di/init.rb

@@ -4,7 +4,7 @@
 # enable dynamic instrumentation for third-party libraries used by the
 # application.
 
-require_relative '../di'
+require_relative 'base'
 
 # Code tracking is required for line probes to work; see the comments
 # on the activate_tracking methods in di.rb for further details.