From bcee44ad33d8a84822a8df068d47812594c445a3 Mon Sep 17 00:00:00 2001
From: Yorick Peterse <yorickpeterse@gmail.com>
Date: Tue, 15 Dec 2015 17:23:23 +0100
Subject: [PATCH] Instrument all ActiveRecord model methods

This works by searching the raw source code for any references to
commonly used ActiveRecord methods. While not bulletproof it saves us
from having to list hundreds of methods by hand. It also ensures that
(most) newly added methods are instrumented automatically.

This _only_ instruments models defined in app/models, should a model
reside somewhere else (e.g. somewhere in lib/) it _won't_ be
instrumented.
---
 Gemfile                        |  1 +
 Gemfile.lock                   |  1 +
 config/initializers/metrics.rb | 21 +++++++++++++++++++++
 3 files changed, 23 insertions(+)

diff --git a/Gemfile b/Gemfile
index 90db2c43006..e9e5c7df075 100644
--- a/Gemfile
+++ b/Gemfile
@@ -210,6 +210,7 @@ gem 'net-ssh',            '~> 3.0.1'
 
 # Metrics
 group :metrics do
+  gem 'method_source', '~> 0.8', require: false
   gem 'influxdb', '~> 0.2', require: false
   gem 'connection_pool', '~> 2.0', require: false
 end
diff --git a/Gemfile.lock b/Gemfile.lock
index 7d592ba93a7..3f301111224 100644
--- a/Gemfile.lock
+++ b/Gemfile.lock
@@ -887,6 +887,7 @@ DEPENDENCIES
   kaminari (~> 0.16.3)
   letter_opener (~> 1.1.2)
   mail_room (~> 0.6.1)
+  method_source (~> 0.8)
   minitest (~> 5.7.0)
   mousetrap-rails (~> 1.4.6)
   mysql2 (~> 0.3.16)
diff --git a/config/initializers/metrics.rb b/config/initializers/metrics.rb
index 0ac4299dcba..a47d2bf59a6 100644
--- a/config/initializers/metrics.rb
+++ b/config/initializers/metrics.rb
@@ -2,6 +2,7 @@ if Gitlab::Metrics.enabled?
   require 'influxdb'
   require 'socket'
   require 'connection_pool'
+  require 'method_source'
 
   # These are manually require'd so the classes are registered properly with
   # ActiveSupport.
@@ -18,6 +19,26 @@ if Gitlab::Metrics.enabled?
     end
   end
 
+  # This instruments all methods residing in app/models that (appear to) use any
+  # of the ActiveRecord methods. This has to take place _after_ initializing as
+  # for some unknown reason calling eager_load! earlier breaks Devise.
+  Gitlab::Application.config.after_initialize do
+    Rails.application.eager_load!
+
+    models = Rails.root.join('app', 'models').to_s
+
+    regex = Regexp.union(
+      ActiveRecord::Querying.public_instance_methods(false).map(&:to_s)
+    )
+
+    Gitlab::Metrics::Instrumentation.
+      instrument_class_hierarchy(ActiveRecord::Base) do |_, method|
+        loc = method.source_location
+
+        loc && loc[0].start_with?(models) && method.source =~ regex
+      end
+  end
+
   Gitlab::Metrics::Instrumentation.configure do |config|
     config.instrument_instance_methods(Gitlab::Shell)