Hijack AR::Base.connection in the DB load balancer

This changes the load balancing logic so we hijack
`ActiveRecord::Base.connection` instead of hijacking "connection" on a
per model basis. This ensures _all_ code that might use it (including
Rails internals) goes through the database load balancer. This in turn
should make the application more resistant to failures on the primary.

Fixes https://gitlab.com/gitlab-org/gitlab-ee/issues/3191

changelogs/unreleased-ee/db-load-balancing-hijack-connection.yml

+---
+title: >
+  Ensure all database queries are routed through the database load balancer when
+  load balancing is enabled
+merge_request: 2707
+author:
+type: changed

lib/gitlab/database/load_balancing.rb

@@ -61,15 +61,10 @@ module Gitlab
       def self.configure_proxy
         self.proxy = ConnectionProxy.new(hosts)
 
-        # ActiveRecordProxy's methods are made available as class methods in
-        # ActiveRecord::Base, while still allowing the use of `super`.
+        # This hijacks the "connection" method to ensure both
+        # `ActiveRecord::Base.connection` and all models use the same load
+        # balancing proxy.
         ActiveRecord::Base.singleton_class.prepend(ActiveRecordProxy)
-
-        # The above will only patch newly defined models, so we also need to
-        # patch existing ones.
-        active_record_models.each do |model|
-          model.singleton_class.prepend(ModelProxy)
-        end
       end
 
       def self.active_record_models

lib/gitlab/database/load_balancing/active_record_proxy.rb

 module Gitlab
   module Database
     module LoadBalancing
-      # Module injected into ActiveRecord::Base to allow proxying of subclasses.
+      # Module injected into ActiveRecord::Base to allow hijacking of the
+      # "connection" method.
       module ActiveRecordProxy
-        def inherited(by)
-          super(by)
-
-          # The methods in ModelProxy will become available as class methods for
-          # the class defined in `by`.
-          by.singleton_class.prepend(ModelProxy)
+        def connection
+          LoadBalancing.proxy
         end
       end
     end

lib/gitlab/database/load_balancing/load_balancer.rb

@@ -6,8 +6,8 @@ module Gitlab
       # Each host in the load balancer uses the same credentials as the primary
       # database.
       #
-      # This class *requires* that `ActiveRecord::Base.connection` always
-      # returns a connection to the primary.
+      # This class *requires* that `ActiveRecord::Base.retrieve_connection`
+      # always returns a connection to the primary.
       class LoadBalancer
         CACHE_KEY = :gitlab_load_balancer_host
 
@@ -63,7 +63,7 @@ module Gitlab
           # Instead of immediately grinding to a halt we'll retry the operation
           # a few times.
           retry_with_backoff do
-            yield ActiveRecord::Base.connection
+            yield ActiveRecord::Base.retrieve_connection
           end
         end
 

lib/gitlab/database/load_balancing/model_proxy.rb

-module Gitlab
-  module Database
-    module LoadBalancing
-      # Modle injected into models in order to redirect connections to a
-      # ConnectionProxy.
-      module ModelProxy
-        def connection
-          LoadBalancing.proxy
-        end
-      end
-    end
-  end
-end

spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb

 require 'spec_helper'
 
 describe Gitlab::Database::LoadBalancing::ActiveRecordProxy do
-  describe '#inherited' do
-    it 'adds the ModelProxy module to the singleton class' do
-      base = Class.new do
+  describe '#connection' do
+    it 'returns a connection proxy' do
+      dummy = Class.new do
         include Gitlab::Database::LoadBalancing::ActiveRecordProxy
       end
 
-      model = Class.new(base)
+      proxy = double(:proxy)
 
-      expect(model.included_modules).to include(described_class)
+      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
+        .and_return(proxy)
+
+      expect(dummy.new.connection).to eq(proxy)
     end
   end
 end

spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb

@@ -86,14 +86,14 @@ describe Gitlab::Database::LoadBalancing::LoadBalancer do
       expect(lb).to receive(:read_write).and_call_original
 
       expect { |b| lb.read(&b) }
-        .to yield_with_args(ActiveRecord::Base.connection)
+        .to yield_with_args(ActiveRecord::Base.retrieve_connection)
     end
   end
 
   describe '#read_write' do
     it 'yields a connection for a write' do
       expect { |b| lb.read_write(&b) }
-        .to yield_with_args(ActiveRecord::Base.connection)
+        .to yield_with_args(ActiveRecord::Base.retrieve_connection)
     end
 
     it 'uses a retry with exponential backoffs' do

spec/lib/gitlab/database/load_balancing/model_proxy_spec.rb

-require 'spec_helper'
-
-describe Gitlab::Database::LoadBalancing::ModelProxy do
-  describe '#connection' do
-    it 'returns a connection proxy' do
-      dummy = Class.new do
-        include Gitlab::Database::LoadBalancing::ModelProxy
-      end
-
-      proxy = double(:proxy)
-
-      expect(Gitlab::Database::LoadBalancing).to receive(:proxy)
-        .and_return(proxy)
-
-      expect(dummy.new.connection).to eq(proxy)
-    end
-  end
-end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -106,17 +106,9 @@ describe Gitlab::Database::LoadBalancing do
     end
 
     it 'configures the connection proxy' do
-      model = double(:model)
-
       expect(ActiveRecord::Base.singleton_class).to receive(:prepend)
         .with(Gitlab::Database::LoadBalancing::ActiveRecordProxy)
 
-      expect(described_class).to receive(:active_record_models)
-        .and_return([model])
-
-      expect(model.singleton_class).to receive(:prepend)
-        .with(Gitlab::Database::LoadBalancing::ModelProxy)
-
       described_class.configure_proxy
     end
   end