Allow percentage rollout of load balancer of decomposed database

This introduces a `GITLAB_USE_MODEL_LOAD_BALANCING`
and `use_model_load_balancing` to enable percentage rollout
of used connection when accessing data.

The usage pattern is to redirect 1% of new connections using CI DB
to use a Load Balancer replicas (and presumbly later primary connection).

This is achieved by injecting a FF evaluation logic as part of LB setup
to re-define how connection is fetched.

.gitlab/ci/rails.gitlab-ci.yml

@@ -27,6 +27,7 @@
   variables:
     DECOMPOSED_DB: "true"
     GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: "main"
+    GITLAB_USE_MODEL_LOAD_BALANCING: "true"
 
 .rspec-base:
   extends: .rails-job-base

config/feature_flags/development/use_model_load_balancing.yml

+---
+name: use_model_load_balancing
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/73631
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/344797
+milestone: '14.5'
+type: development
+group: group::sharding
+default_enabled: false

lib/gitlab/database/load_balancing/setup.rb

@@ -5,7 +5,7 @@ module Gitlab
     module LoadBalancing
       # Class for setting up load balancing of a specific model.
       class Setup
-        attr_reader :configuration
+        attr_reader :model, :configuration
 
         def initialize(model, start_service_discovery: false)
           @model = model
@@ -15,8 +15,9 @@ module Gitlab
 
         def setup
           configure_connection
-          setup_load_balancer
+          setup_connection_proxy
           setup_service_discovery
+          setup_feature_flag_to_model_load_balancing
         end
 
         def configure_connection
@@ -36,28 +37,82 @@ module Gitlab
           @model.establish_connection(hash_config)
         end
 
-        def setup_load_balancer
-          lb = LoadBalancer.new(configuration)
-
+        def setup_connection_proxy
           # We just use a simple `class_attribute` here so we don't need to
           # inject any modules and/or expose unnecessary methods.
-          @model.class_attribute(:connection)
-          @model.class_attribute(:sticking)
+          setup_class_attribute(:connection, ConnectionProxy.new(load_balancer))
+          setup_class_attribute(:sticking, Sticking.new(load_balancer))
+        end
+
+        # TODO: This is temporary code to gradually redirect traffic to use
+        # a dedicated DB replicas, or DB primaries (depending on configuration)
+        # This implements a sticky behavior for the current request if enabled.
+        #
+        # This is needed for Phase 3 and Phase 4 of application rollout
+        # https://gitlab.com/groups/gitlab-org/-/epics/6160#progress
+        #
+        # If `GITLAB_USE_MODEL_LOAD_BALANCING` is set, its value is preferred
+        # Otherwise, a `use_model_load_balancing` FF value is used
+        def setup_feature_flag_to_model_load_balancing
+          return if active_record_base?
 
-          @model.connection = ConnectionProxy.new(lb)
-          @model.sticking = Sticking.new(lb)
+          @model.singleton_class.prepend(ModelLoadBalancingFeatureFlagMixin)
         end
 
         def setup_service_discovery
           return unless configuration.service_discovery_enabled?
 
-          lb = @model.connection.load_balancer
-          sv = ServiceDiscovery.new(lb, **configuration.service_discovery)
+          sv = ServiceDiscovery.new(load_balancer, **configuration.service_discovery)
 
           sv.perform_service_discovery
 
           sv.start if @start_service_discovery
         end
+
+        def load_balancer
+          @load_balancer ||= LoadBalancer.new(configuration)
+        end
+
+        private
+
+        def setup_class_attribute(attribute, value)
+          @model.class_attribute(attribute)
+          @model.public_send("#{attribute}=", value) # rubocop:disable GitlabSecurity/PublicSend
+        end
+
+        def active_record_base?
+          @model == ActiveRecord::Base
+        end
+
+        module ModelLoadBalancingFeatureFlagMixin
+          extend ActiveSupport::Concern
+
+          def use_model_load_balancing?
+            # Cache environment variable and return env variable first if defined
+            use_model_load_balancing_env = Gitlab::Utils.to_boolean(ENV["GITLAB_USE_MODEL_LOAD_BALANCING"])
+
+            unless use_model_load_balancing_env.nil?
+              return use_model_load_balancing_env
+            end
+
+            # Check a feature flag using RequestStore (if active)
+            return false unless Gitlab::SafeRequestStore.active?
+
+            Gitlab::SafeRequestStore.fetch(:use_model_load_balancing) do
+              Feature.enabled?(:use_model_load_balancing, default_enabled: :yaml)
+            end
+          end
+
+          # rubocop:disable Database/MultipleDatabases
+          def connection
+            use_model_load_balancing? ? super : ActiveRecord::Base.connection
+          end
+
+          def sticking
+            use_model_load_balancing? ? super : ActiveRecord::Base.sticking
+          end
+          # rubocop:enable Database/MultipleDatabases
+        end
       end
     end
   end

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

@@ -8,8 +8,9 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
       setup = described_class.new(ActiveRecord::Base)
 
       expect(setup).to receive(:configure_connection)
-      expect(setup).to receive(:setup_load_balancer)
+      expect(setup).to receive(:setup_connection_proxy)
       expect(setup).to receive(:setup_service_discovery)
+      expect(setup).to receive(:setup_feature_flag_to_model_load_balancing)
 
       setup.setup
     end
@@ -44,7 +45,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
     end
   end
 
-  describe '#setup_load_balancer' do
+  describe '#setup_connection_proxy' do
     it 'sets up the load balancer' do
       model = Class.new(ActiveRecord::Base)
       setup = described_class.new(model)
@@ -58,7 +59,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
         .with(setup.configuration)
         .and_return(lb)
 
-      setup.setup_load_balancer
+      setup.setup_connection_proxy
 
       expect(model.connection.load_balancer).to eq(lb)
       expect(model.sticking)
@@ -81,7 +82,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
         model = ActiveRecord::Base
         setup = described_class.new(model)
         sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-        lb = model.connection.load_balancer
 
         allow(setup.configuration)
           .to receive(:service_discovery_enabled?)
@@ -89,7 +89,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
 
         allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
           .to receive(:new)
-          .with(lb, setup.configuration.service_discovery)
+          .with(setup.load_balancer, setup.configuration.service_discovery)
           .and_return(sv)
 
         expect(sv).to receive(:perform_service_discovery)
@@ -102,7 +102,6 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
         model = ActiveRecord::Base
         setup = described_class.new(model, start_service_discovery: true)
         sv = instance_spy(Gitlab::Database::LoadBalancing::ServiceDiscovery)
-        lb = model.connection.load_balancer
 
         allow(setup.configuration)
           .to receive(:service_discovery_enabled?)
@@ -110,7 +109,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
 
         allow(Gitlab::Database::LoadBalancing::ServiceDiscovery)
           .to receive(:new)
-          .with(lb, setup.configuration.service_discovery)
+          .with(setup.load_balancer, setup.configuration.service_discovery)
           .and_return(sv)
 
         expect(sv).to receive(:perform_service_discovery)
@@ -120,4 +119,172 @@ RSpec.describe Gitlab::Database::LoadBalancing::Setup do
       end
     end
   end
+
+  describe '#setup_feature_flag_to_model_load_balancing', :reestablished_active_record_base do
+    using RSpec::Parameterized::TableSyntax
+
+    where do
+      {
+        "with model LB enabled it picks a dedicated CI connection" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: 'true',
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
+          request_store_active: false,
+          ff_use_model_load_balancing: nil,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'ci_replica', write: 'ci' }
+          }
+        },
+        "with model LB enabled and re-use of primary connection it uses CI connection for reads" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: 'true',
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
+          request_store_active: false,
+          ff_use_model_load_balancing: nil,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'ci_replica', write: 'main' }
+          }
+        },
+        "with model LB disabled it fallbacks to use main" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: 'false',
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
+          request_store_active: false,
+          ff_use_model_load_balancing: nil,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'main_replica', write: 'main' }
+          }
+        },
+        "with model LB disabled, but re-use configured it fallbacks to use main" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: 'false',
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
+          request_store_active: false,
+          ff_use_model_load_balancing: nil,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'main_replica', write: 'main' }
+          }
+        },
+        "with FF disabled without RequestStore it uses main" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
+          request_store_active: false,
+          ff_use_model_load_balancing: false,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'main_replica', write: 'main' }
+          }
+        },
+        "with FF enabled without RequestStore sticking of FF does not work, so it fallbacks to use main" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
+          request_store_active: false,
+          ff_use_model_load_balancing: true,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'main_replica', write: 'main' }
+          }
+        },
+        "with FF disabled with RequestStore it uses main" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
+          request_store_active: true,
+          ff_use_model_load_balancing: false,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'main_replica', write: 'main' }
+          }
+        },
+        "with FF enabled with RequestStore it sticks FF and uses CI connection" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: nil,
+          request_store_active: true,
+          ff_use_model_load_balancing: true,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'ci_replica', write: 'ci' }
+          }
+        },
+        "with re-use and FF enabled with RequestStore it sticks FF and uses CI connection for reads" => {
+          env_GITLAB_USE_MODEL_LOAD_BALANCING: nil,
+          env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci: 'main',
+          request_store_active: true,
+          ff_use_model_load_balancing: true,
+          expectations: {
+            main: { read: 'main_replica', write: 'main' },
+            ci: { read: 'ci_replica', write: 'main' }
+          }
+        }
+      }
+    end
+
+    with_them do
+      let(:ci_class) do
+        Class.new(ActiveRecord::Base) do
+          def self.name
+            'Ci::ApplicationRecordTemporary'
+          end
+
+          establish_connection ActiveRecord::DatabaseConfigurations::HashConfig.new(
+            Rails.env,
+            'ci',
+            ActiveRecord::Base.connection_db_config.configuration_hash
+          )
+        end
+      end
+
+      let(:models) do
+        {
+          main: ActiveRecord::Base,
+          ci: ci_class
+        }
+      end
+
+      around do |example|
+        if request_store_active
+          Gitlab::WithRequestStore.with_request_store do
+            RequestStore.clear!
+
+            example.run
+          end
+        else
+          example.run
+        end
+      end
+
+      before do
+        # Rewrite `class_attribute` to use rspec mocking and prevent modifying the objects
+        allow_next_instance_of(described_class) do |setup|
+          allow(setup).to receive(:configure_connection)
+
+          allow(setup).to receive(:setup_class_attribute) do |attribute, value|
+            allow(setup.model).to receive(attribute) { value }
+          end
+        end
+
+        stub_env('GITLAB_USE_MODEL_LOAD_BALANCING', env_GITLAB_USE_MODEL_LOAD_BALANCING)
+        stub_env('GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci', env_GITLAB_LOAD_BALANCING_REUSE_PRIMARY_ci)
+        stub_feature_flags(use_model_load_balancing: ff_use_model_load_balancing)
+      end
+
+      it 'results match expectations' do
+        result = models.transform_values do |model|
+          # Make load balancer to force init with a dedicated replicas connections
+          described_class.new(model).tap do |subject|
+            subject.configuration.hosts = [subject.configuration.replica_db_config.host]
+            subject.setup
+          end
+
+          load_balancer = model.connection.load_balancer
+
+          {
+            read: load_balancer.read { |connection| connection.pool.db_config.name },
+            write: load_balancer.read_write { |connection| connection.pool.db_config.name }
+          }
+        end
+
+        expect(result).to eq(expectations)
+      end
+    end
+  end
 end