From a27800deb28f5fdb6cd12b8e0d4af2e483187122 Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Mon, 24 May 2021 13:42:15 +1200
Subject: [PATCH 1/8] Allow reads to DB_POOL_HEADROOM to proceed normally

---
 ee/spec/tasks/geo_rake_spec.rb | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/ee/spec/tasks/geo_rake_spec.rb b/ee/spec/tasks/geo_rake_spec.rb
index de6669365afee7..aa7e77b17dae9f 100644
--- a/ee/spec/tasks/geo_rake_spec.rb
+++ b/ee/spec/tasks/geo_rake_spec.rb
@@ -428,8 +428,7 @@
     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 @@
     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')
 
-- 
GitLab


From 0f4e48307ea97909526d68144d1eb2fb67f74ad0 Mon Sep 17 00:00:00 2001
From: Stan Hu <stanhu@gmail.com>
Date: Mon, 3 May 2021 23:49:35 -0700
Subject: [PATCH 2/8] Move all core EE load balancing code to CE

This moves all files from:

ee/lib/gitlab/database/load_balancing* -> lib/gitlab/database
ee/spec/lib/gitlab/database/load_balancing* -> spec/lib/gitlab/database
---
 config/initializers/load_balancing.rb         | 28 +++++++++----------
 .../gitlab/database/load_balancing.rb         |  0
 .../load_balancing/active_record_proxy.rb     |  0
 .../load_balancing/connection_proxy.rb        |  0
 .../gitlab/database/load_balancing/host.rb    |  0
 .../database/load_balancing/host_list.rb      |  0
 .../database/load_balancing/load_balancer.rb  |  0
 .../gitlab/database/load_balancing/logger.rb  |  0
 .../load_balancing/rack_middleware.rb         |  0
 .../database/load_balancing/resolver.rb       |  0
 .../load_balancing/service_discovery.rb       |  0
 .../gitlab/database/load_balancing/session.rb |  0
 .../sidekiq_client_middleware.rb              |  0
 .../sidekiq_server_middleware.rb              |  0
 .../database/load_balancing/srv_resolver.rb   |  0
 .../database/load_balancing/sticking.rb       |  0
 {ee/spec => spec}/fixtures/dns/a_rr.json      |  0
 .../a_with_aaaa_rr_in_additional_section.json |  0
 {ee/spec => spec}/fixtures/dns/aaaa_rr.json   |  0
 .../srv_with_a_rr_in_additional_section.json  |  0
 .../active_record_proxy_spec.rb               |  0
 .../load_balancing/connection_proxy_spec.rb   |  0
 .../database/load_balancing/host_list_spec.rb |  0
 .../database/load_balancing/host_spec.rb      |  0
 .../load_balancing/load_balancer_spec.rb      |  0
 .../load_balancing/rack_middleware_spec.rb    |  0
 .../database/load_balancing/resolver_spec.rb  |  0
 .../load_balancing/service_discovery_spec.rb  |  0
 .../database/load_balancing/session_spec.rb   |  0
 .../sidekiq_client_middleware_spec.rb         |  0
 .../sidekiq_server_middleware_spec.rb         |  0
 .../load_balancing/srv_resolver_spec.rb       |  2 +-
 .../database/load_balancing/sticking_spec.rb  |  0
 .../gitlab/database/load_balancing_spec.rb    |  0
 ..._balancing_configuration_shared_context.rb |  0
 35 files changed, 14 insertions(+), 16 deletions(-)
 rename {ee/lib => lib}/gitlab/database/load_balancing.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/active_record_proxy.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/connection_proxy.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/host.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/host_list.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/load_balancer.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/logger.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/rack_middleware.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/resolver.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/service_discovery.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/session.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/sidekiq_client_middleware.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/sidekiq_server_middleware.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/srv_resolver.rb (100%)
 rename {ee/lib => lib}/gitlab/database/load_balancing/sticking.rb (100%)
 rename {ee/spec => spec}/fixtures/dns/a_rr.json (100%)
 rename {ee/spec => spec}/fixtures/dns/a_with_aaaa_rr_in_additional_section.json (100%)
 rename {ee/spec => spec}/fixtures/dns/aaaa_rr.json (100%)
 rename {ee/spec => spec}/fixtures/dns/srv_with_a_rr_in_additional_section.json (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/connection_proxy_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/host_list_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/host_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/load_balancer_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/rack_middleware_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/resolver_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/service_discovery_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/session_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/srv_resolver_spec.rb (95%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing/sticking_spec.rb (100%)
 rename {ee/spec => spec}/lib/gitlab/database/load_balancing_spec.rb (100%)
 rename {ee/spec => spec}/support/shared_contexts/load_balancing_configuration_shared_context.rb (100%)

diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index 7502a6299aeae9..e3884953edd540 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -2,24 +2,22 @@
 
 # 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?
-      Gitlab::Database.disable_prepared_statements
+if Gitlab::Database.cached_table_exists?('licenses')
+  if Gitlab::Database::LoadBalancing.enable?
+    Gitlab::Database.disable_prepared_statements
 
-      Gitlab::Application.configure do |config|
-        config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
-      end
-
-      Gitlab::Database::LoadBalancing.configure_proxy
+    Gitlab::Application.configure do |config|
+      config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+    end
 
-      # This needs to be executed after fork of clustered processes
-      Gitlab::Cluster::LifecycleEvents.on_worker_start do
-        # Service discovery must be started after configuring the proxy, as service
-        # discovery depends on this.
-        Gitlab::Database::LoadBalancing.start_service_discovery
-      end
+    Gitlab::Database::LoadBalancing.configure_proxy
 
+    # This needs to be executed after fork of clustered processes
+    Gitlab::Cluster::LifecycleEvents.on_worker_start do
+      # Service discovery must be started after configuring the proxy, as service
+      # discovery depends on this.
+      Gitlab::Database::LoadBalancing.start_service_discovery
     end
+
   end
 end
diff --git a/ee/lib/gitlab/database/load_balancing.rb b/lib/gitlab/database/load_balancing.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing.rb
rename to lib/gitlab/database/load_balancing.rb
diff --git a/ee/lib/gitlab/database/load_balancing/active_record_proxy.rb b/lib/gitlab/database/load_balancing/active_record_proxy.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/active_record_proxy.rb
rename to lib/gitlab/database/load_balancing/active_record_proxy.rb
diff --git a/ee/lib/gitlab/database/load_balancing/connection_proxy.rb b/lib/gitlab/database/load_balancing/connection_proxy.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/connection_proxy.rb
rename to lib/gitlab/database/load_balancing/connection_proxy.rb
diff --git a/ee/lib/gitlab/database/load_balancing/host.rb b/lib/gitlab/database/load_balancing/host.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/host.rb
rename to lib/gitlab/database/load_balancing/host.rb
diff --git a/ee/lib/gitlab/database/load_balancing/host_list.rb b/lib/gitlab/database/load_balancing/host_list.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/host_list.rb
rename to lib/gitlab/database/load_balancing/host_list.rb
diff --git a/ee/lib/gitlab/database/load_balancing/load_balancer.rb b/lib/gitlab/database/load_balancing/load_balancer.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/load_balancer.rb
rename to lib/gitlab/database/load_balancing/load_balancer.rb
diff --git a/ee/lib/gitlab/database/load_balancing/logger.rb b/lib/gitlab/database/load_balancing/logger.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/logger.rb
rename to lib/gitlab/database/load_balancing/logger.rb
diff --git a/ee/lib/gitlab/database/load_balancing/rack_middleware.rb b/lib/gitlab/database/load_balancing/rack_middleware.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/rack_middleware.rb
rename to lib/gitlab/database/load_balancing/rack_middleware.rb
diff --git a/ee/lib/gitlab/database/load_balancing/resolver.rb b/lib/gitlab/database/load_balancing/resolver.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/resolver.rb
rename to lib/gitlab/database/load_balancing/resolver.rb
diff --git a/ee/lib/gitlab/database/load_balancing/service_discovery.rb b/lib/gitlab/database/load_balancing/service_discovery.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/service_discovery.rb
rename to lib/gitlab/database/load_balancing/service_discovery.rb
diff --git a/ee/lib/gitlab/database/load_balancing/session.rb b/lib/gitlab/database/load_balancing/session.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/session.rb
rename to lib/gitlab/database/load_balancing/session.rb
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
rename to lib/gitlab/database/load_balancing/sidekiq_client_middleware.rb
diff --git a/ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb b/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
rename to lib/gitlab/database/load_balancing/sidekiq_server_middleware.rb
diff --git a/ee/lib/gitlab/database/load_balancing/srv_resolver.rb b/lib/gitlab/database/load_balancing/srv_resolver.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/srv_resolver.rb
rename to lib/gitlab/database/load_balancing/srv_resolver.rb
diff --git a/ee/lib/gitlab/database/load_balancing/sticking.rb b/lib/gitlab/database/load_balancing/sticking.rb
similarity index 100%
rename from ee/lib/gitlab/database/load_balancing/sticking.rb
rename to lib/gitlab/database/load_balancing/sticking.rb
diff --git a/ee/spec/fixtures/dns/a_rr.json b/spec/fixtures/dns/a_rr.json
similarity index 100%
rename from ee/spec/fixtures/dns/a_rr.json
rename to spec/fixtures/dns/a_rr.json
diff --git a/ee/spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json b/spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json
similarity index 100%
rename from ee/spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json
rename to spec/fixtures/dns/a_with_aaaa_rr_in_additional_section.json
diff --git a/ee/spec/fixtures/dns/aaaa_rr.json b/spec/fixtures/dns/aaaa_rr.json
similarity index 100%
rename from ee/spec/fixtures/dns/aaaa_rr.json
rename to spec/fixtures/dns/aaaa_rr.json
diff --git a/ee/spec/fixtures/dns/srv_with_a_rr_in_additional_section.json b/spec/fixtures/dns/srv_with_a_rr_in_additional_section.json
similarity index 100%
rename from ee/spec/fixtures/dns/srv_with_a_rr_in_additional_section.json
rename to spec/fixtures/dns/srv_with_a_rr_in_additional_section.json
diff --git a/ee/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb
rename to spec/lib/gitlab/database/load_balancing/active_record_proxy_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb b/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
rename to spec/lib/gitlab/database/load_balancing/connection_proxy_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/host_list_spec.rb b/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/host_list_spec.rb
rename to spec/lib/gitlab/database/load_balancing/host_list_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/host_spec.rb b/spec/lib/gitlab/database/load_balancing/host_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/host_spec.rb
rename to spec/lib/gitlab/database/load_balancing/host_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
rename to spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
rename to spec/lib/gitlab/database/load_balancing/rack_middleware_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/resolver_spec.rb
rename to spec/lib/gitlab/database/load_balancing/resolver_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb b/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
rename to spec/lib/gitlab/database/load_balancing/service_discovery_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/session_spec.rb b/spec/lib/gitlab/database/load_balancing/session_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/session_spec.rb
rename to spec/lib/gitlab/database/load_balancing/session_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
rename to spec/lib/gitlab/database/load_balancing/sidekiq_client_middleware_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb b/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
rename to spec/lib/gitlab/database/load_balancing/sidekiq_server_middleware_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb b/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
similarity index 95%
rename from ee/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
rename to spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
index 029dc426516db3..6ac0608d4851de 100644
--- a/ee/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/srv_resolver_spec.rb
@@ -52,7 +52,7 @@
   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)
 
diff --git a/ee/spec/lib/gitlab/database/load_balancing/sticking_spec.rb b/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing/sticking_spec.rb
rename to spec/lib/gitlab/database/load_balancing/sticking_spec.rb
diff --git a/ee/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/load_balancing_spec.rb
rename to spec/lib/gitlab/database/load_balancing_spec.rb
diff --git a/ee/spec/support/shared_contexts/load_balancing_configuration_shared_context.rb b/spec/support/shared_contexts/load_balancing_configuration_shared_context.rb
similarity index 100%
rename from ee/spec/support/shared_contexts/load_balancing_configuration_shared_context.rb
rename to spec/support/shared_contexts/load_balancing_configuration_shared_context.rb
-- 
GitLab


From b3161695edf3735552b98855cfb9171c932b8d4d Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Mon, 24 May 2021 10:45:10 +1200
Subject: [PATCH 3/8] LB feature is available in FOSS but is not ready

As we have some code in ee/ uses LB, in overriden methods
and may have the (incorrect) assumption that LB is not present in FOSS
---
 lib/gitlab/database/load_balancing.rb         | 18 +--------
 .../gitlab/database/load_balancing_spec.rb    | 39 ++++++-------------
 2 files changed, 13 insertions(+), 44 deletions(-)

diff --git a/lib/gitlab/database/load_balancing.rb b/lib/gitlab/database/load_balancing.rb
index 76e67e27f91190..ea4628ab757f54 100644
--- a/lib/gitlab/database/load_balancing.rb
+++ b/lib/gitlab/database/load_balancing.rb
@@ -102,23 +102,9 @@ def self.configured?
         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
diff --git a/spec/lib/gitlab/database/load_balancing_spec.rb b/spec/lib/gitlab/database/load_balancing_spec.rb
index fc318be42aa5e4..c3dcfa3eb4a542 100644
--- a/spec/lib/gitlab/database/load_balancing_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing_spec.rb
@@ -5,6 +5,10 @@
 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 @@
   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 @@
       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)
+        stub_env('ENABLE_LOAD_BALANCING_FOR_FOSS', 'false')
       end
-    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 @@
   end
 
   describe '.configured?' do
-    let!(:license) { create(:license, plan: ::License::PREMIUM_PLAN) }
-
     before do
       clear_load_balancing_configuration
     end
@@ -245,17 +240,6 @@
 
       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 @@
       end
 
       before do
-        stub_licensed_features(db_load_balancing: true)
         # Preloading testing class
         model.singleton_class.prepend ::Gitlab::Database::LoadBalancing::ActiveRecordProxy
 
-- 
GitLab


From 5b15b1b7655107d843b434d9a54bd1999ccdb60d Mon Sep 17 00:00:00 2001
From: Stan Hu <stanhu@gmail.com>
Date: Mon, 3 May 2021 23:58:22 -0700
Subject: [PATCH 4/8] Remove EE-specific ApplicationRecord code

This moves the database load balancing functionality into CE.
---
 app/models/application_record.rb              | 10 ++--
 .../models/ee/application_record_helpers.rb   | 19 --------
 .../ee/application_record_helpers_spec.rb     | 47 -------------------
 spec/models/application_record_spec.rb        | 42 +++++++++++++++++
 4 files changed, 47 insertions(+), 71 deletions(-)
 delete mode 100644 ee/app/models/ee/application_record_helpers.rb
 delete mode 100644 ee/spec/models/ee/application_record_helpers_spec.rb

diff --git a/app/models/application_record.rb b/app/models/application_record.rb
index 5e5bc00458efd0..a93348a3b270c5 100644
--- a/app/models/application_record.rb
+++ b/app/models/application_record.rb
@@ -53,10 +53,12 @@ def self.safe_find_or_create_by!(*args, &block)
   # 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)
-    transaction(requires_new: true) do
-      connection.exec_query("SET LOCAL statement_timeout = #{timeout_ms}")
+    ::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
+        yield
+      end
     end
   end
 
@@ -85,5 +87,3 @@ def self.declarative_enum(enum_mod)
     enum(enum_mod.key => values)
   end
 end
-
-ApplicationRecord.prepend_mod_with('ApplicationRecordHelpers')
diff --git a/ee/app/models/ee/application_record_helpers.rb b/ee/app/models/ee/application_record_helpers.rb
deleted file mode 100644
index d023935d52e869..00000000000000
--- a/ee/app/models/ee/application_record_helpers.rb
+++ /dev/null
@@ -1,19 +0,0 @@
-# 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
diff --git a/ee/spec/models/ee/application_record_helpers_spec.rb b/ee/spec/models/ee/application_record_helpers_spec.rb
deleted file mode 100644
index 1a2c29de5f57d3..00000000000000
--- a/ee/spec/models/ee/application_record_helpers_spec.rb
+++ /dev/null
@@ -1,47 +0,0 @@
-# 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
diff --git a/spec/models/application_record_spec.rb b/spec/models/application_record_spec.rb
index 24de46cb536c96..85a6717d259b87 100644
--- a/spec/models/application_record_spec.rb
+++ b/spec/models/application_record_spec.rb
@@ -132,5 +132,47 @@
         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
-- 
GitLab


From 1387e6bfa1f16b87098e9d7f56e9513a91e74d7b Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Mon, 24 May 2021 13:34:40 +1200
Subject: [PATCH 5/8] Move LB Consistency to Core

---
 ee/lib/ee/gitlab/database/consistency.rb      | 20 -------------------
 lib/gitlab/database/consistency.rb            | 20 +++++--------------
 scripts/rspec_helpers.sh                      |  2 +-
 .../lib/gitlab/database/consistency_spec.rb   |  0
 4 files changed, 6 insertions(+), 36 deletions(-)
 delete mode 100644 ee/lib/ee/gitlab/database/consistency.rb
 rename {ee/spec => spec}/lib/gitlab/database/consistency_spec.rb (100%)

diff --git a/ee/lib/ee/gitlab/database/consistency.rb b/ee/lib/ee/gitlab/database/consistency.rb
deleted file mode 100644
index a3ccda7526f902..00000000000000
--- a/ee/lib/ee/gitlab/database/consistency.rb
+++ /dev/null
@@ -1,20 +0,0 @@
-# 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
diff --git a/lib/gitlab/database/consistency.rb b/lib/gitlab/database/consistency.rb
index e99ea7a323275c..17c16640e4c95b 100644
--- a/lib/gitlab/database/consistency.rb
+++ b/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')
diff --git a/scripts/rspec_helpers.sh b/scripts/rspec_helpers.sh
index b1a618270b0961..0484cabca82e0c 100644
--- a/scripts/rspec_helpers.sh
+++ b/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}"
diff --git a/ee/spec/lib/gitlab/database/consistency_spec.rb b/spec/lib/gitlab/database/consistency_spec.rb
similarity index 100%
rename from ee/spec/lib/gitlab/database/consistency_spec.rb
rename to spec/lib/gitlab/database/consistency_spec.rb
-- 
GitLab


From 2d72120fc293090451e9190bc23cb20094c6b544 Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Mon, 24 May 2021 14:23:24 +1200
Subject: [PATCH 6/8] Fix regexp cop that was previously excempted

---
 .rubocop_manual_todo.yml                                      | 2 --
 spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb | 2 +-
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/.rubocop_manual_todo.yml b/.rubocop_manual_todo.yml
index 9bdf1aa0082a60..0493dc47b86ecf 100644
--- a/.rubocop_manual_todo.yml
+++ b/.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'
@@ -2915,7 +2914,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'
diff --git a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
index a16d77a447df18..59f7016538053d 100644
--- a/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
+++ b/spec/lib/gitlab/database/load_balancing/load_balancer_spec.rb
@@ -294,7 +294,7 @@ def twice_wrapped_exception(top, middle, original)
 
   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
-- 
GitLab


From 7ce20dfb2bb4848538dd0c49aa60a27faf265c4c Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Tue, 25 May 2021 12:21:00 +1200
Subject: [PATCH 7/8] Do not check licenses table as we don't refer to License
 anymore

---
 config/initializers/load_balancing.rb | 27 +++++++++++----------------
 1 file changed, 11 insertions(+), 16 deletions(-)

diff --git a/config/initializers/load_balancing.rb b/config/initializers/load_balancing.rb
index e3884953edd540..a8a6e87179ae7c 100644
--- a/config/initializers/load_balancing.rb
+++ b/config/initializers/load_balancing.rb
@@ -1,23 +1,18 @@
 # frozen_string_literal: true
 
-# We need to run this initializer after migrations are done so it doesn't fail on CI
+if Gitlab::Database::LoadBalancing.enable?
+  Gitlab::Database.disable_prepared_statements
 
-if Gitlab::Database.cached_table_exists?('licenses')
-  if Gitlab::Database::LoadBalancing.enable?
-    Gitlab::Database.disable_prepared_statements
-
-    Gitlab::Application.configure do |config|
-      config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
-    end
-
-    Gitlab::Database::LoadBalancing.configure_proxy
+  Gitlab::Application.configure do |config|
+    config.middleware.use(Gitlab::Database::LoadBalancing::RackMiddleware)
+  end
 
-    # This needs to be executed after fork of clustered processes
-    Gitlab::Cluster::LifecycleEvents.on_worker_start do
-      # Service discovery must be started after configuring the proxy, as service
-      # discovery depends on this.
-      Gitlab::Database::LoadBalancing.start_service_discovery
-    end
+  Gitlab::Database::LoadBalancing.configure_proxy
 
+  # This needs to be executed after fork of clustered processes
+  Gitlab::Cluster::LifecycleEvents.on_worker_start do
+    # Service discovery must be started after configuring the proxy, as service
+    # discovery depends on this.
+    Gitlab::Database::LoadBalancing.start_service_discovery
   end
 end
-- 
GitLab


From c19df23717e3ee81c5125e205e42daa44e88db5f Mon Sep 17 00:00:00 2001
From: Thong Kuah <tkuah@gitlab.com>
Date: Tue, 25 May 2021 12:36:27 +1200
Subject: [PATCH 8/8] Move disable_prepared_statements to Core

We use this method in initializer for Load Balancing
---
 ee/lib/ee/gitlab/database.rb           |  6 ------
 ee/spec/lib/ee/gitlab/database_spec.rb | 22 ----------------------
 lib/gitlab/database.rb                 |  5 +++++
 spec/lib/gitlab/database_spec.rb       | 22 ++++++++++++++++++++++
 4 files changed, 27 insertions(+), 28 deletions(-)

diff --git a/ee/lib/ee/gitlab/database.rb b/ee/lib/ee/gitlab/database.rb
index 9ec5cc53f2b405..6d72783cba9f08 100644
--- a/ee/lib/ee/gitlab/database.rb
+++ b/ee/lib/ee/gitlab/database.rb
@@ -17,12 +17,6 @@ def healthy?
           !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?
 
diff --git a/ee/spec/lib/ee/gitlab/database_spec.rb b/ee/spec/lib/ee/gitlab/database_spec.rb
index 6c8cf098c0be33..bc2f9036b7fa3c 100644
--- a/ee/spec/lib/ee/gitlab/database_spec.rb
+++ b/ee/spec/lib/ee/gitlab/database_spec.rb
@@ -60,28 +60,6 @@
     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
diff --git a/lib/gitlab/database.rb b/lib/gitlab/database.rb
index 0305433aef4694..3583338c58a219 100644
--- a/lib/gitlab/database.rb
+++ b/lib/gitlab/database.rb
@@ -89,6 +89,11 @@ def self.human_adapter_name
       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
diff --git a/spec/lib/gitlab/database_spec.rb b/spec/lib/gitlab/database_spec.rb
index 663c8d69328b4c..2b31f3b4dee19e 100644
--- a/spec/lib/gitlab/database_spec.rb
+++ b/spec/lib/gitlab/database_spec.rb
@@ -65,6 +65,28 @@
     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? }
 
-- 
GitLab