Enable caching for lb per request

- Temporary ensure caching under the feature flag
- Fix specs for load_balancer

ee/config/feature_flags/development/query_cache_for_load_balancing.yml

+---
+name: query_cache_for_load_balancing
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/46765
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/276203
+milestone: '13.7'
+type: development
+group: group::memory
+default_enabled: false

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

@@ -7,7 +7,7 @@ module Gitlab
       class Host
         attr_reader :pool, :last_checked_at, :intervals, :load_balancer, :host, :port
 
-        delegate :connection, :release_connection, to: :pool
+        delegate :connection, :release_connection, :enable_query_cache!, :disable_query_cache!, to: :pool
 
         CONNECTION_ERRORS =
           if defined?(PG)

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

@@ -12,6 +12,7 @@ module Gitlab
       # always returns a connection to the primary.
       class LoadBalancer
         CACHE_KEY = :gitlab_load_balancer_host
+        ENSURE_CACHING_KEY = 'ensure_caching'
 
         attr_reader :host_list
 
@@ -28,6 +29,8 @@ module Gitlab
           conflict_retried = 0
 
           while host
+            ensure_caching!
+
             begin
               return yield host.connection
             rescue => error
@@ -95,7 +98,12 @@ module Gitlab
 
         # Releases the host and connection for the current thread.
         def release_host
-          RequestStore[CACHE_KEY]&.release_connection
+          if host = RequestStore[CACHE_KEY]
+            host.disable_query_cache!
+            host.release_connection
+          end
+
+          RequestStore.delete(ENSURE_CACHING_KEY)
           RequestStore.delete(CACHE_KEY)
         end
 
@@ -169,6 +177,23 @@ module Gitlab
             error.is_a?(PG::TRSerializationFailure)
           end
         end
+
+        private
+
+        # TODO:
+        # Move enable_query_cache! to ConnectionPool (https://gitlab.com/gitlab-org/gitlab/-/blob/master/lib/gitlab/database.rb#L223)
+        # when the feature flag is removed in https://gitlab.com/gitlab-org/gitlab/-/issues/276203.
+        def ensure_caching!
+          # Feature (Flipper gem) reads the data from the database, and it would cause the infinite loop here.
+          # We need to ensure that the code below is executed only once, until the feature flag is removed.
+          return if RequestStore[ENSURE_CACHING_KEY]
+
+          RequestStore[ENSURE_CACHING_KEY] = true
+
+          if Feature.enabled?(:query_cache_for_load_balancing)
+            host.enable_query_cache!
+          end
+        end
       end
     end
   end

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

@@ -2,16 +2,15 @@
 
 require 'spec_helper'
 
-RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer do
+RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer, :request_store do
+  let(:pool_spec) { ActiveRecord::Base.connection_pool.spec }
+  let(:pool) { ActiveRecord::ConnectionAdapters::ConnectionPool.new(pool_spec) }
+
   let(:lb) { described_class.new(%w(localhost localhost)) }
 
   before do
     allow(Gitlab::Database).to receive(:create_connection_pool)
-      .and_return(ActiveRecord::Base.connection_pool)
-  end
-
-  after do
-    RequestStore.delete(described_class::CACHE_KEY)
+      .and_return(pool)
   end
 
   def raise_and_wrap(wrapper, original)
@@ -52,8 +51,28 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer do
 
       allow(lb).to receive(:host).and_return(host)
       expect(host).to receive(:connection).and_return(connection)
+      expect(host).to receive(:enable_query_cache!).once
 
       expect { |b| lb.read(&b) }.to yield_with_args(connection)
+
+      expect(RequestStore[described_class::ENSURE_CACHING_KEY]).to be true
+    end
+
+    context 'when :query_cache_for_load_balancing feature flag is disabled' do
+      before do
+        stub_feature_flags(query_cache_for_load_balancing: false)
+      end
+
+      it 'yields a connection for a read without enabling query cache' do
+        connection = double(:connection)
+        host = double(:host)
+
+        allow(lb).to receive(:host).and_return(host)
+        expect(host).to receive(:connection).and_return(connection)
+        expect(host).not_to receive(:enable_query_cache!)
+
+        expect { |b| lb.read(&b) }.to yield_with_args(connection)
+      end
     end
 
     it 'marks hosts that are offline' do
@@ -142,10 +161,14 @@ RSpec.describe Gitlab::Database::LoadBalancing::LoadBalancer do
 
   describe '#release_host' do
     it 'releases the host and its connection' do
-      lb.host
+      host = lb.host
+
+      expect(host).to receive(:disable_query_cache!)
+
       lb.release_host
 
       expect(RequestStore[described_class::CACHE_KEY]).to be_nil
+      expect(RequestStore[described_class::ENSURE_CACHING_KEY]).to be_nil
     end
   end