DEBUG-2334 repair instrumentation of virtual and later-defined methods

lib/datadog/di/instrumenter.rb

@@ -92,8 +92,16 @@ def hook_method(probe, &block)
         cls = symbolize_class_name(probe.type_name)
         serializer = self.serializer
         method_name = probe.method_name
-        target_method = cls.instance_method(method_name)
-        loc = target_method.source_location
+        loc = begin
+          cls.instance_method(method_name).source_location
+        rescue NameError
+          # The target method is not defined.
+          # This could be because it will be explicitly defined later
+          # (since classes can be reopened in Ruby)
+          # or the method is virtual (provided by a method_missing handler).
+          # In these cases we do not have a source location for the
+          # target method here.
+        end
         rate_limiter = probe.rate_limiter
 
         mod = Module.new do
@@ -111,8 +119,19 @@ def hook_method(probe, &block)
               # The method itself is not part of the stack trace because
               # we are getting the stack trace from outside of the method.
               # Add the method in manually as the top frame.
-              method_frame = Location.new(loc.first, loc.last, method_name)
-              caller_locs = [method_frame] + caller_locations # steep:ignore
+              method_frame = if loc
+                [Location.new(loc.first, loc.last, method_name)]
+              else
+                # For virtual and lazily-defined methods, we do not have
+                # the original source location here, and they won't be
+                # included in the stack trace currently.
+                # TODO when begin/end trace points are added for local
+                # variable capture in method probes, we should be able
+                # to obtain actual method execution location and use
+                # that location here.
+                []
+              end
+              caller_locs = method_frame + caller_locations # steep:ignore
               # TODO capture arguments at exit
               # & is to stop steep complaints, block is always present here.
               block&.call(probe: probe, rv: rv, duration: duration, caller_locations: caller_locs,

spec/datadog/di/integration/instrumentation_spec.rb

@@ -131,7 +131,7 @@ def run_test
           end
         end
 
-        context 'when method is defined after probe is added to probe manager' do
+        context 'when class with target method is defined after probe is added to probe manager' do
           let(:probe) do
             Datadog::DI::Probe.new(id: "1234", type: :log,
               type_name: 'InstrumentationDelayedTestClass', method_name: 'test_method',
@@ -170,6 +170,94 @@ def test_method
             )
           end
         end
+
+        context 'when class exists without target method and method is defined after probe is added to probe manager' do
+          let(:probe) do
+            Datadog::DI::Probe.new(id: "1234", type: :log,
+              type_name: 'InstrumentationDelayedPartialTestClass', method_name: 'test_method',
+              capture_snapshot: false,)
+          end
+
+          it 'invokes probe and creates expected snapshot' do
+            class InstrumentationDelayedPartialTestClass # rubocop:disable Lint/ConstantDefinitionInBlock
+              # test_method should not be defined here
+            end
+
+            expect(component.transport).to receive(:send_request).at_least(:once)
+            expect(probe_manager.add_probe(probe)).to be true
+
+            class InstrumentationDelayedPartialTestClass # rubocop:disable Lint/ConstantDefinitionInBlock
+              def test_method
+                43
+              end
+            end
+
+            payload = nil
+            expect(component.probe_notifier_worker).to receive(:add_snapshot) do |payload_|
+              payload = payload_
+            end
+
+            expect(InstrumentationDelayedPartialTestClass.new.test_method).to eq(43)
+            component.probe_notifier_worker.flush
+
+            snapshot = payload.fetch(:"debugger.snapshot")
+            expect(snapshot).to match(
+              id: String,
+              timestamp: Integer,
+              evaluationErrors: [],
+              probe: {id: '1234', version: 0, location: {
+                method: 'test_method', type: 'InstrumentationDelayedPartialTestClass',
+              }},
+              language: 'ruby',
+              # TODO the stack trace here does not contain the target method
+              # as the first frame - see the comment in Instrumenter.
+              stack: Array,
+              captures: nil,
+            )
+          end
+        end
+
+        context 'when class exists and target method virtual' do
+          let(:probe) do
+            Datadog::DI::Probe.new(id: "1234", type: :log,
+              type_name: 'InstrumentationVirtualTestClass', method_name: 'test_method',
+              capture_snapshot: false,)
+          end
+
+          it 'invokes probe and creates expected snapshot' do
+            class InstrumentationVirtualTestClass # rubocop:disable Lint/ConstantDefinitionInBlock
+              def method_missing(name) # rubocop:disable Style/MissingRespondToMissing
+                name
+              end
+            end
+
+            expect(component.transport).to receive(:send_request).at_least(:once)
+            expect(probe_manager.add_probe(probe)).to be true
+
+            payload = nil
+            expect(component.probe_notifier_worker).to receive(:add_snapshot) do |payload_|
+              payload = payload_
+            end
+
+            expect(InstrumentationVirtualTestClass.new.test_method).to eq(:test_method)
+            component.probe_notifier_worker.flush
+
+            snapshot = payload.fetch(:"debugger.snapshot")
+            expect(snapshot).to match(
+              id: String,
+              timestamp: Integer,
+              evaluationErrors: [],
+              probe: {id: '1234', version: 0, location: {
+                method: 'test_method', type: 'InstrumentationVirtualTestClass',
+              }},
+              language: 'ruby',
+              # TODO the stack trace here does not contain the target method
+              # as the first frame - see the comment in Instrumenter.
+              stack: Array,
+              captures: nil,
+            )
+          end
+        end
       end
 
       context 'enriched probe' do