Commit a55a37fb authored by Nikola Milojevic's avatar Nikola Milojevic 2️⃣ Committed by Kamil Trzciński
Browse files

Enable caching for lb per request 

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

ee/config/feature_flags/development/query_cache_for_load_balancing.yml

0 → 100644
+8 −0
Original line number Diff line number Diff line 
---

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

+1 −1
Original line number Diff line number Diff line
@@ -7,7 +7,7 @@ module LoadBalancing 
      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

+26 −1
Original line number Diff line number Diff line
@@ -12,6 +12,7 @@ module LoadBalancing 
      # 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 @@ def read(&block) 
          conflict_retried = 0



          while host

            ensure_caching!



            begin

              return yield host.connection

            rescue => error
@@ -95,7 +98,12 @@ def host 


        # 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 @@ def serialization_failure?(error) 
            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

+30 −7
Original line number Diff line number Diff line
@@ -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 @@ def twice_wrapped_exception(top, middle, original) 


      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 @@ def twice_wrapped_exception(top, middle, original) 


  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