DEBUG-2334 correctly instrument methods that yield

lib/datadog/di/instrumenter.rb

@@ -106,7 +106,7 @@ def hook_method(probe, &block)
         settings = self.settings
 
         mod = Module.new do
-          define_method(method_name) do |*args, **kwargs| # steep:ignore
+          define_method(method_name) do |*args, **kwargs, &target_block| # steep:ignore
             if rate_limiter.nil? || rate_limiter.allow?
               # Arguments may be mutated by the method, therefore
               # they need to be serialized prior to method invocation.
@@ -120,14 +120,14 @@ def hook_method(probe, &block)
               duration = Benchmark.realtime do # steep:ignore
                 rv = if args.any?
                   if kwargs.any?
-                    super(*args, **kwargs)
+                    super(*args, **kwargs, &target_block)
                   else
-                    super(*args)
+                    super(*args, &target_block)
                   end
                 elsif kwargs.any?
-                  super(**kwargs)
+                  super(**kwargs, &target_block)
                 else
-                  super()
+                  super(&target_block)
                 end
               end
               # The method itself is not part of the stack trace because

spec/datadog/di/hook_method.rb

@@ -15,6 +15,10 @@ def hook_test_method_with_pos_and_kwarg(arg, kwarg:)
     [arg, kwarg]
   end
 
+  def yielding(arg)
+    yield arg
+  end
+
   def recursive(depth)
     if depth > 0
       recursive(depth - 1) + '-'

spec/datadog/di/instrumenter_spec.rb

@@ -77,6 +77,31 @@
       end
     end
 
+    context 'when target method yields to a block' do
+      let(:probe_args) do
+        {type_name: 'HookTestClass', method_name: 'yielding'}
+      end
+
+      it 'invokes callback' do
+        instrumenter.hook_method(probe) do |payload|
+          observed_calls << payload
+        end
+
+        yielded_value = nil
+        expect(HookTestClass.new.yielding('hello') do |value|
+          yielded_value = value
+          [value]
+        end).to eq ['hello']
+
+        expect(yielded_value).to eq('hello')
+
+        expect(observed_calls.length).to eq 1
+        expect(observed_calls.first.keys.sort).to eq call_keys
+        expect(observed_calls.first[:rv]).to eq ['hello']
+        expect(observed_calls.first[:duration]).to be_a(Float)
+      end
+    end
+
     context 'positional args' do
       context 'without snapshot capture' do
         let(:probe_args) do