From a9ce5845019c7d221571193acd588aa3d153b8a6 Mon Sep 17 00:00:00 2001
From: Andrey <andrey.marchenko@datadoghq.com>
Date: Thu, 21 Nov 2024 14:42:20 +0100
Subject: [PATCH] working auto-instrumentation

---
 lib/datadog/ci/contrib/instrumentation.rb  | 57 ++++++++++++++++------
 sig/datadog/ci/contrib/instrumentation.rbs |  7 +++
 2 files changed, 50 insertions(+), 14 deletions(-)

diff --git a/lib/datadog/ci/contrib/instrumentation.rb b/lib/datadog/ci/contrib/instrumentation.rb
index 9811ac8e..e212aca7 100644
--- a/lib/datadog/ci/contrib/instrumentation.rb
+++ b/lib/datadog/ci/contrib/instrumentation.rb
@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "datadog/core/utils/only_once"
+
 module Datadog
   module CI
     module Contrib
@@ -24,18 +26,7 @@ def self.register_integration(integration_class)
         def self.auto_instrument
           Datadog.logger.debug("Auto instrumenting all integrations...")
 
-          auto_instrumented_integrations = []
-          @registry.each do |name, integration|
-            # ignore integrations that are not in the Gemfile or have incompatible versions
-            next unless integration.compatible?
-
-            # late instrumented integrations will be patched when the test session starts
-            next if integration.late_instrument?
-
-            Datadog.logger.debug("#{name} should be auto instrumented")
-            auto_instrumented_integrations << integration
-          end
-
+          auto_instrumented_integrations = fetch_auto_instrumented_integrations
           if auto_instrumented_integrations.empty?
             Datadog.logger.warn(
               "Auto instrumentation was requested, but no available integrations were found. " \
@@ -44,15 +35,26 @@ def self.auto_instrument
             return
           end
 
-          script_compiled_tracepoint = TracePoint.new(:script_compiled) do
+          script_compiled_tracepoint = TracePoint.new(:script_compiled) do |tp|
+            all_patched = true
+
             auto_instrumented_integrations.each do |integration|
               next if integration.patched?
+
+              all_patched = false
               next unless integration.loaded?
 
-              Datadog.logger.debug("#{integration.class} is loaded")
+              auto_configure_datadog
 
+              Datadog.logger.debug("#{integration.class} is loaded")
               patch_integration(integration)
             end
+
+            if all_patched
+              Datadog.logger.debug("All expected integrations are patched, disabling the script_compiled tracepoint")
+
+              tp.disable
+            end
           end
           script_compiled_tracepoint.enable
         end
@@ -126,6 +128,33 @@ def self.patch_integration(integration, with_dependencies: false)
             Datadog.logger.debug("Attention: #{integration.class} is not patched (#{patch_results})")
           end
         end
+
+        def self.fetch_auto_instrumented_integrations
+          @registry.filter_map do |name, integration|
+            # ignore integrations that are not in the Gemfile or have incompatible versions
+            next unless integration.compatible?
+
+            # late instrumented integrations will be patched when the test session starts
+            next if integration.late_instrument?
+
+            Datadog.logger.debug("#{name} should be auto instrumented")
+            integration
+          end
+        end
+
+        def self.auto_configure_datadog
+          configure_once.run do
+            Datadog.logger.debug("Applying Datadog configuration in CI mode...")
+            Datadog.configure do |c|
+              c.ci.enabled = true
+              c.tracing.enabled = true
+            end
+          end
+        end
+
+        def self.configure_once
+          @configure_once ||= Datadog::Core::Utils::OnlyOnce.new
+        end
       end
     end
   end
diff --git a/sig/datadog/ci/contrib/instrumentation.rbs b/sig/datadog/ci/contrib/instrumentation.rbs
index 8d168107..8d820d3e 100644
--- a/sig/datadog/ci/contrib/instrumentation.rbs
+++ b/sig/datadog/ci/contrib/instrumentation.rbs
@@ -7,6 +7,8 @@ module Datadog
 
         self.@registry: Hash[Symbol, untyped]
 
+        self.@configure_once: Datadog::Core::Utils::OnlyOnce
+
         def self.registry: () -> Hash[Symbol, untyped]
 
         def self.auto_instrument: () -> void
@@ -23,6 +25,11 @@ module Datadog
 
         def self.patch_integration: (Contrib::Integration integration, ?with_dependencies: bool) -> void
 
+        def self.fetch_auto_instrumented_integrations: () -> Array[Contrib::Integration]
+
+        def self.auto_configure_datadog: () -> void
+
+        def self.configure_once: () -> Datadog::Core::Utils::OnlyOnce
       end
     end
   end