From 66a997a91403eef62ffd9fb789e899619d021a26 Mon Sep 17 00:00:00 2001
From: Yorick Peterse <yorickpeterse@gmail.com>
Date: Mon, 4 Jan 2016 12:14:36 +0100
Subject: [PATCH] Track total query/view timings in transactions

---
 config/initializers/metrics.rb                |  1 +
 lib/gitlab/metrics/subscribers/action_view.rb |  1 +
 .../metrics/subscribers/active_record.rb      | 22 ++++++++++++
 .../metrics/subscribers/action_view_spec.rb   |  3 ++
 .../metrics/subscribers/active_record_spec.rb | 35 +++++++++++++++++++
 5 files changed, 62 insertions(+)
 create mode 100644 lib/gitlab/metrics/subscribers/active_record.rb
 create mode 100644 spec/lib/gitlab/metrics/subscribers/active_record_spec.rb

diff --git a/config/initializers/metrics.rb b/config/initializers/metrics.rb
index ebb20be7283..52ace27b7ae 100644
--- a/config/initializers/metrics.rb
+++ b/config/initializers/metrics.rb
@@ -6,6 +6,7 @@ if Gitlab::Metrics.enabled?
   # These are manually require'd so the classes are registered properly with
   # ActiveSupport.
   require 'gitlab/metrics/subscribers/action_view'
+  require 'gitlab/metrics/subscribers/active_record'
 
   Gitlab::Application.configure do |config|
     config.middleware.use(Gitlab::Metrics::RackMiddleware)
diff --git a/lib/gitlab/metrics/subscribers/action_view.rb b/lib/gitlab/metrics/subscribers/action_view.rb
index 7e0dcf99d92..7c0105d543a 100644
--- a/lib/gitlab/metrics/subscribers/action_view.rb
+++ b/lib/gitlab/metrics/subscribers/action_view.rb
@@ -19,6 +19,7 @@ module Gitlab
           values = values_for(event)
           tags   = tags_for(event)
 
+          current_transaction.increment(:view_duration, event.duration)
           current_transaction.add_metric(SERIES, values, tags)
         end
 
diff --git a/lib/gitlab/metrics/subscribers/active_record.rb b/lib/gitlab/metrics/subscribers/active_record.rb
new file mode 100644
index 00000000000..8008b3bc895
--- /dev/null
+++ b/lib/gitlab/metrics/subscribers/active_record.rb
@@ -0,0 +1,22 @@
+module Gitlab
+  module Metrics
+    module Subscribers
+      # Class for tracking the total query duration of a transaction.
+      class ActiveRecord < ActiveSupport::Subscriber
+        attach_to :active_record
+
+        def sql(event)
+          return unless current_transaction
+
+          current_transaction.increment(:sql_duration, event.duration)
+        end
+
+        private
+
+        def current_transaction
+          Transaction.current
+        end
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/metrics/subscribers/action_view_spec.rb b/spec/lib/gitlab/metrics/subscribers/action_view_spec.rb
index bca76ca5a69..699d50f770a 100644
--- a/spec/lib/gitlab/metrics/subscribers/action_view_spec.rb
+++ b/spec/lib/gitlab/metrics/subscribers/action_view_spec.rb
@@ -28,6 +28,9 @@ describe Gitlab::Metrics::Subscribers::ActionView do
         line: 4
       }
 
+      expect(transaction).to receive(:increment).
+        with(:view_duration, 2.1)
+
       expect(transaction).to receive(:add_metric).
         with(described_class::SERIES, values, tags)
 
diff --git a/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
new file mode 100644
index 00000000000..9ecedd934c6
--- /dev/null
+++ b/spec/lib/gitlab/metrics/subscribers/active_record_spec.rb
@@ -0,0 +1,35 @@
+require 'spec_helper'
+
+describe Gitlab::Metrics::Subscribers::ActiveRecord do
+  let(:transaction) { Gitlab::Metrics::Transaction.new('rspec') }
+  let(:subscriber)  { described_class.new }
+
+  let(:event) do
+    double(:event, duration: 0.2,
+           payload:  { sql: 'SELECT * FROM users WHERE id = 10' })
+  end
+
+  describe '#sql' do
+    describe 'without a current transaction' do
+      it 'simply returns' do
+        expect_any_instance_of(Gitlab::Metrics::Transaction).
+          to_not receive(:increment)
+
+        subscriber.sql(event)
+      end
+    end
+
+    describe 'with a current transaction' do
+      it 'increments the :sql_duration value' do
+        expect(subscriber).to receive(:current_transaction).
+          at_least(:once).
+          and_return(transaction)
+
+        expect(transaction).to receive(:increment).
+          with(:sql_duration, 0.2)
+
+        subscriber.sql(event)
+      end
+    end
+  end
+end