Track and act upon the number of executed queries

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.

3 files changed, 149 insertions, 0 deletions

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