From f295a18e948827786b3236ae244775b5aa4ff0b7 Mon Sep 17 00:00:00 2001
From: Sam <sam.saffron@gmail.com>
Date: Wed, 28 Feb 2018 10:45:11 +1100
Subject: [PATCH] FIX: stop double counting net calls in logs

---
 lib/method_profiler.rb                  | 13 ++++++++++++-
 lib/middleware/request_tracker.rb       |  2 +-
 spec/components/method_profiler_spec.rb | 16 ++++++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/lib/method_profiler.rb b/lib/method_profiler.rb
index 9ad59b007ff..f624ea8f977 100644
--- a/lib/method_profiler.rb
+++ b/lib/method_profiler.rb
@@ -1,7 +1,16 @@
 # see https://samsaffron.com/archive/2017/10/18/fastest-way-to-profile-a-method-in-ruby
 class MethodProfiler
-  def self.patch(klass, methods, name)
+  def self.patch(klass, methods, name, no_recurse: false)
     patches = methods.map do |method_name|
+
+      recurse_protection = ""
+      if no_recurse
+        recurse_protection = <<~RUBY
+          return #{method_name}__mp_unpatched(*args, &blk) if @mp_recurse_protect_#{method_name}
+          @mp_recurse_protect_#{method_name} = true
+        RUBY
+      end
+
       <<~RUBY
       unless defined?(#{method_name}__mp_unpatched)
         alias_method :#{method_name}__mp_unpatched, :#{method_name}
@@ -9,6 +18,7 @@ class MethodProfiler
           unless prof = Thread.current[:_method_profiler]
             return #{method_name}__mp_unpatched(*args, &blk)
           end
+          #{recurse_protection}
           begin
             start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
             #{method_name}__mp_unpatched(*args, &blk)
@@ -16,6 +26,7 @@ class MethodProfiler
             data = (prof[:#{name}] ||= {duration: 0.0, calls: 0})
             data[:duration] += Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
             data[:calls] += 1
+            #{"@mp_recurse_protect_#{method_name} = false" if no_recurse}
           end
         end
       end
diff --git a/lib/middleware/request_tracker.rb b/lib/middleware/request_tracker.rb
index d53b029ee82..543a5c60dc2 100644
--- a/lib/middleware/request_tracker.rb
+++ b/lib/middleware/request_tracker.rb
@@ -27,7 +27,7 @@ class Middleware::RequestTracker
 
       MethodProfiler.patch(Net::HTTP, [
         :request
-      ], :net)
+      ], :net, no_recurse: true)
 
       MethodProfiler.patch(Excon::Connection, [
         :request
diff --git a/spec/components/method_profiler_spec.rb b/spec/components/method_profiler_spec.rb
index cbbfcff65a8..480280a5fae 100644
--- a/spec/components/method_profiler_spec.rb
+++ b/spec/components/method_profiler_spec.rb
@@ -5,6 +5,22 @@ describe MethodProfiler do
   class Sneetch
     def beach
     end
+
+    def recurse(count = 5)
+      if count > 0
+        recurse(count - 1)
+      end
+    end
+  end
+
+  it "can bypass recursion on demand" do
+    MethodProfiler.patch(Sneetch, [:recurse], :recurse, no_recurse: true)
+
+    MethodProfiler.start
+    Sneetch.new.recurse
+    result = MethodProfiler.stop
+
+    expect(result[:recurse][:calls]).to eq(1)
   end
 
   it "can transfer data between threads" do