From 65fb81d4cff4ac378a8d60460ee8298a9f1e3178 Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <156273877+p-datadog@users.noreply.github.com>
Date: Thu, 12 Dec 2024 12:07:50 -0500
Subject: [PATCH] DEBUG-3176 DI: report an error when instrumenting a loaded
 file that is not in registry (#4208)

Co-authored-by: Oleg Pudeyev <code@olegp.name>
---
 lib/datadog/di/error.rb              |  5 +++++
 lib/datadog/di/instrumenter.rb       | 24 +++++++++++++++++++++++-
 sig/datadog/di/error.rbs             |  2 ++
 sig/datadog/di/instrumenter.rbs      |  1 +
 spec/datadog/di/instrumenter_spec.rb | 14 ++++++++++++++
 5 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/lib/datadog/di/error.rb b/lib/datadog/di/error.rb
index d5f305617cd..c1785f71ba7 100644
--- a/lib/datadog/di/error.rb
+++ b/lib/datadog/di/error.rb
@@ -27,6 +27,11 @@ class AgentCommunicationError < Error
       class DITargetNotDefined < Error
       end
 
+      # Attempting to instrument a line and the file containing the line
+      # was loaded prior to code tracking being enabled.
+      class DITargetNotInRegistry < Error
+      end
+
       # Raised when trying to install a probe whose installation failed
       # earlier in the same process. This exception should contain the
       # original exception report from initial installation attempt.
diff --git a/lib/datadog/di/instrumenter.rb b/lib/datadog/di/instrumenter.rb
index b9238c9ad74..b1f8f0f599a 100644
--- a/lib/datadog/di/instrumenter.rb
+++ b/lib/datadog/di/instrumenter.rb
@@ -246,11 +246,12 @@ def hook_line(probe, &block)
               #
               # If the requested file is not in code tracker's registry,
               # or the code tracker does not exist at all,
-              # do not attempt to instrumnet now.
+              # do not attempt to instrument now.
               # The caller should add the line to the list of pending lines
               # to instrument and install the hook when the file in
               # question is loaded (and hopefully, by then code tracking
               # is active, otherwise the line will never be instrumented.)
+              raise_if_probe_in_loaded_features(probe)
               raise Error::DITargetNotDefined, "File not in code tracker registry: #{probe.file}"
             end
           end
@@ -258,6 +259,7 @@ def hook_line(probe, &block)
           # Same as previous comment, if untargeted trace points are not
           # explicitly defined, and we do not have code tracking, do not
           # instrument the method.
+          raise_if_probe_in_loaded_features(probe)
           raise Error::DITargetNotDefined, "File not in code tracker registry: #{probe.file}"
         end
 
@@ -374,6 +376,26 @@ def unhook(probe)
 
       attr_reader :lock
 
+      def raise_if_probe_in_loaded_features(probe)
+        return unless probe.file
+
+        # If the probe file is in the list of loaded files
+        # (as per $LOADED_FEATURES, using either exact or suffix match),
+        # raise an error indicating that
+        # code tracker is missing the loaded file because the file
+        # won't be loaded again (DI only works in production environments
+        # that do not normally reload code).
+        if $LOADED_FEATURES.include?(probe.file)
+          raise Error::DITargetNotInRegistry, "File loaded but is not in code tracker registry: #{probe.file}"
+        end
+        # Ths is an expensive check
+        $LOADED_FEATURES.each do |path|
+          if Utils.path_matches_suffix?(path, probe.file)
+            raise Error::DITargetNotInRegistry, "File matching probe path (#{probe.file}) was loaded and is not in code tracker registry: #{path}"
+          end
+        end
+      end
+
       # TODO test that this resolves qualified names e.g. A::B
       def symbolize_class_name(cls_name)
         Object.const_get(cls_name)
diff --git a/sig/datadog/di/error.rbs b/sig/datadog/di/error.rbs
index 351fa7bb85a..4812de578a7 100644
--- a/sig/datadog/di/error.rbs
+++ b/sig/datadog/di/error.rbs
@@ -7,6 +7,8 @@ module Datadog
       end
       class DITargetNotDefined < Error
       end
+      class DITargetNotInRegistry < Error
+      end
       class ProbePreviouslyFailed < Error
       end
       class MultiplePathsMatch < Error
diff --git a/sig/datadog/di/instrumenter.rbs b/sig/datadog/di/instrumenter.rbs
index cfaa56cac64..f87c2818cf4 100644
--- a/sig/datadog/di/instrumenter.rbs
+++ b/sig/datadog/di/instrumenter.rbs
@@ -48,6 +48,7 @@ module Datadog
 
       attr_reader lock: untyped
       def symbolize_class_name: (untyped cls_name) -> untyped
+      def raise_if_probe_in_loaded_features: (Probe probe) -> void
     end
   end
 end
diff --git a/spec/datadog/di/instrumenter_spec.rb b/spec/datadog/di/instrumenter_spec.rb
index 821e6f685df..695d8567386 100644
--- a/spec/datadog/di/instrumenter_spec.rb
+++ b/spec/datadog/di/instrumenter_spec.rb
@@ -983,6 +983,20 @@
         expect(observed_calls.length).to eq 1
         expect(observed_calls.first).to be_a(Hash)
       end
+
+      context 'when instrumenting a line in loaded but not tracked file' do
+        let(:probe) do
+          Datadog::DI::Probe.new(file: 'hook_line.rb', line_no: 3,
+            id: 1, type: :log)
+        end
+
+        it 'raises DITargetNotInRegistry' do
+          expect do
+            instrumenter.hook_line(probe) do |payload|
+            end
+          end.to raise_error(Datadog::DI::Error::DITargetNotInRegistry, /File matching probe path.*was loaded and is not in code tracker registry/)
+        end
+      end
     end
 
     context 'when method is recursive' do