Measure Ruby blocks using Gitlab::Metrics

This allows measuring of timings of arbitrary Ruby blocks, this allows
for more fine grained performance monitoring. Custom values and tags can
also be attached to a block.

2 files changed, 67 insertions, 0 deletions

diff --git a/lib/gitlab/metrics.rb b/lib/gitlab/metrics.rb
index 88a265c6af2..4a3f47b5a95 100644
--- a/lib/gitlab/metrics.rb
+++ b/lib/gitlab/metrics.rb
@@ -70,6 +70,32 @@ module Gitlab
       value.to_s.gsub('=', '\\=')
     end
 
+    # Measures the execution time of a block.
+    #
+    # Example:
+    #
+    #     Gitlab::Metrics.measure(:find_by_username_timings) do
+    #       User.find_by_username(some_username)
+    #     end
+    #
+    # series - The name of the series to store the data in.
+    # values - A Hash containing extra values to add to the metric.
+    # tags - A Hash containing extra tags to add to the metric.
+    #
+    # Returns the value yielded by the supplied block.
+    def self.measure(series, values = {}, tags = {})
+      return yield unless Transaction.current
+
+      start = Time.now.to_f
+      retval = yield
+      duration = (Time.now.to_f - start) * 1000.0
+      values = values.merge(duration: duration)
+
+      Transaction.current.add_metric(series, values, tags)
+
+      retval
+    end
+
     # When enabled this should be set before being used as the usual pattern
     # "@foo ||= bar" is _not_ thread-safe.
     if enabled?
diff --git a/spec/lib/gitlab/metrics_spec.rb b/spec/lib/gitlab/metrics_spec.rb
index 93dd958168d..8f63a5f2043 100644
--- a/spec/lib/gitlab/metrics_spec.rb
+++ b/spec/lib/gitlab/metrics_spec.rb
@@ -60,4 +60,45 @@ describe Gitlab::Metrics do
       expect(described_class.escape_value(10)).to eq('10')
     end
   end
+
+  describe '.measure' do
+    context 'without a transaction' do
+      it 'returns the return value of the block' do
+        val = Gitlab::Metrics.measure(:foo) { 10 }
+
+        expect(val).to eq(10)
+      end
+    end
+
+    context 'with a transaction' do
+      let(:transaction) { Gitlab::Metrics::Transaction.new }
+
+      before do
+        allow(Gitlab::Metrics::Transaction).to receive(:current).
+          and_return(transaction)
+      end
+
+      it 'adds a metric to the current transaction' do
+        expect(transaction).to receive(:add_metric).
+          with(:foo, { duration: a_kind_of(Numeric) }, { tag: 'value' })
+
+        Gitlab::Metrics.measure(:foo, {}, tag: 'value') { 10 }
+      end
+
+      it 'supports adding of custom values' do
+        values = { duration: a_kind_of(Numeric), number: 10 }
+
+        expect(transaction).to receive(:add_metric).
+          with(:foo, values, { tag: 'value' })
+
+        Gitlab::Metrics.measure(:foo, { number: 10 }, tag: 'value') { 10 }
+      end
+
+      it 'returns the return value of the block' do
+        val = Gitlab::Metrics.measure(:foo) { 10 }
+
+        expect(val).to eq(10)
+      end
+    end
+  end
 end