replace monkey patching Kernel with tracepoint

DataDog · Sep 11, 2024 · 49a4a8f · 49a4a8f
1 parent d3bb89d
commit 49a4a8f
Show file tree

Hide file tree

Showing 6 changed files with 51 additions and 32 deletions.
diff --git a/lib/datadog/ci/auto_instrument.rb b/lib/datadog/ci/auto_instrument.rb
@@ -2,7 +2,4 @@
 require "datadog/ci/contrib/contrib"
 require "datadog/ci/contrib/kernel"
 
-::Kernel.prepend(Datadog::CI::Contrib::Kernel)
-::Kernel.singleton_class.prepend(Datadog::CI::Contrib::Kernel)
-
 Datadog::CI::Contrib.auto_instrument!
diff --git a/lib/datadog/ci/contrib/contrib.rb b/lib/datadog/ci/contrib/contrib.rb
@@ -23,12 +23,16 @@ def self.auto_instrument!
             integration.requires.each do |require_path|
               Datadog.logger.debug("Registering on require hook for #{require_path}...")
 
-              ::Kernel.on_require(require_path) do
+              ::Datadog::CI::Contrib.on_require(require_path) do
                 configure_ci_with_framework(name)
               end
+
+              ::Datadog::CI::Contrib.register(require_path, integration)
             end
           end
         end
+
+        Datadog::CI::Contrib.enable_trace_requires
       end
 
       def self.configure_ci_with_framework(framework)

diff --git a/lib/datadog/ci/contrib/cucumber/integration.rb b/lib/datadog/ci/contrib/cucumber/integration.rb
@@ -21,7 +21,7 @@ def self.version
           end
 
           def self.loaded?
-            !defined?(::Cucumber).nil? && !defined?(::Cucumber::Runtime).nil?
+            !defined?(::Cucumber).nil? && !defined?(::Cucumber::Runtime).nil? && !defined?(::Cucumber::Configuration).nil?
           end
 
           def self.compatible?

diff --git a/lib/datadog/ci/contrib/kernel.rb b/lib/datadog/ci/contrib/kernel.rb
@@ -1,42 +1,57 @@
 module Datadog
   module CI
     module Contrib
-      module Kernel
-        def require(name)
-          just_loaded = super
+      class Instance
+        def initialize
+          @on_require = {}
+          @integrations = {}
+        end
 
-          @@dd_instance.require(name) if just_loaded
+        def require(path)
+          Datadog.logger.debug { "Path: #{path}" }
+          @on_require.keys.each do |script_name|
+            if path.include?(script_name) && @integrations[script_name].class.loaded?
+              Datadog.logger.debug { "Gem '#{script_name}' loaded. Configuring integration." }
 
-          just_loaded
+              Contrib.disable_trace_requires
+              @on_require[script_name].call
+            end
+          end
+        rescue => e
+          Datadog.logger.debug do
+            "Failed to execute callback for gem: #{e.class.name} #{e.message} at #{Array(e.backtrace).join("\n")}"
+          end
         end
 
         def on_require(gem, &block)
-          @@dd_instance.on_require(gem, &block)
+          @on_require[gem] = block
         end
 
-        class Instance
-          def initialize
-            @on_require = {}
-          end
+        def register(gem, integration)
+          @integrations[gem] = integration
+        end
+      end
 
-          def require(name)
-            if @on_require.include?(name)
-              Datadog.logger.debug { "Gem '#{name}' loaded. Invoking callback." }
+      @@dd_instance = Instance.new
 
-              @on_require[name].call
-            end
-          rescue => e
-            Datadog.logger.debug do
-              "Failed to execute callback for gem '#{name}': #{e.class.name} #{e.message} at #{Array(e.backtrace).join("\n")}"
-            end
-          end
+      def self.on_require(gem, &block)
+        @@dd_instance.on_require(gem, &block)
+      end
 
-          def on_require(gem, &block)
-            @on_require[gem] = block
-          end
+      def self.register(gem, integration)
+        @@dd_instance.register(gem, integration)
+      end
+
+      def self.enable_trace_requires
+        @@trp = TracePoint.new(:script_compiled) do |tp|
+          @@dd_instance.require(tp.instruction_sequence.path)
         end
 
-        @@dd_instance = Instance.new
+        @@trp.enable
+      end
+
+      def self.disable_trace_requires
+        @@trp.disable
       end
     end
   end

diff --git a/lib/datadog/ci/contrib/minitest/integration.rb b/lib/datadog/ci/contrib/minitest/integration.rb
@@ -21,7 +21,8 @@ def self.version
           end
 
           def self.loaded?
-            !defined?(::Minitest).nil?
+            !defined?(::Minitest).nil? && !defined?(::Minitest::Runnable).nil? && !defined?(::Minitest::Test).nil? &&
+              !defined?(::Minitest::CompositeReporter).nil?
           end
 
           def self.compatible?

diff --git a/lib/datadog/ci/contrib/rspec/integration.rb b/lib/datadog/ci/contrib/rspec/integration.rb
@@ -22,15 +22,17 @@ def self.version
 
           def self.loaded?
             !defined?(::RSpec).nil? && !defined?(::RSpec::Core).nil? &&
-              !defined?(::RSpec::Core::Example).nil?
+              !defined?(::RSpec::Core::Example).nil? &&
+              !defined?(::RSpec::Core::Runner).nil? &&
+              !defined?(::RSpec::Core::ExampleGroup).nil?
           end
 
           def self.compatible?
             super && version >= MINIMUM_VERSION
           end
 
           def requires
-            ["rspec/core"]
+            ["rspec"]
           end
 
           # TODO: rename the following 2 methods: the difference is not about auto or on session start: