Track and act upon the number of executed queriesquery-counts

This ensures that we have more visibility in the number of SQL queries
that are executed in web requests. The current threshold is hardcoded to
100 as we will rarely (maybe once or twice) change it.

In production and development we use Sentry if enabled, in the test
environment we raise an error. This feature is also only enabled in
production/staging when running on GitLab.com as it's not very useful to
other users.

4 files changed, 185 insertions, 0 deletions

diff --git a/lib/gitlab/query_limiting.rb b/lib/gitlab/query_limiting.rb
new file mode 100644
index 00000000000..f64f1757144
--- /dev/null
+++ b/lib/gitlab/query_limiting.rb
@@ -0,0 +1,36 @@
+module Gitlab
+  module QueryLimiting
+    # Returns true if we should enable tracking of query counts.
+    #
+    # This is only enabled in production/staging if we're running on GitLab.com.
+    # This ensures we don't produce any errors that users can't do anything
+    # about themselves.
+    def self.enable?
+      Gitlab.com? || Rails.env.development? || Rails.env.test?
+    end
+
+    # Allows the current request to execute any number of SQL queries.
+    #
+    # This method should _only_ be used when there's a corresponding issue to
+    # reduce the number of queries.
+    #
+    # The issue URL is only meant to push developers into creating an issue
+    # instead of blindly whitelisting offending blocks of code.
+    def self.whitelist(issue_url)
+      return unless enable_whitelist?
+
+      unless issue_url.start_with?('https://')
+        raise(
+          ArgumentError,
+          'You must provide a valid issue URL in order to whitelist a block of code'
+        )
+      end
+
+      Transaction&.current&.whitelisted = true
+    end
+
+    def self.enable_whitelist?
+      Rails.env.development? || Rails.env.test?
+    end
+  end
+end
diff --git a/lib/gitlab/query_limiting/active_support_subscriber.rb b/lib/gitlab/query_limiting/active_support_subscriber.rb
new file mode 100644
index 00000000000..66049c94ec6
--- /dev/null
+++ b/lib/gitlab/query_limiting/active_support_subscriber.rb
@@ -0,0 +1,11 @@
+module Gitlab
+  module QueryLimiting
+    class ActiveSupportSubscriber < ActiveSupport::Subscriber
+      attach_to :active_record
+
+      def sql(*)
+        Transaction.current&.increment
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/query_limiting/middleware.rb b/lib/gitlab/query_limiting/middleware.rb
new file mode 100644
index 00000000000..949ae79a047
--- /dev/null
+++ b/lib/gitlab/query_limiting/middleware.rb
@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Gitlab
+  module QueryLimiting
+    # Middleware for reporting (or raising) when a request performs more than a
+    # certain amount of database queries.
+    class Middleware
+      CONTROLLER_KEY = 'action_controller.instance'.freeze
+      ENDPOINT_KEY = 'api.endpoint'.freeze
+
+      def initialize(app)
+        @app = app
+      end
+
+      def call(env)
+        transaction, retval = Transaction.run do
+          @app.call(env)
+        end
+
+        transaction.action = action_name(env)
+        transaction.act_upon_results
+
+        retval
+      end
+
+      def action_name(env)
+        if env[CONTROLLER_KEY]
+          action_for_rails(env)
+        elsif env[ENDPOINT_KEY]
+          action_for_grape(env)
+        end
+      end
+
+      private
+
+      def action_for_rails(env)
+        controller = env[CONTROLLER_KEY]
+        action = "#{controller.class.name}##{controller.action_name}"
+
+        if controller.content_type == 'text/html'
+          action
+        else
+          "#{action} (#{controller.content_type})"
+        end
+      end
+
+      def action_for_grape(env)
+        endpoint = env[ENDPOINT_KEY]
+        route = endpoint.route rescue nil
+
+        "#{route.request_method} #{route.path}" if route
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/query_limiting/transaction.rb b/lib/gitlab/query_limiting/transaction.rb
new file mode 100644
index 00000000000..3cbafadb0d0
--- /dev/null
+++ b/lib/gitlab/query_limiting/transaction.rb
@@ -0,0 +1,83 @@
+module Gitlab
+  module QueryLimiting
+    class Transaction
+      THREAD_KEY = :__gitlab_query_counts_transaction
+
+      attr_accessor :count, :whitelisted
+
+      # The name of the action (e.g. `UsersController#show`) that is being
+      # executed.
+      attr_accessor :action
+
+      # The maximum number of SQL queries that can be executed in a request. For
+      # the sake of keeping things simple we hardcode this value here, it's not
+      # supposed to be changed very often anyway.
+      THRESHOLD = 100
+
+      # Error that is raised whenever exceeding the maximum number of queries.
+      ThresholdExceededError = Class.new(StandardError)
+
+      def self.current
+        Thread.current[THREAD_KEY]
+      end
+
+      # Starts a new transaction and returns it and the blocks' return value.
+      #
+      # Example:
+      #
+      #     transaction, retval = Transaction.run do
+      #       10
+      #     end
+      #
+      #     retval # => 10
+      def self.run
+        transaction = new
+        Thread.current[THREAD_KEY] = transaction
+
+        [transaction, yield]
+      ensure
+        Thread.current[THREAD_KEY] = nil
+      end
+
+      def initialize
+        @action = nil
+        @count = 0
+        @whitelisted = false
+      end
+
+      # Sends a notification based on the number of executed SQL queries.
+      def act_upon_results
+        return unless threshold_exceeded?
+
+        error = ThresholdExceededError.new(error_message)
+
+        if raise_error?
+          raise(error)
+        else
+          # Raven automatically logs to the Rails log if disabled, thus we don't
+          # need to manually log anything in case Sentry support is not enabled.
+          Raven.capture_exception(error)
+        end
+      end
+
+      def increment
+        @count += 1 unless whitelisted
+      end
+
+      def raise_error?
+        Rails.env.test?
+      end
+
+      def threshold_exceeded?
+        count > THRESHOLD
+      end
+
+      def error_message
+        header = 'Too many SQL queries were executed'
+        header += " in #{action}" if action
+
+        "#{header}: a maximum of #{THRESHOLD} is allowed but #{count} SQL queries were executed"
+      end
+    end
+  end
+end