Merge branch 'sh-port-ee-load-balancing-ce' into 'master'

[RUN ALL RSPEC] Move core EE load balancing code to CE

See merge request gitlab-org/gitlab!60894

.rubocop_manual_todo.yml

@@ -761,7 +761,6 @@ RSpec/TimecopFreeze:
     - 'ee/spec/lib/gitlab/analytics/cycle_analytics/summary/group/stage_time_summary_spec.rb'
     - 'ee/spec/lib/gitlab/analytics/type_of_work/tasks_by_type_spec.rb'
     - 'ee/spec/lib/gitlab/auth/group_saml/sso_enforcer_spec.rb'
-    - 'ee/spec/lib/gitlab/database/load_balancing/host_spec.rb'
     - 'ee/spec/lib/gitlab/geo/base_request_spec.rb'
     - 'ee/spec/lib/gitlab/geo/event_gap_tracking_spec.rb'
     - 'ee/spec/lib/gitlab/geo/git_push_http_spec.rb'
@@ -2900,7 +2899,6 @@ Style/RegexpLiteralMixedPreserve:
     - 'ee/spec/controllers/groups/groups_controller_spec.rb'
     - 'ee/spec/features/groups/saml_enforcement_spec.rb'
     - 'ee/spec/features/markdown/metrics_spec.rb'
-    - 'ee/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb'
     - 'ee/spec/services/jira/requests/issues/list_service_spec.rb'
     - 'lib/api/invitations.rb'
     - 'lib/gitlab/ci/pipeline/expression/lexeme/pattern.rb'

app/models/application_record.rb

@@ -53,12 +53,14 @@ class ApplicationRecord < ActiveRecord::Base
   # Start a new transaction with a shorter-than-usual statement timeout. This is
   # currently one third of the default 15-second timeout
   def self.with_fast_read_statement_timeout(timeout_ms = 5000)
+    ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
       transaction(requires_new: true) do
         connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
 
         yield
       end
     end
+  end
 
   def self.safe_find_or_create_by(*args, &block)
     safe_ensure_unique(retries: 1) do
@@ -85,5 +87,3 @@ class ApplicationRecord < ActiveRecord::Base
     enum(enum_mod.key => values)
   end
 end
-
-ApplicationRecord.prepend_mod_with('ApplicationRecordHelpers')

config/initializers/load_balancing.rb

 # frozen_string_literal: true
 
-# We need to run this initializer after migrations are done so it doesn't fail on CI
-
-Gitlab.ee do
-  if Gitlab::Database.cached_table_exists?('licenses')
-    if Gitlab::Database::LoadBalancing.enable?
+if Gitlab::Database::LoadBalancing.enable?
   Gitlab::Database.disable_prepared_statements
 
   Gitlab::Application.configure do |config|
@@ -19,7 +15,4 @@ Gitlab.ee do
     # discovery depends on this.
     Gitlab::Database::LoadBalancing.start_service_discovery
   end
-
-    end
-  end
 end

ee/app/models/ee/application_record_helpers.rb

-# frozen_string_literal: true
-
-module EE
-  # Intentionally pick a different name, to prevent naming conflicts
-  module ApplicationRecordHelpers
-    extend ActiveSupport::Concern
-
-    class_methods do
-      extend ::Gitlab::Utils::Override
-
-      override :with_fast_read_statement_timeout
-      def with_fast_read_statement_timeout(timeout_ms = 5000)
-        ::Gitlab::Database::LoadBalancing::Session.current.fallback_to_replicas_for_ambiguous_queries do
-          super
-        end
-      end
-    end
-  end
-end

ee/lib/ee/gitlab/database.rb

@@ -17,12 +17,6 @@ module EE
           !Postgresql::ReplicationSlot.lag_too_great?
         end
 
-        # Disables prepared statements for the current database connection.
-        def disable_prepared_statements
-          config = ::Gitlab::Database.config.merge(prepared_statements: false)
-          ActiveRecord::Base.establish_connection(config)
-        end
-
         def geo_uncached_queries(&block)
           raise 'No block given' unless block_given?
 

ee/lib/ee/gitlab/database/consistency.rb

-# frozen_string_literal: true
-
-module EE
-  module Gitlab
-    module Database
-      module Consistency
-        extend ::Gitlab::Utils::Override
-        ##
-        # In EE we are disabling the database load balancing for calls that
-        # require read consistency after recent writes.
-        #
-        override :with_read_consistency
-        def with_read_consistency(&block)
-          ::Gitlab::Database::LoadBalancing::Session
-            .current.use_primary(&block)
-        end
-      end
-    end
-  end
-end

ee/spec/lib/ee/gitlab/database_spec.rb

@@ -60,28 +60,6 @@ RSpec.describe Gitlab::Database do
     end
   end
 
-  describe '.disable_prepared_statements' do
-    around do |example|
-      original_config = ::Gitlab::Database.config
-
-      example.run
-
-      ActiveRecord::Base.establish_connection(original_config)
-    end
-
-    it 'disables prepared statements' do
-      ActiveRecord::Base.establish_connection(::Gitlab::Database.config.merge(prepared_statements: true))
-      expect(ActiveRecord::Base.connection.prepared_statements).to eq(true)
-
-      expect(ActiveRecord::Base).to receive(:establish_connection)
-        .with(a_hash_including({ 'prepared_statements' => false })).and_call_original
-
-      described_class.disable_prepared_statements
-
-      expect(ActiveRecord::Base.connection.prepared_statements).to eq(false)
-    end
-  end
-
   describe '.geo_uncached_queries' do
     context 'when no block is given' do
       it 'raises error' do

ee/spec/models/ee/application_record_helpers_spec.rb

-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe ApplicationRecord do
-  describe '.with_fast_read_statement_timeout' do
-    let(:session) { double(:session) }
-
-    before do
-      allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
-      allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries).and_yield
-    end
-
-    it 'yields control' do
-      expect do |blk|
-        described_class.with_fast_read_statement_timeout(&blk)
-      end.to yield_control.once
-    end
-
-    context 'when the query runs faster than configured timeout' do
-      it 'executes the query without error' do
-        result = nil
-
-        expect do
-          described_class.with_fast_read_statement_timeout(100) do
-            result = described_class.connection.exec_query('SELECT 1')
-          end
-        end.not_to raise_error
-
-        expect(result).not_to be_nil
-      end
-    end
-
-    # This query hangs for 10ms and then gets cancelled.  As there is no
-    # other way to test the timeout for sure, 10ms of waiting seems to be
-    # reasonable!
-    context 'when the query runs longer than configured timeout' do
-      it 'cancels the query and raiss an exception' do
-        expect do
-          described_class.with_fast_read_statement_timeout(10) do
-            described_class.connection.exec_query('SELECT pg_sleep(0.1)')
-          end
-        end.to raise_error(ActiveRecord::QueryCanceled)
-      end
-    end
-  end
-end

ee/spec/tasks/geo_rake_spec.rb

@@ -428,8 +428,7 @@ RSpec.describe 'geo rake tasks', :geo do
     end
 
     it 'removes orphaned registries taking into account TO_PROJECT_ID' do
-      allow(ENV).to receive(:[]).with('FROM_PROJECT_ID').and_return(nil)
-      allow(ENV).to receive(:[]).with('TO_PROJECT_ID').and_return(@orphaned.project_id)
+      stub_env('FROM_PROJECT_ID' => nil, 'TO_PROJECT_ID' => @orphaned.project_id)
 
       run_rake_task('geo:run_orphaned_project_registry_cleaner')
 
@@ -439,8 +438,7 @@ RSpec.describe 'geo rake tasks', :geo do
     end
 
     it 'removes orphaned registries taking into account FROM_PROJECT_ID' do
-      allow(ENV).to receive(:[]).with('FROM_PROJECT_ID').and_return(@orphaned1.project_id)
-      allow(ENV).to receive(:[]).with('TO_PROJECT_ID').and_return(nil)
+      stub_env('FROM_PROJECT_ID' => @orphaned1.project_id, 'TO_PROJECT_ID' => nil)
 
       run_rake_task('geo:run_orphaned_project_registry_cleaner')
 

lib/gitlab/database.rb

@@ -89,6 +89,11 @@ module Gitlab
       end
     end
 
+    # Disables prepared statements for the current database connection.
+    def self.disable_prepared_statements
+      ActiveRecord::Base.establish_connection(config.merge(prepared_statements: false))
+    end
+
     # @deprecated
     def self.postgresql?
       adapter_name.casecmp('postgresql') == 0

lib/gitlab/database/consistency.rb

@@ -4,28 +4,18 @@ module Gitlab
   module Database
     ##
     # This class is used to make it possible to ensure read consistency in
-    # GitLab EE without the need of overriding a lot of methods / classes /
+    # GitLab without the need of overriding a lot of methods / classes /
     # classs.
     #
-    # This is a CE class that does nothing in CE, because database load
-    # balancing is EE-only feature, but you can still use it in CE. It will
-    # start ensuring read consistency once it is overridden in EE.
-    #
-    # Using this class in CE helps to avoid creeping discrepancy between CE /
-    # EE only to force usage of the primary database in EE.
-    #
     class Consistency
       ##
-      # In CE there is no database load balancing, so all reads are expected to
-      # be consistent by the ACID guarantees of a single PostgreSQL instance.
-      #
-      # This method is overridden in EE.
+      # Within the block, disable the database load balancing for calls that
+      # require read consistency after recent writes.
       #
       def self.with_read_consistency(&block)
-        yield
+        ::Gitlab::Database::LoadBalancing::Session
+          .current.use_primary(&block)
       end
     end
   end
 end
-
-::Gitlab::Database::Consistency.singleton_class.prepend_mod_with('Gitlab::Database::Consistency')

ee/lib/gitlab/database/load_balancing.rb → lib/gitlab/database/load_balancing.rb

@@ -102,23 +102,9 @@ module Gitlab
         hosts.any? || service_discovery_enabled?
       end
 
+      # Temporarily disabled for FOSS until move from EE to FOSS is complete
       def self.feature_available?
-        # If this method is called in any subscribers listening to
-        # sql.active_record, the SQL call below may cause infinite recursion.
-        # So, the memoization variable must have 3 states
-        # - First call: @feature_available is undefined
-        #   -> Set @feature_available to false
-        #   -> Trigger SQL
-        #   -> SQL subscriber triggers this method again
-        #     -> return false
-        #   -> Set @feature_available  to true
-        #   -> return true
-        # - Second call: return @feature_available right away
-        return @feature_available if defined?(@feature_available)
-
-        @feature_available = false
-        @feature_available = Gitlab::Database.cached_table_exists?('licenses') &&
-                             ::License.feature_available?(:db_load_balancing)
+        Gitlab.ee? || Gitlab::Utils.to_boolean(ENV['ENABLE_LOAD_BALANCING_FOR_FOSS'], default: false)
       end
 
       def self.start_service_discovery

scripts/rspec_helpers.sh

@@ -85,7 +85,7 @@ function rspec_db_library_code() {
   local db_files="spec/lib/gitlab/database/ spec/support/helpers/database/"
 
   if [[ -d "ee/" ]]; then
-    db_files="${db_files} ee/spec/lib/gitlab/database/ ee/spec/lib/ee/gitlab/database_spec.rb"
+    db_files="${db_files} ee/spec/lib/ee/gitlab/database_spec.rb"
   fi
 
   rspec_simple_job "-- ${db_files}"

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

@@ -294,7 +294,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
 
   describe '#primary_write_location' do
     it 'returns a String in the right format' do
-      expect(lb.primary_write_location).to match(/[A-F0-9]{1,8}\/[A-F0-9]{1,8}/)
+      expect(lb.primary_write_location).to match(%r{[A-F0-9]{1,8}/[A-F0-9]{1,8}})
     end
 
     it 'raises an error if the write location could not be retrieved' do

ee/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb → spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb

@@ -52,7 +52,7 @@ RSpec.describe Gitlab::Database::LoadBalancing::SrvResolver do
   end
 
   def dns_response_packet_from_fixture(fixture_name)
-    fixture         = File.read(Rails.root + "ee/spec/fixtures/dns/#{fixture_name}.json")
+    fixture         = File.read(Rails.root + "spec/fixtures/dns/#{fixture_name}.json")
     encoded_payload = Gitlab::Json.parse(fixture)['payload']
     payload         = Base64.decode64(encoded_payload)
 

ee/spec/lib/gitlab/database/load_balancing_spec.rb → spec/lib/gitlab/database/load_balancing_spec.rb

@@ -5,6 +5,10 @@ require 'spec_helper'
 RSpec.describe Gitlab::Database::LoadBalancing do
   include_context 'clear DB Load Balancing configuration'
 
+  before do
+    stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'true')
+  end
+
   describe '.proxy' do
     context 'when configured' do
       before do
@@ -127,8 +131,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.enable?' do
-    let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
-
     before do
       clear_load_balancing_configuration
       allow(described_class).to receive(:hosts).and_return(%w(foo))
@@ -181,27 +183,22 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
     end
 
-    context 'without a license' do
+    context 'FOSS' do
       before do
-        License.destroy_all # rubocop: disable Cop/DestroyAll
-        clear_load_balancing_configuration
-      end
+        allow(Gitlab).to receive(:ee?).and_return(false)
 
-      it 'is disabled' do
-        expect(described_class.enable?).to eq(false)
-      end
+        stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'false')
       end
 
-    context 'with an EES license' do
-      let!(:license) { create(:license, plan: ::License::STARTER_PLAN) }
-
       it 'is disabled' do
         expect(described_class.enable?).to eq(false)
       end
     end
 
-    context 'with an EEP license' do
-      let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
+    context 'EE' do
+      before do
+        allow(Gitlab).to receive(:ee?).and_return(true)
+      end
 
       it 'is enabled' do
         allow(described_class).to receive(:hosts).and_return(%w(foo))
@@ -213,8 +210,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
   end
 
   describe '.configured?' do
-    let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
-
     before do
       clear_load_balancing_configuration
     end
@@ -245,17 +240,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
 
       expect(described_class.configured?).to eq(false)
     end
-
-    context 'without a license' do
-      before do
-        License.destroy_all # rubocop: disable Cop/DestroyAll
-        clear_load_balancing_configuration
-      end
-
-      it 'is not configured' do
-        expect(described_class.configured?).to eq(false)
-      end
-    end
   end
 
   describe '.configure_proxy' do
@@ -444,7 +428,6 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       end
 
       before do
-        stub_licensed_features(db_load_balancing: true)
         # Preloading testing class
         model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
 

spec/lib/gitlab/database_spec.rb

@@ -65,6 +65,28 @@ RSpec.describe Gitlab::Database do
     end
   end
 
+  describe '.disable_prepared_statements' do
+    around do |example|
+      original_config = ::Gitlab::Database.config
+
+      example.run
+
+      ActiveRecord::Base.establish_connection(original_config)
+    end
+
+    it 'disables prepared statements' do
+      ActiveRecord::Base.establish_connection(::Gitlab::Database.config.merge(prepared_statements: true))
+      expect(ActiveRecord::Base.connection.prepared_statements).to eq(true)
+
+      expect(ActiveRecord::Base).to receive(:establish_connection)
+        .with(a_hash_including({ 'prepared_statements' => false })).and_call_original
+
+      described_class.disable_prepared_statements
+
+      expect(ActiveRecord::Base.connection.prepared_statements).to eq(false)
+    end
+  end
+
   describe '.postgresql?' do
     subject { described_class.postgresql? }
 

spec/models/application_record_spec.rb

@@ -132,5 +132,47 @@ RSpec.describe ApplicationRecord do
         end.to raise_error(ActiveRecord::QueryCanceled)
       end
     end
+
+    context 'with database load balancing' do
+      let(:session) { double(:session) }
+
+      before do
+        allow(::Gitlab::Database::LoadBalancing::Session).to receive(:current).and_return(session)
+        allow(session).to receive(:fallback_to_replicas_for_ambiguous_queries).and_yield
+      end
+
+      it 'yields control' do
+        expect do |blk|
+          described_class.with_fast_read_statement_timeout(&blk)
+        end.to yield_control.once
+      end
+
+      context 'when the query runs faster than configured timeout' do
+        it 'executes the query without error' do
+          result = nil
+
+          expect do
+            described_class.with_fast_read_statement_timeout(100) do
+              result = described_class.connection.exec_query('SELECT 1')
+            end
+          end.not_to raise_error
+
+          expect(result).not_to be_nil
+        end
+      end
+
+      # This query hangs for 10ms and then gets cancelled.  As there is no
+      # other way to test the timeout for sure, 10ms of waiting seems to be
+      # reasonable!
+      context 'when the query runs longer than configured timeout' do
+        it 'cancels the query and raiss an exception' do
+          expect do
+            described_class.with_fast_read_statement_timeout(10) do
+              described_class.connection.exec_query('SELECT pg_sleep(0.1)')
+            end
+          end.to raise_error(ActiveRecord::QueryCanceled)
+        end
+      end
+    end
   end
 end