Run serivce discovery on load balancing configuration

This solves an issue where load balancing would be configured but then
queries would run before service discovery ran. This could lead to
read queries falling back to the database primary.

See https://gitlab.com/gitlab-org/gitlab/-/issues/323726

Changelog: fixed

lib/gitlab/database/load_balancing.rb

@@ -108,6 +108,11 @@ module Gitlab
       # Configures proxying of requests.
       def self.configure_proxy(proxy = ConnectionProxy.new(hosts))
         ActiveRecord::Base.load_balancing_proxy = proxy
+
+        # Populate service discovery immediately if it is configured
+        if service_discovery_enabled?
+          ServiceDiscovery.new(service_discovery_configuration).perform_service_discovery
+        end
       end
 
       def self.active_record_models

lib/gitlab/database/load_balancing/service_discovery.rb

@@ -63,18 +63,14 @@ module Gitlab
         end
 
         def start
-          # We run service discovery once in the current thread so that the application's main thread
-          # does not race this thread to use the results of initial service discovery.
-          next_sleep_duration = perform_service_discovery
-
           Thread.new do
             loop do
+              next_sleep_duration = perform_service_discovery
+
               # We slightly randomize the sleep() interval. This should reduce
               # the likelihood of _all_ processes refreshing at the same time,
               # possibly putting unnecessary pressure on the DNS server.
               sleep(next_sleep_duration + rand(MAX_SLEEP_ADJUSTMENT))
-
-              next_sleep_duration = perform_service_discovery
             end
           end
         end

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

@@ -57,16 +57,13 @@ RSpec.describe Gitlab::Database::LoadBalancing::ServiceDiscovery do
         .and_yield
     end
 
-    it 'runs service discovery once before starting the worker thread' do
-      expect(service).to receive(:perform_service_discovery).ordered.and_return(5)
-
+    it 'starts service discovery in a new thread' do
       expect(Thread).to receive(:new).ordered.and_call_original # Thread starts
 
+      expect(service).to receive(:perform_service_discovery).ordered.and_return(5)
       expect(service).to receive(:rand).ordered.and_return(2)
       expect(service).to receive(:sleep).ordered.with(7) # Sleep runs after thread starts
 
-      expect(service).to receive(:perform_service_discovery).ordered.and_return(1)
-
       service.start.join
     end
   end

spec/lib/gitlab/database/load_balancing_spec.rb

@@ -212,6 +212,21 @@ RSpec.describe Gitlab::Database::LoadBalancing do
       expect(ActiveRecord::Base).to have_received(:load_balancing_proxy=)
         .with(Gitlab::Database::LoadBalancing::ConnectionProxy)
     end
+
+    context 'when service discovery is enabled' do
+      let(:service_discovery) { double(Gitlab::Database::LoadBalancing::ServiceDiscovery) }
+
+      it 'runs initial service discovery when configuring the connection proxy' do
+        allow(described_class)
+          .to receive(:configuration)
+          .and_return('discover' => { 'record' => 'foo' })
+
+        expect(Gitlab::Database::LoadBalancing::ServiceDiscovery).to receive(:new).and_return(service_discovery)
+        expect(service_discovery).to receive(:perform_service_discovery)
+
+        described_class.configure_proxy
+      end
+    end
   end
 
   describe '.active_record_models' do