Merge branch '354016-resolve-growing-query-with-more-invites-3' into 'master'

Refactor bulk max access load for general use

See merge request gitlab-org/gitlab!82318

app/models/concerns/bulk_member_access_load.rb

 # frozen_string_literal: true
 
-# Returns and caches in thread max member access for a resource
-#
 module BulkMemberAccessLoad
   extend ActiveSupport::Concern
 
   included do
-    # Determine the maximum access level for a group of resources in bulk.
-    #
-    # Returns a Hash mapping resource ID -> maximum access level.
-    def max_member_access_for_resource_ids(resource_klass, resource_ids, &block)
-      raise 'Block is mandatory' unless block_given?
-
-      memoization_index = self.id
-      memoization_class = self.class
-
-      resource_ids = resource_ids.uniq
-      memo_id = "#{memoization_class}:#{memoization_index}"
-      access = load_access_hash(resource_klass, memo_id)
-
-      # Look up only the IDs we need
-      resource_ids -= access.keys
-
-      return access if resource_ids.empty?
-
-      resource_access = yield(resource_ids)
-
-      access.merge!(resource_access)
-
-      missing_resource_ids = resource_ids - resource_access.keys
-
-      missing_resource_ids.each do |resource_id|
-        access[resource_id] = Gitlab::Access::NO_ACCESS
-      end
-
-      access
-    end
-
     def merge_value_to_request_store(resource_klass, resource_id, value)
-      max_member_access_for_resource_ids(resource_klass, [resource_id]) do
+      Gitlab::SafeRequestLoader.execute(resource_key: max_member_access_for_resource_key(resource_klass),
+                                        resource_ids: [resource_id],
+                                        default_value: Gitlab::Access::NO_ACCESS) do
         { resource_id => value }
       end
     end
 
-    private
-
-    def max_member_access_for_resource_key(klass, memoization_index)
-      "max_member_access_for_#{klass.name.underscore.pluralize}:#{memoization_index}"
-    end
-
-    def load_access_hash(resource_klass, memo_id)
-      return {} unless Gitlab::SafeRequestStore.active?
-
-      key = max_member_access_for_resource_key(resource_klass, memo_id)
-      Gitlab::SafeRequestStore[key] ||= {}
-
-      Gitlab::SafeRequestStore[key]
+    def max_member_access_for_resource_key(klass)
+      "max_member_access_for_#{klass.name.underscore.pluralize}:#{self.class}:#{self.id}"
     end
   end
 end

app/models/group.rb

@@ -816,7 +816,9 @@ class Group < Namespace
   private
 
   def max_member_access(user_ids)
-    max_member_access_for_resource_ids(User, user_ids) do |user_ids|
+    Gitlab::SafeRequestLoader.execute(resource_key: max_member_access_for_resource_key(User),
+                                      resource_ids: user_ids,
+                                      default_value: Gitlab::Access::NO_ACCESS) do |user_ids|
       members_with_parents.where(user_id: user_ids).group(:user_id).maximum(:access_level)
     end
   end

app/models/project_team.rb

@@ -179,7 +179,9 @@ class ProjectTeam
   #
   # Returns a Hash mapping user ID -> maximum access level.
   def max_member_access_for_user_ids(user_ids)
-    project.max_member_access_for_resource_ids(User, user_ids) do |user_ids|
+    Gitlab::SafeRequestLoader.execute(resource_key: project.max_member_access_for_resource_key(User),
+                                      resource_ids: user_ids,
+                                      default_value: Gitlab::Access::NO_ACCESS) do |user_ids|
       project.project_authorizations
              .where(user: user_ids)
              .group(:user_id)

app/models/user.rb

@@ -1862,7 +1862,9 @@ class User < ApplicationRecord
   #
   # Returns a Hash mapping project ID -> maximum access level.
   def max_member_access_for_project_ids(project_ids)
-    max_member_access_for_resource_ids(Project, project_ids) do |project_ids|
+    Gitlab::SafeRequestLoader.execute(resource_key: max_member_access_for_resource_key(Project),
+                                      resource_ids: project_ids,
+                                      default_value: Gitlab::Access::NO_ACCESS) do |project_ids|
       project_authorizations.where(project: project_ids)
                             .group(:project_id)
                             .maximum(:access_level)
@@ -1877,7 +1879,9 @@ class User < ApplicationRecord
   #
   # Returns a Hash mapping project ID -> maximum access level.
   def max_member_access_for_group_ids(group_ids)
-    max_member_access_for_resource_ids(Group, group_ids) do |group_ids|
+    Gitlab::SafeRequestLoader.execute(resource_key: max_member_access_for_resource_key(Group),
+                                      resource_ids: group_ids,
+                                      default_value: Gitlab::Access::NO_ACCESS) do |group_ids|
       group_members.where(source: group_ids).group(:source_id).maximum(:access_level)
     end
   end

lib/gitlab/safe_request_loader.rb

+# frozen_string_literal: true
+
+module Gitlab
+  class SafeRequestLoader
+    def self.execute(args, &block)
+      new(**args).execute(&block)
+    end
+
+    def initialize(resource_key:, resource_ids:, default_value: nil)
+      @resource_key = resource_key
+      @resource_ids = resource_ids.uniq
+      @default_value = default_value
+      @resource_data = {}
+    end
+
+    def execute(&block)
+      raise ArgumentError, 'Block is mandatory' unless block_given?
+
+      load_resource_data
+      remove_loaded_resource_ids
+
+      update_resource_data(&block)
+
+      resource_data
+    end
+
+    private
+
+    attr_reader :resource_key, :resource_ids, :default_value, :resource_data, :missing_resource_ids
+
+    def load_resource_data
+      @resource_data = Gitlab::SafeRequestStore.fetch(resource_key) { resource_data }
+    end
+
+    def remove_loaded_resource_ids
+      # Look up only the IDs we need
+      @missing_resource_ids = resource_ids - resource_data.keys
+    end
+
+    def update_resource_data(&block)
+      return if missing_resource_ids.blank?
+
+      reloaded_resource_data = yield(missing_resource_ids)
+
+      @resource_data.merge!(reloaded_resource_data)
+
+      mark_absent_values
+    end
+
+    def mark_absent_values
+      absent = (missing_resource_ids - resource_data.keys).to_h { [_1, default_value] }
+      @resource_data.merge!(absent)
+    end
+  end
+end

spec/lib/gitlab/safe_request_loader_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::SafeRequestLoader, :aggregate_failures do
+  let(:resource_key) { '_key_' }
+  let(:resource_ids) { [] }
+  let(:args) { { resource_key: resource_key, resource_ids: resource_ids } }
+  let(:block) { proc { {} } }
+
+  describe '.execute', :request_store do
+    let(:resource_data) { { 'foo' => 'bar' } }
+
+    before do
+      Gitlab::SafeRequestStore[resource_key] = resource_data
+    end
+
+    subject(:execute_instance) { described_class.execute(**args, &block) }
+
+    it 'gets data from the store and returns it' do
+      expect(execute_instance.keys).to contain_exactly(*resource_data.keys)
+      expect(execute_instance).to match(a_hash_including(resource_data))
+      expect_store_to_be_updated
+    end
+  end
+
+  describe '#execute' do
+    subject(:execute_instance) { described_class.new(**args).execute(&block) }
+
+    context 'without a block' do
+      let(:block) { nil }
+
+      it 'raises an error' do
+        expect { execute_instance }.to raise_error(ArgumentError, 'Block is mandatory')
+      end
+    end
+
+    context 'when a resource_id is nil' do
+      let(:block) { proc { {} } }
+      let(:resource_ids) { [nil] }
+
+      it 'contains resource_data with nil key' do
+        expect(execute_instance.keys).to contain_exactly(nil)
+        expect(execute_instance).to match(a_hash_including(nil => nil))
+      end
+    end
+
+    context 'with SafeRequestStore considerations' do
+      let(:resource_data) { { 'foo' => 'bar' } }
+
+      before do
+        Gitlab::SafeRequestStore[resource_key] = resource_data
+      end
+
+      context 'when request store is active', :request_store do
+        it 'gets data from the store' do
+          expect(execute_instance.keys).to contain_exactly(*resource_data.keys)
+          expect(execute_instance).to match(a_hash_including(resource_data))
+          expect_store_to_be_updated
+        end
+
+        context 'with already loaded resource_ids', :request_store do
+          let(:resource_key) { 'foo_data' }
+          let(:existing_resource_data) { { 'foo' => 'zoo' } }
+          let(:block) { proc { { 'foo' => 'bar' } } }
+          let(:resource_ids) { ['foo'] }
+
+          before do
+            Gitlab::SafeRequestStore[resource_key] = existing_resource_data
+          end
+
+          it 'does not re-fetch data if resource_id already exists' do
+            expect(execute_instance.keys).to contain_exactly(*resource_ids)
+            expect(execute_instance).to match(a_hash_including(existing_resource_data))
+            expect_store_to_be_updated
+          end
+
+          context 'with mixture of new and existing resource_ids' do
+            let(:existing_resource_data) { { 'foo' => 'bar' } }
+            let(:resource_ids) { %w[foo bar] }
+
+            context 'when block does not filter for only the missing resource_ids' do
+              let(:block) { proc { { 'foo' => 'zoo', 'bar' => 'foo' } } }
+
+              it 'overwrites existing keyed data with results from the block' do
+                expect(execute_instance.keys).to contain_exactly(*resource_ids)
+                expect(execute_instance).to match(a_hash_including(block.call))
+                expect_store_to_be_updated
+              end
+            end
+
+            context 'when passing the missing resource_ids to a block that filters for them' do
+              let(:block) { proc { |rids| { 'foo' => 'zoo', 'bar' => 'foo' }.select { |k, _v| rids.include?(k) } } }
+
+              it 'only updates resource_data with keyed items that did not exist' do
+                expect(execute_instance.keys).to contain_exactly(*resource_ids)
+                expect(execute_instance).to match(a_hash_including({ 'foo' => 'bar', 'bar' => 'foo' }))
+                expect_store_to_be_updated
+              end
+            end
+
+            context 'with default_value for resource_ids that did not exist in the results' do
+              context 'when default_value is provided' do
+                let(:args) { { resource_key: resource_key, resource_ids: resource_ids, default_value: '_value_' } }
+
+                it 'populates a default value' do
+                  expect(execute_instance.keys).to contain_exactly(*resource_ids)
+                  expect(execute_instance).to match(a_hash_including({ 'foo' => 'bar', 'bar' => '_value_' }))
+                  expect_store_to_be_updated
+                end
+              end
+
+              context 'when default_value is not provided' do
+                it 'populates a default_value of nil' do
+                  expect(execute_instance.keys).to contain_exactly(*resource_ids)
+                  expect(execute_instance).to match(a_hash_including({ 'foo' => 'bar', 'bar' => nil }))
+                  expect_store_to_be_updated
+                end
+              end
+            end
+          end
+        end
+      end
+
+      context 'when request store is not active' do
+        let(:block) { proc { { 'foo' => 'bar' } } }
+        let(:resource_ids) { ['foo'] }
+
+        it 'has no data added from the store' do
+          expect(execute_instance).to eq(block.call)
+        end
+
+        context 'with mixture of new and existing resource_ids' do
+          let(:resource_ids) { %w[foo bar] }
+
+          context 'when block does not filter out existing resource_data keys' do
+            let(:block) { proc { { 'foo' => 'zoo', 'bar' => 'foo' } } }
+
+            it 'overwrites existing keyed data with results from the block' do
+              expect(execute_instance.keys).to contain_exactly(*resource_ids)
+              expect(execute_instance).to match(a_hash_including(block.call))
+            end
+          end
+
+          context 'when passing the missing resource_ids to a block that filters for them' do
+            let(:block) { proc { |rids| { 'foo' => 'zoo', 'bar' => 'foo' }.select { |k, _v| rids.include?(k) } } }
+
+            it 'only updates resource_data with keyed items that did not exist' do
+              expect(execute_instance.keys).to contain_exactly(*resource_ids)
+              expect(execute_instance).to match(a_hash_including({ 'foo' => 'zoo', 'bar' => 'foo' }))
+            end
+          end
+
+          context 'with default_value for resource_ids that did not exist in the results' do
+            context 'when default_value is provided' do
+              let(:args) { { resource_key: resource_key, resource_ids: resource_ids, default_value: '_value_' } }
+
+              it 'populates a default value' do
+                expect(execute_instance.keys).to contain_exactly(*resource_ids)
+                expect(execute_instance).to match(a_hash_including({ 'foo' => 'bar', 'bar' => '_value_' }))
+              end
+            end
+
+            context 'when default_value is not provided' do
+              it 'populates a default_value of nil' do
+                expect(execute_instance.keys).to contain_exactly(*resource_ids)
+                expect(execute_instance).to match(a_hash_including({ 'foo' => 'bar', 'bar' => nil }))
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+
+  def expect_store_to_be_updated
+    expect(execute_instance).to match(a_hash_including(Gitlab::SafeRequestStore[resource_key]))
+    expect(execute_instance.keys).to contain_exactly(*Gitlab::SafeRequestStore[resource_key].keys)
+  end
+end