DI: Fix line instrumentation when target line raises an exception (#4900)

Author: p-datadog

lib/datadog/di/instrumenter.rb

@@ -307,33 +307,52 @@ def hook_line(probe, &block)
             # are invoked for *each* line of Ruby executed.
             # TODO find out exactly when the path in trace point is relative.
             # Looks like this is the case when line trace point is not targeted?
-            if iseq || tp.lineno == probe.line_no && (
+            continue = iseq || tp.lineno == probe.line_no && (
               probe.file == tp.path || probe.file_matches?(tp.path)
             )
-              continue = if condition = probe.condition
-                locals = Instrumenter.get_local_variables(tp)
-                context = EL::Context.new(locals: locals, target: tp.self)
-                condition.satisfied?(context)
-              else
-                true
+
+            # We set the trace point on :return to be able to instrument
+            # 'end' lines. This also causes the trace point to be invoked on
+            # non-'end' lines when a line raises an exception, since the
+            # exception causes the method to stop executing and stack unwends.
+            # We do not want two invocations of the trace point.
+            # Therefore, if a trace point is invoked with a :line event,
+            # mark it as such and ignore subsequent :return events.
+            continue &&= if probe.executed_on_line?
+              tp.event == :line
+            else
+              if tp.event == :line
+                probe.executed_on_line!
               end
-              if continue and rate_limiter.nil? || rate_limiter.allow? # standard:disable Style/AndOr
-                serialized_locals = if probe.capture_snapshot?
-                  serializer.serialize_vars(locals || Instrumenter.get_local_variables(tp),
-                    depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
-                    attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count,)
-                end
-                if probe.capture_snapshot?
-                  serializer.serialize_value(tp.self,
-                    depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
-                    attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count,)
-                end
-                # & is to stop steep complaints, block is always present here.
-                block&.call(probe: probe,
-                  serialized_locals: serialized_locals,
-                  target_self: tp.self,
-                  path: tp.path, caller_locations: caller_locations)
+              true
+            end
+
+            continue &&= if condition = probe.condition
+              locals = Instrumenter.get_local_variables(tp)
+              context = EL::Context.new(locals: locals, target: tp.self)
+              condition.satisfied?(context)
+            else
+              true
+            end
+
+            continue &&= rate_limiter.nil? || rate_limiter.allow? # standard:disable Style/AndOr
+
+            if continue
+              serialized_locals = if probe.capture_snapshot?
+                serializer.serialize_vars(locals || Instrumenter.get_local_variables(tp),
+                  depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
+                  attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count,)
               end
+              if probe.capture_snapshot?
+                serializer.serialize_value(tp.self,
+                  depth: probe.max_capture_depth || settings.dynamic_instrumentation.max_capture_depth,
+                  attribute_count: probe.max_capture_attribute_count || settings.dynamic_instrumentation.max_capture_attribute_count,)
+              end
+              # & is to stop steep complaints, block is always present here.
+              block&.call(probe: probe,
+                serialized_locals: serialized_locals,
+                target_self: tp.self,
+                path: tp.path, caller_locations: caller_locations)
             end
           rescue => exc
             raise if settings.dynamic_instrumentation.internal.propagate_all_exceptions

lib/datadog/di/probe.rb

@@ -190,6 +190,15 @@ def file_matches?(path)
       def emitting_notified?
         !!@emitting_notified
       end
+
+      def executed_on_line?
+        !!@executed_on_line
+      end
+
+      def executed_on_line!
+        # TODO lock?
+        @executed_on_line = true
+      end
     end
   end
 end

sig/datadog/di/probe.rbs

@@ -20,6 +20,8 @@ module Datadog
       @capture_snapshot: bool
 
       @rate_limiter: Datadog::Core::RateLimiter
+      
+      @executed_on_line: bool?
 
       def initialize: (id: String, type: Symbol, ?file: String?, ?line_no: Integer?, ?type_name: String?, ?method_name: String?, ?template: String?, ?capture_snapshot: bool,
 	?max_capture_depth: Integer, ?max_capture_attribute_count: Integer?, ?rate_limit: Integer) -> void
@@ -52,6 +54,9 @@ module Datadog
       def location: () -> ::String
       def file_matches?: (String path) -> bool
 
+      def executed_on_line!: -> void
+      def executed_on_line?: -> bool
+      
       attr_accessor instrumentation_module: Module?
       attr_accessor instrumentation_trace_point: TracePoint?
       attr_accessor instrumented_path: String?

spec/datadog/di/hook_line_load.rb

@@ -1,9 +1,9 @@
 # Comment line - not executable
 
 class HookLineLoadTestClass
-  def test_method       # Line 2
-    42                  # Line 3
-  end
+  def test_method       # Line 4
+    42                  # Line 5
+  end                   # Line 6
 
   def test_method_with_local
     local = 42 # standard:disable Style/RedundantAssignment
@@ -16,15 +16,24 @@ def test_method_with_arg(arg)
 end
 
 class HookLineIvarLoadTestClass
+  class TestException < StandardError
+  end
+
   def initialize
     @ivar = 42
   end
 
   def test_method
     1337                 # Line 24
   end
+
+  def test_exception
+    local = 42
+    raise TestException, 'Intentional exception'       # Line 32
+    local               # Line 33 # standard:disable Lint/UnreachableCode
+  end
 end
 
-unless (actual = File.read(__FILE__).count("\n")) == 30
-  raise "Wrong number of lines in hook_line_load.rb: actual #{actual}, expected 30"
+unless (actual = File.read(__FILE__).count("\n")) == 39
+  raise "Wrong number of lines in hook_line_load.rb: actual #{actual}, expected 39"
 end

spec/datadog/di/instrumenter_spec.rb

@@ -886,6 +886,37 @@
       %i[caller_locations path probe serialized_locals target_self]
     end
 
+    shared_examples 'multiple invocations' do
+      # Since the instrumentation mutates the state of the probe,
+      # verify that the state mutation is not breaking the instrumentation.
+      context 'when the code is executed multiple times' do
+        before do
+          load File.join(File.dirname(__FILE__), 'hook_line_load.rb')
+        end
+
+        let(:probe) do
+          Datadog::DI::Probe.new(file: 'hook_line_load.rb', line_no: 5,
+            id: 1, type: :log, rate_limit: rate_limit)
+        end
+
+        it 'invokes the instrumentation every time' do
+          expect_any_instance_of(TracePoint).to receive(:enable).and_call_original
+
+          instrumenter.hook_line(probe) do |payload|
+            observed_calls << payload
+          end
+
+          HookLineLoadTestClass.new.test_method
+          HookLineLoadTestClass.new.test_method
+
+          expect(observed_calls.length).to eq 2
+
+          expect(observed_calls[0].keys.sort).to eq(call_keys)
+          expect(observed_calls[1].keys.sort).to eq(call_keys)
+        end
+      end
+    end
+
     context 'when called without a block' do
       let(:probe) do
         instance_double(Datadog::DI::Probe)
@@ -999,7 +1030,7 @@
       end
     end
 
-    context 'line inside of method' do
+    context 'line inside of method without code tracking' do
       before do
         # We need untargeted trace points for this test since the line
         # being instrumented has already been loaded.
@@ -1035,6 +1066,8 @@
           expect(File.basename(frame.path)).to eq 'hook_line.rb'
         end
       end
+
+      include_examples 'multiple invocations'
     end
 
     context 'when hooking same line twice with identical but different probes' do
@@ -1116,6 +1149,51 @@
         expect(observed_calls.first).to be_a(Hash)
       end
 
+      context 'end line of a method' do
+        before do
+          load File.join(File.dirname(__FILE__), 'hook_line_load.rb')
+        end
+
+        let(:probe) do
+          Datadog::DI::Probe.new(file: 'hook_line_load.rb', line_no: 6,
+            id: 1, type: :log, rate_limit: rate_limit)
+        end
+
+        it 'invokes the instrumentation' do
+          expect_any_instance_of(TracePoint).to receive(:enable).and_call_original
+
+          instrumenter.hook_line(probe) do |payload|
+            observed_calls << payload
+          end
+
+          HookLineLoadTestClass.new.test_method
+
+          expect(observed_calls.length).to eq 1
+
+          expect(observed_calls[0].keys.sort).to eq(call_keys)
+        end
+
+        # Since the instrumentation mutates the state of the probe,
+        # verify that the state mutation is not breaking the instrumentation.
+        context 'when the code is executed multiple times' do
+          it 'invokes the instrumentation every time' do
+            expect_any_instance_of(TracePoint).to receive(:enable).and_call_original
+
+            instrumenter.hook_line(probe) do |payload|
+              observed_calls << payload
+            end
+
+            HookLineLoadTestClass.new.test_method
+            HookLineLoadTestClass.new.test_method
+
+            expect(observed_calls.length).to eq 2
+
+            expect(observed_calls[0].keys.sort).to eq(call_keys)
+            expect(observed_calls[1].keys.sort).to eq(call_keys)
+          end
+        end
+      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,
@@ -1129,6 +1207,8 @@
           end.to raise_error(Datadog::DI::Error::DITargetNotInRegistry, /File matching probe path.*was loaded and is not in code tracker registry/)
         end
       end
+
+      include_examples 'multiple invocations'
     end
 
     context 'when method is recursive' do
@@ -1208,6 +1288,39 @@
       end
     end
 
+    context 'when the instrumented line raises an exception' do
+      include_context 'with code tracking'
+
+      before do
+        load File.join(File.dirname(__FILE__), 'hook_line_load.rb')
+      end
+
+      let(:probe) do
+        Datadog::DI::Probe.new(file: 'hook_line_load.rb', line_no: 32,
+          id: 1, type: :log, rate_limit: rate_limit)
+      end
+
+      let(:payload) do
+        expect_any_instance_of(TracePoint).to receive(:enable).and_call_original
+
+        instrumenter.hook_line(probe) do |payload|
+          observed_calls << payload
+        end
+
+        expect do
+          HookLineIvarLoadTestClass.new.test_exception
+        end.to raise_error(HookLineIvarLoadTestClass::TestException)
+
+        expect(observed_calls.length).to eq 1
+        observed_calls.first
+      end
+
+      it 'invokes callback with expected keys' do
+        expect(payload).to be_a(Hash)
+        expect(payload.keys.sort).to eq(call_keys)
+      end
+    end
+
     context 'when there is a condition' do
       include_context 'with code tracking'