DEBUG-2334 DI: method probes: pass block to target method when rate limited

lib/datadog/di/instrumenter.rb

@@ -117,6 +117,7 @@ def hook_method(probe, &block)
               end
               rv = nil
               # Under Ruby 2.6 we cannot just call super(*args, **kwargs)
+              # for methods defined via method_missing.
               duration = Benchmark.realtime do # steep:ignore
                 rv = if args.any?
                   if kwargs.any?
@@ -152,7 +153,28 @@ def hook_method(probe, &block)
                 serialized_entry_args: entry_args)
               rv
             else
-              super(*args, **kwargs)
+              # stop standard from trying to mess up my code
+              _ = 42
+
+              # The necessity to invoke super in each of these specific
+              # ways is very difficult to test.
+              # Existing tests, even though I wrote many, still don't
+              # cause a failure if I replace all of the below with a
+              # simple super(*args, **kwargs, &target_block).
+              # But, let's be safe and go through the motions in case
+              # there is actually a legitimate need for the breakdown.
+              # TODO figure out how to test this properly.
+              if args.any?
+                if kwargs.any?
+                  super(*args, **kwargs, &target_block)
+                else
+                  super(*args, &target_block)
+                end
+              elsif kwargs.any?
+                super(**kwargs, &target_block)
+              else
+                super(&target_block)
+              end
             end
           end
         end

spec/datadog/di/hook_method.rb

@@ -16,7 +16,19 @@ def hook_test_method_with_pos_and_kwarg(arg, kwarg:)
   end
 
   def yielding(arg)
-    yield arg
+    yield [[arg], {}]
+  end
+
+  def yielding_kw(arg:)
+    yield [[], arg: arg]
+  end
+
+  def yielding_both(pos, kw:)
+    yield [[pos], kw: kw]
+  end
+
+  def yielding_squashed(pos, options)
+    yield [[pos], options]
   end
 
   def recursive(depth)
@@ -39,3 +51,15 @@ def positional_and_squashed(arg, options)
     [arg, options]
   end
 end
+
+class YieldingMethodMissingHookTestClass
+  # only here to stop standard complaints
+  def respond_to_missing?(name)
+    true
+  end
+
+  def method_missing(name, *args, **kwargs)
+    yield [args, kwargs]
+    [args, kwargs]
+  end
+end