Merge branch 'db-load-balancing-hijack-connection' into 'master'

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

Closes #3191

See merge request !2707

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