Add bulk_stick to load balancing sticking

- Dry specs under the shared example
- Fix missing guards if lb is enabled.
- Use bulk_stick instrad of regular stick for each user
- Fix specs
- Fix changelog entry
- Add changelog entry for sticking improvement

ee/app/services/ee/user_project_access_changed_service.rb

@@ -5,9 +5,7 @@ module EE
     def execute(blocking: true, priority: ::UserProjectAccessChangedService::HIGH_PRIORITY)
       result = super
 
-      @user_ids.each do |id| # rubocop:disable Gitlab/ModuleWithInstanceVariables
-        ::Gitlab::Database::LoadBalancing::Sticking.stick(:user, id)
-      end
+      ::Gitlab::Database::LoadBalancing::Sticking.bulk_stick(:user, @user_ids) # rubocop:disable Gitlab/ModuleWithInstanceVariables
 
       result
     end

ee/changelogs/unreleased/254219-fix-n-1-cache-queries-for-projects-create.yml

+---
+title: Fix N+1 cache queries for load balancing sticking
+merge_request: 43843
+author:
+type: performance

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

@@ -53,7 +53,19 @@ module Gitlab
           Session.current.use_primary!
         end
 
-        def self.mark_primary_write_location(namespace, id)
+        def self.bulk_stick(namespace, ids)
+          return unless LoadBalancing.enable?
+
+          with_primary_write_location do |location|
+            ids.each do |id|
+              set_write_location_for(namespace, id, location)
+            end
+          end
+
+          Session.current.use_primary!
+        end
+
+        def self.with_primary_write_location
           return unless LoadBalancing.configured?
 
           # Load balancing could be enabled for the Web application server,
@@ -68,7 +80,13 @@ module Gitlab
 
           return if location.blank?
 
-          set_write_location_for(namespace, id, location)
+          yield(location)
+        end
+
+        def self.mark_primary_write_location(namespace, id)
+          with_primary_write_location do |location|
+            set_write_location_for(namespace, id, location)
+          end
         end
 
         # Stops sticking to the primary.

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

@@ -127,36 +127,53 @@ RSpec.describe Gitlab::Database::LoadBalancing::Sticking, :redis do
     end
   end
 
-  describe '.stick' do
+  RSpec.shared_examples 'sticking' do
     context 'when sticking is disabled' do
-      it 'does not perform any sticking' do
+      it 'does not perform any sticking', :aggregate_failures do
         expect(described_class).not_to receive(:set_write_location_for)
+        expect(Gitlab::Database::LoadBalancing::Session.current).not_to receive(:use_primary!)
 
-        described_class.stick(:user, 42)
+        described_class.bulk_stick(:user, ids)
       end
     end
 
     context 'when sticking is enabled' do
-      it 'sticks an entity to the primary' do
-        allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-          .and_return(true)
+      before do
         allow(Gitlab::Database::LoadBalancing).to receive(:configured?).and_return(true)
 
         lb = double(:lb, primary_write_location: 'foo')
 
         allow(described_class).to receive(:load_balancer).and_return(lb)
+      end
 
-        expect(described_class).to receive(:set_write_location_for)
-          .with(:user, 42, 'foo')
+      it 'sticks an entity to the primary', :aggregate_failures do
+        ids.each do |id|
+          expect(described_class).to receive(:set_write_location_for)
+                                       .with(:user, id, 'foo')
+        end
 
         expect(Gitlab::Database::LoadBalancing::Session.current)
           .to receive(:use_primary!)
 
-        described_class.stick(:user, 42)
+        subject
       end
     end
   end
 
+  describe '.stick' do
+    it_behaves_like 'sticking' do
+      let(:ids) { [42] }
+      subject { described_class.stick(:user, ids.first) }
+    end
+  end
+
+  describe '.bulk_stick' do
+    it_behaves_like 'sticking' do
+      let(:ids) { [42, 43] }
+      subject { described_class.bulk_stick(:user, ids) }
+    end
+  end
+
   describe '.mark_primary_write_location' do
     context 'when enabled' do
       before do

ee/spec/services/ee/user_project_access_changed_service_spec.rb

@@ -6,21 +6,33 @@ RSpec.describe EE::UserProjectAccessChangedService do
   let(:service) { UserProjectAccessChangedService.new([1, 2]) }
 
   describe '#execute' do
-    it 'sticks all the updated users and returns the original result' do
-      allow(Gitlab::Database::LoadBalancing).to receive(:enable?)
-        .and_return(true)
+    before do
+      allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
 
       expect(AuthorizedProjectsWorker).to receive(:bulk_perform_and_wait)
-        .with([[1], [2]])
-        .and_return(10)
+                                            .with([[1], [2]])
+                                            .and_return(10)
+    end
 
-      [1, 2].each do |id|
-        expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:stick)
-          .with(:user, id)
-          .ordered
-      end
+    it 'sticks all the updated users and returns the original result', :aggregate_failures do
+      expect(Gitlab::Database::LoadBalancing::Sticking).to receive(:bulk_stick).with(:user, [1, 2])
 
       expect(service.execute).to eq(10)
     end
+
+    it 'avoids N+1 cached queries', :use_sql_query_cache, :request_store do
+      # Run this once to establish a baseline
+      control_count = ActiveRecord::QueryRecorder.new(skip_cached: false) do
+        service.execute
+      end
+
+      service = UserProjectAccessChangedService.new([1, 2, 3, 4, 5])
+
+      allow(AuthorizedProjectsWorker).to receive(:bulk_perform_and_wait)
+                                            .with([[1], [2], [3], [4], [5]])
+                                            .and_return(10)
+
+      expect { service.execute }.not_to exceed_all_query_limit(control_count.count)
+    end
   end
 end