From 1f9d011fa7c6359cb02ce2897fc896e28ab107aa Mon Sep 17 00:00:00 2001
From: Maxime Bedard <maxime.bedard@shopify.com>
Date: Mon, 30 Oct 2017 15:16:05 -0400
Subject: [PATCH] Send measure metric even when the block raises

---
 lib/statsd/instrument.rb                    | 38 +++++++++++++--------
 lib/statsd/instrument/metric_expectation.rb |  2 +-
 test/statsd_instrumentation_test.rb         | 16 +++++++++
 test/statsd_test.rb                         | 13 ++++++-
 4 files changed, 53 insertions(+), 16 deletions(-)

diff --git a/lib/statsd/instrument.rb b/lib/statsd/instrument.rb
index 92de80d4..a36cba66 100644
--- a/lib/statsd/instrument.rb
+++ b/lib/statsd/instrument.rb
@@ -56,19 +56,21 @@ def self.generate_metric_name(metric_name, callee, *args)
       metric_name.respond_to?(:call) ? metric_name.call(callee, args).gsub('::', '.') : metric_name.gsub('::', '.')
     end
 
+    def self.duration
+      start = current_timestamp
+      yield
+      current_timestamp - start
+    end
+
     if Process.respond_to?(:clock_gettime)
       # @private
-      def self.duration
-        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-        yield
-        Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
+      def self.current_timestamp
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
       end
     else
       # @private
-      def self.duration
-        start = Time.now
-        yield
-        Time.now - start
+      def self.current_timestamp
+        Time.now
       end
     end
 
@@ -277,17 +279,25 @@ def backend
   #      http_response = StatsD.measure('HTTP.call.duration') do
   #        HTTP.get(url)
   #      end
-  def measure(key, value = nil, *metric_options, &block)
+  def measure(key, value = nil, *metric_options)
     if value.is_a?(Hash) && metric_options.empty?
       metric_options = [value]
       value = value.fetch(:value, nil)
     end
 
-    result = nil
-    value  = 1000 * StatsD::Instrument.duration { result = block.call } if block_given?
-    metric = collect_metric(hash_argument(metric_options).merge(type: :ms, name: key, value: value))
-    result = metric unless block_given?
-    result
+    if block_given?
+      result = nil
+      start = StatsD::Instrument.current_timestamp
+      begin
+        result = yield
+      ensure
+        value = 1000 * (StatsD::Instrument.current_timestamp - start)
+        collect_metric(hash_argument(metric_options).merge(type: :ms, name: key, value: value))
+        result
+      end
+    else
+      collect_metric(hash_argument(metric_options).merge(type: :ms, name: key, value: value))
+    end
   end
 
   # Emits a counter metric.
diff --git a/lib/statsd/instrument/metric_expectation.rb b/lib/statsd/instrument/metric_expectation.rb
index 1cea31cc..38926a2d 100644
--- a/lib/statsd/instrument/metric_expectation.rb
+++ b/lib/statsd/instrument/metric_expectation.rb
@@ -64,4 +64,4 @@ def to_s
   def inspect
     "#<StatsD::Instrument::MetricExpectation #{self.to_s}>"
   end
-end
\ No newline at end of file
+end
diff --git a/test/statsd_instrumentation_test.rb b/test/statsd_instrumentation_test.rb
index 8795fc40..b6d62deb 100644
--- a/test/statsd_instrumentation_test.rb
+++ b/test/statsd_instrumentation_test.rb
@@ -37,6 +37,10 @@ def ssl_post(arg)
     {:success => arg}
   end
 
+  def boom
+    raise "boom!"
+  end
+
   def purchase(arg)
     ssl_post(arg)
   end
@@ -198,6 +202,18 @@ def test_statsd_measure
     ActiveMerchant::UniqueGateway.statsd_remove_measure :ssl_post, 'ActiveMerchant.Gateway.ssl_post'
   end
 
+  def test_statsd_measure_when_exception_is_raised
+    ActiveMerchant::UniqueGateway.statsd_measure :boom, 'ActiveMerchant.Gateway.boom', sample_rate: 0.3
+
+    assert_statsd_measure('ActiveMerchant.Gateway.boom', sample_rate: 0.3) do
+      assert_raises("boom!") do
+        ActiveMerchant::UniqueGateway.new.boom
+      end
+    end
+  ensure
+    ActiveMerchant::UniqueGateway.statsd_remove_measure :boom, 'ActiveMerchant.Gateway.boom'
+  end
+
   def test_statsd_measure_uses_normalized_metric_name
     ActiveMerchant::UniqueGateway.statsd_measure :ssl_post, 'ActiveMerchant::Gateway.ssl_post'
 
diff --git a/test/statsd_test.rb b/test/statsd_test.rb
index f77a781a..0e87345d 100644
--- a/test/statsd_test.rb
+++ b/test/statsd_test.rb
@@ -36,13 +36,24 @@ def test_statsd_measure_with_explicit_value_and_sample_rate
   end
 
   def test_statsd_measure_with_benchmarked_block_duration
-    StatsD::Instrument.stubs(:duration).returns(1.12)
+    StatsD::Instrument.stubs(:current_timestamp).returns(0, 1.12)
     metric = capture_statsd_call do
       StatsD.measure('values.foobar') { 'foo' }
     end
     assert_equal 1120.0, metric.value
   end
 
+  def test_statsd_measure_with_benchmarked_block_duration_even_when_it_raises
+    StatsD::Instrument.stubs(:current_timestamp).returns(0, 1.12)
+    metric = capture_statsd_call do
+      exception = "boom!"
+      assert_raises(exception) do
+        StatsD.measure('values.foobar') { raise exception }
+      end
+    end
+    assert_equal 1120.0, metric.value
+  end
+
   def test_statsd_measure_returns_return_value_of_block
     return_value = StatsD.measure('values.foobar') { 'sarah' }
     assert_equal 'sarah', return_value