Fix excessive DB access in set_assignees mutation

app/graphql/mutations/concerns/mutations/assignable.rb

@@ -13,14 +13,13 @@ module Mutations
       argument :operation_mode,
                Types::MutationOperationModeEnum,
                required: false,
+               default_value: Types::MutationOperationModeEnum.default_mode,
                description: 'The operation to perform. Defaults to REPLACE.'
     end
 
-    def resolve(project_path:, iid:, assignee_usernames:, operation_mode: Types::MutationOperationModeEnum.enum[:replace])
+    def resolve(project_path:, iid:, assignee_usernames:, operation_mode:)
       resource = authorized_find!(project_path: project_path, iid: iid)
 
-      Gitlab::QueryLimiting.disable!('https://gitlab.com/gitlab-org/gitlab/issues/36098') if resource.is_a?(MergeRequest)
-
       update_service_class.new(
         resource.project,
         current_user,
@@ -35,18 +34,22 @@ module Mutations
 
     private
 
-    def assignee_ids(resource, usernames, operation_mode)
-      assignee_ids = []
-      assignee_ids += resource.assignees.map(&:id) if Types::MutationOperationModeEnum.enum.values_at(:remove, :append).include?(operation_mode)
-      user_ids = UsersFinder.new(current_user, username: usernames).execute.map(&:id)
+    def assignee_ids(resource, usernames, mode)
+      new = UsersFinder.new(current_user, username: usernames).execute.map(&:id)
+
+      transform_list(mode, resource, new)
+    end
 
-      if operation_mode == Types::MutationOperationModeEnum.enum[:remove]
-        assignee_ids -= user_ids
-      else
-        assignee_ids |= user_ids
+    def current_assignee_ids(resource)
+      resource.assignees.map(&:id)
     end
 
-      assignee_ids
+    def transform_list(mode, resource, new_values)
+      case mode
+      when 'REPLACE' then new_values
+      when 'APPEND' then current_assignee_ids(resource) | new_values
+      when 'REMOVE' then current_assignee_ids(resource) - new_values
+      end
     end
   end
 end

app/graphql/mutations/merge_requests/set_assignees.rb

@@ -8,7 +8,7 @@ module Mutations
       include Assignable
 
       def update_service_class
-        ::MergeRequests::UpdateService
+        ::MergeRequests::UpdateAssigneesService
       end
     end
   end

app/graphql/types/mutation_operation_mode_enum.rb

@@ -10,5 +10,13 @@ module Types
     value 'REPLACE', 'Performs a replace operation.'
     value 'APPEND', 'Performs an append operation.'
     value 'REMOVE', 'Performs a removal operation.'
+
+    def self.default_mode
+      enum[:replace]
+    end
+
+    def self.transform_modes
+      enum.values_at(:remove, :append)
+    end
   end
 end

app/services/merge_requests/update_assignees_service.rb

+# frozen_string_literal: true
+
+module MergeRequests
+  class UpdateAssigneesService < UpdateService
+    # a stripped down service that only does what it must to update the
+    # assignees, and knows that it does not have to check for other updates.
+    # This saves a lot of queries for irrelevant things that cannot possibly
+    # change in the execution of this service.
+    def execute(merge_request)
+      return unless current_user&.can?(:update_merge_request, merge_request)
+
+      old_ids = merge_request.assignees.map(&:id)
+      return if old_ids.to_set == update_attrs[:assignee_ids].to_set # no-change
+
+      merge_request.update!(**update_attrs)
+
+      # Defer the more expensive operations (handle_assignee_changes) to the background
+      MergeRequests::AssigneesChangeWorker.perform_async(merge_request.id, current_user.id, old_ids)
+    end
+
+    def handle_assignee_changes(merge_request, old_assignees)
+      # exposes private method from super-class
+      users = old_assignees.to_a
+      handle_assignees_change(merge_request, users)
+      execute_hooks(
+        merge_request,
+        'update',
+        old_associations: { assignees: users }
+      )
+    end
+
+    private
+
+    def assignee_ids
+      params.fetch(:assignee_ids).first(1)
+    end
+
+    def update_attrs
+      @attrs ||= { updated_at: Time.current, updated_by: current_user, assignee_ids: assignee_ids }
+    end
+  end
+end
+
+MergeRequests::UpdateAssigneesService.prepend_if_ee('EE::MergeRequests::UpdateAssigneesService')

app/workers/all_queues.yml

@@ -1860,6 +1860,14 @@
   :weight: 1
   :idempotent: true
   :tags: []
+- :name: merge_requests_assignees_change
+  :feature_category: :source_code_management
+  :has_external_dependencies:
+  :urgency: :high
+  :resource_boundary: :unknown
+  :weight: 1
+  :idempotent: true
+  :tags: []
 - :name: merge_requests_delete_source_branch
   :feature_category: :source_code_management
   :has_external_dependencies:

app/workers/merge_requests/assignees_change_worker.rb

+# frozen_string_literal: true
+
+class MergeRequests::AssigneesChangeWorker
+  include ApplicationWorker
+
+  feature_category :source_code_management
+  urgency :high
+  deduplicate :until_executed
+  idempotent!
+
+  def perform(merge_request_id, user_id, old_assignee_ids)
+    merge_request = MergeRequest.find(merge_request_id)
+    current_user = User.find(user_id)
+
+    # if a user was added and then removed, or removed and then added
+    # while waiting for this job to run, assume that nothing happened.
+    users = User.id_in(old_assignee_ids - merge_request.assignee_ids)
+
+    return if users.blank?
+
+    service = ::MergeRequests::UpdateAssigneesService.new(
+      merge_request.target_project,
+      current_user
+    )
+
+    service.handle_assignee_changes(merge_request, users)
+  rescue ActiveRecord::RecordNotFound
+  end
+end

changelogs/unreleased/ajk-36098-fix-excessive-db-activity-in-mr-assign-mutation.yml

+---
+title: Reduce number of queries in mergeRequestSetAssignees GraphQL mutation
+merge_request: 57523
+author:
+type: performance

config/sidekiq_queues.yml

@@ -210,6 +210,8 @@
   - 1
 - - merge_request_reset_approvals
   - 1
+- - merge_requests_assignees_change
+  - 1
 - - merge_requests_delete_source_branch
   - 1
 - - metrics_dashboard_prune_old_annotations

ee/app/services/ee/merge_requests/update_assignees_service.rb

+# frozen_string_literal: true
+
+module EE
+  module MergeRequests
+    module UpdateAssigneesService
+      def assignee_ids
+        if project.licensed_feature_available?(:multiple_merge_request_assignees)
+          params.fetch(:assignee_ids)
+        else
+          super
+        end
+      end
+    end
+  end
+end

ee/app/services/ee/merge_requests/update_service.rb

@@ -18,9 +18,7 @@ module EE
 
         merge_request = super(merge_request)
 
-        if should_remove_old_approvers && merge_request.valid?
-          cleanup_approvers(merge_request, reload: true)
-        end
+        cleanup_approvers(merge_request, reload: true) if should_remove_old_approvers && merge_request.valid?
 
         merge_request.reset_approval_cache!
 

ee/spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb

@@ -5,23 +5,21 @@ require 'spec_helper'
 RSpec.describe 'Setting assignees of a merge request' do
   include GraphqlHelpers
 
-  let(:current_user) { create(:user) }
-  let(:merge_request) { create(:merge_request) }
-  let(:project) { merge_request.project }
-  let(:assignees) { create_list(:user, 3) }
-  let(:extra_assignees) { create_list(:user, 2) }
+  let_it_be(:project) { create(:project, :repository) }
+  let_it_be(:current_user) { create(:user, developer_projects: [project]) }
+  let_it_be(:merge_request) { create(:merge_request, source_project: project) }
+  let_it_be(:assignees) { create_list(:user, 3, developer_projects: [project]) }
+  let_it_be(:extra_assignees) { create_list(:user, 2, developer_projects: [project]) }
+
   let(:input) { { assignee_usernames: assignees.map(&:username) } }
   let(:expected_result) do
     assignees.map { |u| { 'username' => u.username } }
   end
 
-  let(:mutation) do
-    variables = {
-      project_path: project.full_path,
-      iid: merge_request.iid.to_s
-    }
-    graphql_mutation(:merge_request_set_assignees, variables.merge(input),
-                     <<-QL.strip_heredoc
+  def mutation(vars = input, mr = merge_request)
+    variables = vars.merge(project_path: mr.project.full_path, iid: mr.iid.to_s)
+
+    graphql_mutation(:merge_request_set_assignees, variables, <<-QL.strip_heredoc)
       clientMutationId
       errors
       mergeRequest {
@@ -33,7 +31,6 @@ RSpec.describe 'Setting assignees of a merge request' do
         }
       }
     QL
-    )
   end
 
   def mutation_response
@@ -45,11 +42,7 @@ RSpec.describe 'Setting assignees of a merge request' do
   end
 
   before do
-    project.add_developer(current_user)
-    assignees.each do |user|
-      project.add_developer(user)
-    end
-    extra_assignees.each do |user|
+    [current_user, *assignees, *extra_assignees].each do |user|
       project.add_developer(user)
     end
   end
@@ -76,21 +69,50 @@ RSpec.describe 'Setting assignees of a merge request' do
   end
 
   context 'when passing append as true' do
-    let(:input) { { assignee_usernames: assignees.map(&:username), operation_mode: Types::MutationOperationModeEnum.enum[:append] } }
+    let(:mode) { Types::MutationOperationModeEnum.enum[:append] }
+    let(:usernames) { assignees.map(&:username) }
+    let(:input) { { operation_mode: mode } }
+
     let(:expected_result) do
-      assignees.map { |u| { 'username' => u.username } } + extra_assignees.map { |u| { 'username' => u.username } }
+      (assignees + extra_assignees).map { |u| { 'username' => u.username } }
     end
 
     before do
+      merge_request.reload
       merge_request.assignees = extra_assignees
       merge_request.save!
     end
 
     it 'does not remove users not in the list' do
-      post_graphql_mutation(mutation, current_user: current_user)
+      vars = input.merge(assignee_usernames: usernames)
+      post_graphql_mutation(mutation(vars), current_user: current_user)
 
       expect(response).to have_gitlab_http_status(:success)
       expect(mutation_assignee_nodes).to match_array(expected_result)
     end
+
+    describe 'performance' do
+      it 'is scalable' do
+        mr_a = create(:merge_request, :unique_branches, source_project: project)
+        mr_b = create(:merge_request, :unique_branches, source_project: project)
+
+        add_one_assignee = mutation(input.merge(assignee_usernames: usernames.take(1)), mr_a)
+        add_two_assignees = mutation(input.merge(assignee_usernames: usernames.last(2)), mr_b)
+
+        baseline = ActiveRecord::QueryRecorder.new do
+          post_graphql_mutation(add_one_assignee, current_user: current_user)
+        end
+
+        # given the way ActiveRecord implements MergeRequest#assignee_ids=(ids),
+        # we to live with a slight inefficiency here:
+        # For each ID, AR issues:
+        #   - SELECT 1 AS one FROM "merge_request_assignees"...
+        # Followed by:
+        #   - INSERT INTO "merge_request_assignees" ("user_id", "merge_request_id", "created_at")...
+        expect do
+          post_graphql_mutation(add_two_assignees, current_user: current_user)
+        end.not_to exceed_query_limit(baseline.count + 2)
+      end
+    end
   end
 end

ee/spec/support/shared_examples/graphql/mutations/set_multiple_assignees_shared_examples.rb

@@ -9,10 +9,17 @@ RSpec.shared_examples 'a multi-assignable resource' do
 
   describe '#resolve' do
     let_it_be(:assignees) { create_list(:user, 3) }
+
+    let(:mode) { Types::MutationOperationModeEnum.default_mode }
     let(:assignee_usernames) { assignees.map(&:username) }
     let(:mutated_resource) { subject[resource.class.name.underscore.to_sym] }
 
-    subject { mutation.resolve(project_path: resource.project.full_path, iid: resource.iid, assignee_usernames: assignee_usernames) }
+    subject do
+      mutation.resolve(project_path: resource.project.full_path,
+                       iid: resource.iid,
+                       operation_mode: mode,
+                       assignee_usernames: assignee_usernames)
+    end
 
     before do
       assignees.each do |user|

spec/graphql/mutations/issues/set_assignees_spec.rb

@@ -11,7 +11,12 @@ RSpec.describe Mutations::Issues::SetAssignees do
     subject(:mutation) { described_class.new(object: nil, context: { current_user: user }, field: nil) }
 
     describe '#resolve' do
-      subject { mutation.resolve(project_path: issue.project.full_path, iid: issue.iid, assignee_usernames: [assignee.username]) }
+      subject do
+        mutation.resolve(project_path: issue.project.full_path,
+                         iid: issue.iid,
+                         operation_mode: Types::MutationOperationModeEnum.default_mode,
+                         assignee_usernames: [assignee.username])
+      end
 
       it_behaves_like 'permission level for issue mutation is correctly verified'
     end

spec/graphql/mutations/merge_requests/set_assignees_spec.rb

@@ -11,7 +11,12 @@ RSpec.describe Mutations::MergeRequests::SetAssignees do
     subject(:mutation) { described_class.new(object: nil, context: { current_user: user }, field: nil) }
 
     describe '#resolve' do
-      subject { mutation.resolve(project_path: merge_request.project.full_path, iid: merge_request.iid, assignee_usernames: [assignee.username]) }
+      subject do
+        mutation.resolve(project_path: merge_request.project.full_path,
+                         iid: merge_request.iid,
+                         operation_mode: described_class.arguments['operationMode'].default_value,
+                         assignee_usernames: [assignee.username])
+      end
 
       it_behaves_like 'permission level for merge request mutation is correctly verified'
     end

spec/requests/api/graphql/mutations/merge_requests/set_assignees_spec.rb

@@ -5,11 +5,12 @@ require 'spec_helper'
 RSpec.describe 'Setting assignees of a merge request' do
   include GraphqlHelpers
 
-  let(:current_user) { create(:user) }
-  let(:merge_request) { create(:merge_request) }
-  let(:project) { merge_request.project }
-  let(:assignee) { create(:user) }
-  let(:assignee2) { create(:user) }
+  let_it_be(:project) { create(:project, :repository) }
+  let_it_be(:current_user) { create(:user, developer_projects: [project]) }
+  let_it_be(:assignee) { create(:user) }
+  let_it_be(:assignee2) { create(:user) }
+  let_it_be_with_reload(:merge_request) { create(:merge_request, source_project: project) }
+
   let(:input) { { assignee_usernames: [assignee.username] } }
   let(:expected_result) do
     [{ 'username' => assignee.username }]
@@ -44,10 +45,19 @@ RSpec.describe 'Setting assignees of a merge request' do
     mutation_response['mergeRequest']['assignees']['nodes']
   end
 
+  def run_mutation!
+    recorder = ActiveRecord::QueryRecorder.new do
+      post_graphql_mutation(mutation, current_user: current_user)
+    end
+
+    expect(recorder.count).to be <= db_query_limit
+  end
+
   before do
-    project.add_developer(current_user)
     project.add_developer(assignee)
     project.add_developer(assignee2)
+
+    merge_request.update!(assignees: [])
   end
 
   it 'returns an error if the user is not allowed to update the merge request' do
@@ -56,23 +66,29 @@ RSpec.describe 'Setting assignees of a merge request' do
     expect(graphql_errors).not_to be_empty
   end
 
-  it 'does not allow members without the right permission to add assignees' do
-    user = create(:user)
-    project.add_guest(user)
+  context 'when the current user does not have permission to add assignees' do
+    let(:current_user) { create(:user) }
+    let(:db_query_limit) { 27 }
+
+    it 'does not change the assignees' do
+      project.add_guest(current_user)
 
-    post_graphql_mutation(mutation, current_user: user)
+      expect { run_mutation! }.not_to change { merge_request.reset.assignees.pluck(:id) }
 
       expect(graphql_errors).not_to be_empty
     end
+  end
 
   context 'with assignees already assigned' do
+    let(:db_query_limit) { 38 }
+
     before do
       merge_request.assignees = [assignee2]
       merge_request.save!
     end
 
     it 'replaces the assignee' do
-      post_graphql_mutation(mutation, current_user: current_user)
+      run_mutation!
 
       expect(response).to have_gitlab_http_status(:success)
       expect(mutation_assignee_nodes).to match_array(expected_result)
@@ -80,6 +96,7 @@ RSpec.describe 'Setting assignees of a merge request' do
   end
 
   context 'when passing an empty list of assignees' do
+    let(:db_query_limit) { 31 }
     let(:input) { { assignee_usernames: [] } }
 
     before do
@@ -88,7 +105,7 @@ RSpec.describe 'Setting assignees of a merge request' do
     end
 
     it 'removes assignee' do
-      post_graphql_mutation(mutation, current_user: current_user)
+      run_mutation!
 
       expect(response).to have_gitlab_http_status(:success)
       expect(mutation_assignee_nodes).to eq([])
@@ -96,7 +113,9 @@ RSpec.describe 'Setting assignees of a merge request' do
   end
 
   context 'when passing append as true' do
-    let(:input) { { assignee_usernames: [assignee2.username], operation_mode: Types::MutationOperationModeEnum.enum[:append] } }
+    let(:mode) { Types::MutationOperationModeEnum.enum[:append] }
+    let(:input) { { assignee_usernames: [assignee2.username], operation_mode: mode } }
+    let(:db_query_limit) { 20 }
 
     before do
       # In CE, APPEND is a NOOP as you can't have multiple assignees
@@ -108,7 +127,7 @@ RSpec.describe 'Setting assignees of a merge request' do
     end
 
     it 'does not replace the assignee in CE' do
-      post_graphql_mutation(mutation, current_user: current_user)
+      run_mutation!
 
       expect(response).to have_gitlab_http_status(:success)
       expect(mutation_assignee_nodes).to match_array(expected_result)
@@ -116,7 +135,9 @@ RSpec.describe 'Setting assignees of a merge request' do
   end
 
   context 'when passing remove as true' do
-    let(:input) { { assignee_usernames: [assignee.username], operation_mode: Types::MutationOperationModeEnum.enum[:remove] } }
+    let(:db_query_limit) { 31 }
+    let(:mode) { Types::MutationOperationModeEnum.enum[:remove] }
+    let(:input) { { assignee_usernames: [assignee.username], operation_mode: mode } }
     let(:expected_result) { [] }
 
     before do
@@ -125,7 +146,7 @@ RSpec.describe 'Setting assignees of a merge request' do
     end
 
     it 'removes the users in the list, while adding none' do
-      post_graphql_mutation(mutation, current_user: current_user)
+      run_mutation!
 
       expect(response).to have_gitlab_http_status(:success)
       expect(mutation_assignee_nodes).to match_array(expected_result)

spec/services/merge_requests/update_assignees_service_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe MergeRequests::UpdateAssigneesService do
+  include AfterNextHelpers
+
+  let_it_be(:group) { create(:group, :public) }
+  let_it_be(:project) { create(:project, :private, :repository, group: group) }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:user2) { create(:user) }
+  let_it_be(:user3) { create(:user) }
+
+  let_it_be_with_reload(:merge_request) do
+    create(:merge_request, :simple, :unique_branches,
+           title: 'Old title',
+           description: "FYI #{user2.to_reference}",
+           assignee_ids: [user3.id],
+           source_project: project,
+           author: create(:user))
+  end
+
+  before do
+    project.add_maintainer(user)
+    project.add_developer(user2)
+    project.add_developer(user3)
+  end
+
+  let(:service) { described_class.new(project, user, opts) }
+  let(:opts) { { assignee_ids: [user2.id] } }
+
+  describe 'execute' do
+    def update_merge_request
+      service.execute(merge_request)
+      merge_request.reload
+    end
+
+    context 'when the parameters are valid' do
+      it 'updates the MR, and queues the more expensive work for later' do
+        expect(MergeRequests::AssigneesChangeWorker)
+          .to receive(:perform_async)
+          .with(merge_request.id, user.id, [user3.id])
+
+        expect { update_merge_request }
+          .to change(merge_request, :assignees).to([user2])
+          .and change(merge_request, :updated_at)
+          .and change(merge_request, :updated_by).to(user)
+      end
+
+      it 'is more efficient than using the full update-service' do
+        allow(MergeRequests::AssigneesChangeWorker)
+          .to receive(:perform_async)
+          .with(merge_request.id, user.id, [user3.id])
+
+        other_mr = create(:merge_request, :simple, :unique_branches,
+                          title: merge_request.title,
+                          description: merge_request.description,
+                          assignee_ids: merge_request.assignee_ids,
+                          source_project: merge_request.project,
+                          author: merge_request.author)
+
+        update_service = ::MergeRequests::UpdateService.new(project, user, opts)
+
+        expect { service.execute(merge_request) }
+          .to issue_fewer_queries_than { update_service.execute(other_mr) }
+      end
+    end
+  end
+
+  describe '#handle_assignee_changes' do
+    subject { service.handle_assignee_changes(merge_request, [user2]) }
+
+    it 'calls UpdateService#handle_assignee_changes and executes hooks' do
+      expect(service).to receive(:handle_assignees_change).with(merge_request, [user2])
+      expect(merge_request.project).to receive(:execute_hooks).with(anything, :merge_request_hooks)
+      expect(merge_request.project).to receive(:execute_services).with(anything, :merge_request_hooks)
+      expect(service).to receive(:enqueue_jira_connect_messages_for).with(merge_request)
+
+      subject
+    end
+  end
+end

spec/support/shared_examples/graphql/mutations/set_assignees_shared_examples.rb

@@ -10,10 +10,17 @@ RSpec.shared_examples 'an assignable resource' do
   describe '#resolve' do
     let_it_be(:assignee) { create(:user) }
     let_it_be(:assignee2) { create(:user) }
+
     let(:assignee_usernames) { [assignee.username] }
     let(:mutated_resource) { subject[resource.class.name.underscore.to_sym] }
+    let(:mode) { described_class.arguments['operationMode'].default_value }
 
-    subject { mutation.resolve(project_path: resource.project.full_path, iid: resource.iid, assignee_usernames: assignee_usernames) }
+    subject do
+      mutation.resolve(project_path: resource.project.full_path,
+                       iid: resource.iid,
+                       operation_mode: mode,
+                       assignee_usernames: assignee_usernames)
+    end
 
     before do
       resource.project.add_developer(assignee)

spec/workers/merge_requests/assignees_change_worker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe MergeRequests::AssigneesChangeWorker do
+  include AfterNextHelpers
+
+  let_it_be(:merge_request) { create(:merge_request) }
+  let_it_be(:user) { create(:user) }
+  let_it_be(:old_assignees) { create_list(:user, 3) }
+
+  let(:user_ids) { old_assignees.map(&:id).to_a }
+  let(:worker) { described_class.new }
+
+  it_behaves_like 'an idempotent worker' do
+    let(:job_args) { [merge_request.id, user.id, user_ids] }
+  end
+
+  describe '#perform' do
+    context 'with a non-existing merge request' do
+      it 'does nothing' do
+        expect(::MergeRequests::UpdateAssigneesService).not_to receive(:new)
+
+        worker.perform(non_existing_record_id, user.id, user_ids)
+      end
+    end
+
+    context 'with a non-existing user' do
+      it 'does nothing' do
+        expect(::MergeRequests::UpdateAssigneesService).not_to receive(:new)
+
+        worker.perform(merge_request.id, non_existing_record_id, user_ids)
+      end
+    end
+
+    context 'when there are no changes' do
+      it 'does nothing' do
+        expect(::MergeRequests::UpdateAssigneesService).not_to receive(:new)
+
+        worker.perform(merge_request.id, user.id, merge_request.assignee_ids)
+      end
+    end
+
+    context 'when the old users cannot be found' do
+      it 'does nothing' do
+        expect(::MergeRequests::UpdateAssigneesService).not_to receive(:new)
+
+        worker.perform(merge_request.id, user.id, [non_existing_record_id])
+      end
+    end
+
+    it 'gets MergeRequests::UpdateAssigneesService to handle the changes' do
+      expect_next(::MergeRequests::UpdateAssigneesService)
+        .to receive(:handle_assignee_changes).with(merge_request, old_assignees)
+
+      worker.perform(merge_request.id, user.id, user_ids)
+    end
+  end
+end