Preload user status for mentioning users

This change uses BatchLoader to efficiently preload
UserStatus.availability for several different user arrays.

app/services/concerns/users/participable_service.rb

@@ -8,6 +8,8 @@ module Users
       attr_reader :noteable
     end
 
+    private
+
     def noteable_owner
       return [] unless noteable && noteable.author.present?
 
@@ -22,23 +24,29 @@ module Users
     end
 
     def sorted(users)
-      users.uniq.to_a.compact.sort_by(&:username).map do |user|
-        user_as_hash(user)
+      users.uniq.to_a.compact.sort_by(&:username).tap do |users|
+        preload_status(users)
       end
     end
 
     def groups
-      group_counts = GroupMember
-                       .of_groups(current_user.authorized_groups)
-                       .non_request
-                       .count_users_by_group_id
+      current_user.authorized_groups.with_route.sort_by(&:path)
+    end
 
-      current_user.authorized_groups.with_route.sort_by(&:path).map do |group|
-        group_as_hash(group, group_counts)
-      end
+    def render_participants_as_hash(participants)
+      participants.map(&method(:participant_as_hash))
     end
 
-    private
+    def participant_as_hash(participant)
+      case participant
+      when Group
+        group_as_hash(participant)
+      when User
+        user_as_hash(participant)
+      else
+        participant
+      end
+    end
 
     def user_as_hash(user)
       {
@@ -46,12 +54,11 @@ module Users
         username: user.username,
         name: user.name,
         avatar_url: user.avatar_url,
-        availability: nil
+        availability: lazy_user_availability(user).itself # calling #itself to avoid returning a BatchLoader instance
       }
-      # Return nil for availability for now due to https://gitlab.com/gitlab-org/gitlab/-/issues/285442
     end
 
-    def group_as_hash(group, group_counts)
+    def group_as_hash(group)
       {
         type: group.class.name,
         username: group.full_path,
@@ -61,5 +68,27 @@ module Users
         mentionsDisabled: group.mentions_disabled
       }
     end
+
+    def group_counts
+      @group_counts ||= GroupMember
+        .of_groups(current_user.authorized_groups)
+        .non_request
+        .count_users_by_group_id
+    end
+
+    def preload_status(users)
+      users.each { |u| lazy_user_availability(u) }
+    end
+
+    def lazy_user_availability(user)
+      BatchLoader.for(user.id).batch do |user_ids, loader|
+        user_ids.each_slice(1_000) do |sliced_user_ids|
+          UserStatus
+            .select(:user_id, :availability)
+            .primary_key_in(sliced_user_ids)
+            .each { |status| loader.call(status.user_id, status.availability) }
+        end
+      end
+    end
   end
 end

app/services/projects/participants_service.rb

@@ -14,7 +14,7 @@ module Projects
         groups +
         project_members
 
-      participants.uniq
+      render_participants_as_hash(participants.uniq)
     end
 
     def project_members

ee/app/services/groups/participants_service.rb

@@ -7,8 +7,14 @@ module Groups
     def execute(noteable)
       @noteable = noteable
 
-      participants = noteable_owner + participants_in_noteable + all_members + groups + group_members
-      participants.uniq
+      participants =
+        noteable_owner +
+        participants_in_noteable +
+        all_members +
+        groups +
+        group_members
+
+      render_participants_as_hash(participants.uniq)
     end
 
     def all_members

ee/spec/services/groups/participants_service_spec.rb

@@ -42,11 +42,11 @@ RSpec.describe Groups::ParticipantsService do
     it 'returns all participants' do
       service = described_class.new(group, user)
       service.instance_variable_set(:@noteable, epic)
-      result = service.participants_in_noteable
+      result = service.execute(epic)
 
       expected_users = (@users + [user]).map(&method(:user_to_autocompletable))
 
-      expect(result).to match_array(expected_users)
+      expect(result).to include(*expected_users)
     end
   end
 
@@ -55,6 +55,7 @@ RSpec.describe Groups::ParticipantsService do
     let(:group) { create(:group, parent: parent_group) }
     let(:subgroup) { create(:group_with_members, parent: group) }
     let(:subproject) { create(:project, group: subgroup) }
+    let(:epic) { create(:epic, group: group, author: user) }
 
     it 'returns all members in parent groups, sub-groups, and sub-projects' do
       parent_group.add_developer(create(:user))
@@ -63,31 +64,31 @@ RSpec.describe Groups::ParticipantsService do
 
       service = described_class.new(group, user)
       service.instance_variable_set(:@noteable, epic)
-      result = service.group_members
+      result = service.execute(epic)
 
       expected_users = (group.self_and_hierarchy.flat_map(&:users) + subproject.users)
-                       .map(&method(:user_to_autocompletable))
+        .map(&method(:user_to_autocompletable))
 
       expect(expected_users.count).to eq(5)
-      expect(result).to match_array(expected_users)
+      expect(result).to include(*expected_users)
     end
   end
 
-  describe '#groups' do
-    describe 'avatar_url' do
-      let(:groups) { described_class.new(group, user).groups }
+  describe 'group items' do
+    subject(:group_items) { described_class.new(group, user).execute(epic).select { |hash| hash[:type].eql?('Group') } }
 
+    describe 'avatar_url' do
       it 'returns a URL for the avatar' do
-        expect(groups.size).to eq 1
-        expect(groups.first[:avatar_url]).to eq("/uploads/-/system/group/avatar/#{group.id}/dk.png")
+        expect(group_items.size).to eq 1
+        expect(group_items.first[:avatar_url]).to eq("/uploads/-/system/group/avatar/#{group.id}/dk.png")
       end
 
       it 'returns a relative URL for the avatar' do
         stub_config_setting(relative_url_root: '/gitlab')
         stub_config_setting(url: Settings.send(:build_gitlab_url))
 
-        expect(groups.size).to eq 1
-        expect(groups.first[:avatar_url]).to eq("/gitlab/uploads/-/system/group/avatar/#{group.id}/dk.png")
+        expect(group_items.size).to eq 1
+        expect(group_items.first[:avatar_url]).to eq("/gitlab/uploads/-/system/group/avatar/#{group.id}/dk.png")
       end
     end
   end

spec/services/projects/participants_service_spec.rb

@@ -3,57 +3,90 @@
 require 'spec_helper'
 
 RSpec.describe Projects::ParticipantsService do
-  describe '#groups' do
-    let_it_be(:user) { create(:user) }
-    let_it_be(:project) { create(:project, :public) }
-    let(:service) { described_class.new(project, user) }
+  describe '#execute' do
+    let(:user) { create(:user) }
+    let(:project) { create(:project, :public) }
+    let(:noteable) { create(:issue, project: project) }
 
-    it 'avoids N+1 queries' do
-      group_1 = create(:group)
-      group_1.add_owner(user)
+    def run_service
+      described_class.new(project, user).execute(noteable)
+    end
 
-      service.groups # Run general application warmup queries
-      control_count = ActiveRecord::QueryRecorder.new { service.groups }.count
+    before do
+      project.add_developer(user)
 
-      group_2 = create(:group)
-      group_2.add_owner(user)
+      run_service # warmup, runs table cache queries and create queries
+      BatchLoader::Executor.clear_current
+    end
 
-      expect { service.groups }.not_to exceed_query_limit(control_count)
+    it 'avoids N+1 UserDetail queries' do
+      project.add_developer(create(:user))
+
+      control_count = ActiveRecord::QueryRecorder.new { run_service.to_a }.count
+
+      BatchLoader::Executor.clear_current
+
+      project.add_developer(create(:user, status: build(:user_status, availability: :busy)))
+
+      expect { run_service.to_a }.not_to exceed_query_limit(control_count)
     end
 
-    it 'returns correct user counts for groups' do
+    it 'avoids N+1 groups queries' do
       group_1 = create(:group)
       group_1.add_owner(user)
-      group_1.add_owner(create(:user))
+
+      control_count = ActiveRecord::QueryRecorder.new { run_service }.count
+
+      BatchLoader::Executor.clear_current
 
       group_2 = create(:group)
       group_2.add_owner(user)
-      create(:group_member, :access_request, group: group_2, user: create(:user))
 
-      expect(service.groups).to contain_exactly(
-        a_hash_including(name: group_1.full_name, count: 2),
-        a_hash_including(name: group_2.full_name, count: 1)
-      )
+      expect { run_service }.not_to exceed_query_limit(control_count)
     end
 
-    describe 'avatar_url' do
-      let(:group) { create(:group, avatar: fixture_file_upload('spec/fixtures/dk.png')) }
+    describe 'group items' do
+      subject(:group_items) { run_service.select { |hash| hash[:type].eql?('Group') } }
 
-      before do
-        group.add_owner(user)
-      end
+      describe 'group user counts' do
+        let(:group_1) { create(:group) }
+        let(:group_2) { create(:group) }
+
+        before do
+          group_1.add_owner(user)
+          group_1.add_owner(create(:user))
 
-      it 'returns an url for the avatar' do
-        expect(service.groups.size).to eq 1
-        expect(service.groups.first[:avatar_url]).to eq("/uploads/-/system/group/avatar/#{group.id}/dk.png")
+          group_2.add_owner(user)
+          create(:group_member, :access_request, group: group_2, user: create(:user))
+        end
+
+        it 'returns correct user counts for groups' do
+          expect(group_items).to contain_exactly(
+            a_hash_including(name: group_1.full_name, count: 2),
+            a_hash_including(name: group_2.full_name, count: 1)
+          )
+        end
       end
 
-      it 'returns an url for the avatar with relative url' do
-        stub_config_setting(relative_url_root: '/gitlab')
-        stub_config_setting(url: Settings.send(:build_gitlab_url))
+      describe 'avatar_url' do
+        let(:group) { create(:group, avatar: fixture_file_upload('spec/fixtures/dk.png')) }
 
-        expect(service.groups.size).to eq 1
-        expect(service.groups.first[:avatar_url]).to eq("/gitlab/uploads/-/system/group/avatar/#{group.id}/dk.png")
+        before do
+          group.add_owner(user)
+        end
+
+        it 'returns an url for the avatar' do
+          expect(group_items.size).to eq 1
+          expect(group_items.first[:avatar_url]).to eq("/uploads/-/system/group/avatar/#{group.id}/dk.png")
+        end
+
+        it 'returns an url for the avatar with relative url' do
+          stub_config_setting(relative_url_root: '/gitlab')
+          stub_config_setting(url: Settings.send(:build_gitlab_url))
+
+          expect(group_items.size).to eq 1
+          expect(group_items.first[:avatar_url]).to eq("/gitlab/uploads/-/system/group/avatar/#{group.id}/dk.png")
+        end
       end
     end
   end