From 0206476ae2a2d659fd0fb42338050253a9a91439 Mon Sep 17 00:00:00 2001
From: Sean McGivern <sean@gitlab.com>
Date: Thu, 7 Jun 2018 12:14:27 +0100
Subject: [PATCH] Fix some N+1s when calculating notification recipients

First N+1: we may have loaded a user's notification settings already, but not
have loaded their sources. Because we're iterating through, we'd potentially
load sources that are completely unrelated, just because they belong to this
user.

Second N+1: we do a separate query for each user who could be subscribed to or
unsubcribed from the target. It's actually more efficient in this case to get
all subscriptions at once, as we will need to check most of them.

We can fix both by the slightly unpleasant means of checking IDs manually,
rather than object equality.
---
 app/models/notification_recipient.rb          |  2 +-
 app/models/user.rb                            |  5 ++-
 .../n-plus-one-notification-recipients.yml    |  5 +++
 .../notification_recipient_service_spec.rb    | 36 +++++++++++++++++++
 4 files changed, 46 insertions(+), 2 deletions(-)
 create mode 100644 changelogs/unreleased/n-plus-one-notification-recipients.yml
 create mode 100644 spec/services/notification_recipient_service_spec.rb

diff --git a/app/models/notification_recipient.rb b/app/models/notification_recipient.rb
index 2c3580bbdc6..79458bd048a 100644
--- a/app/models/notification_recipient.rb
+++ b/app/models/notification_recipient.rb
@@ -64,7 +64,7 @@ class NotificationRecipient
     return false unless @target
     return false unless @target.respond_to?(:subscriptions)
 
-    subscription = @target.subscriptions.find_by_user_id(@user.id)
+    subscription = @target.subscriptions.find { |subscription| subscription.user_id == @user.id }
     subscription && !subscription.subscribed
   end
 
diff --git a/app/models/user.rb b/app/models/user.rb
index e219ab800ad..8e0dc91b2a7 100644
--- a/app/models/user.rb
+++ b/app/models/user.rb
@@ -1038,7 +1038,10 @@ class User < ActiveRecord::Base
 
   def notification_settings_for(source)
     if notification_settings.loaded?
-      notification_settings.find { |notification| notification.source == source }
+      notification_settings.find do |notification|
+        notification.source_type == source.class.base_class.name &&
+          notification.source_id == source.id
+      end
     else
       notification_settings.find_or_initialize_by(source: source)
     end
diff --git a/changelogs/unreleased/n-plus-one-notification-recipients.yml b/changelogs/unreleased/n-plus-one-notification-recipients.yml
new file mode 100644
index 00000000000..91c31e4c930
--- /dev/null
+++ b/changelogs/unreleased/n-plus-one-notification-recipients.yml
@@ -0,0 +1,5 @@
+---
+title: Fix some sources of excessive query counts when calculating notification recipients
+merge_request:
+author:
+type: performance
diff --git a/spec/services/notification_recipient_service_spec.rb b/spec/services/notification_recipient_service_spec.rb
new file mode 100644
index 00000000000..340d4585e0c
--- /dev/null
+++ b/spec/services/notification_recipient_service_spec.rb
@@ -0,0 +1,36 @@
+require 'spec_helper'
+
+describe NotificationRecipientService do
+  let(:service) { described_class }
+  let(:assignee) { create(:user) }
+  let(:project) { create(:project, :public) }
+  let(:other_projects) { create_list(:project, 5, :public) }
+
+  describe '#build_new_note_recipients' do
+    let(:issue) { create(:issue, project: project, assignees: [assignee]) }
+    let(:note) { create(:note_on_issue, noteable: issue, project_id: issue.project_id) }
+
+    def create_watcher
+      watcher = create(:user)
+      create(:notification_setting, project: project, user: watcher, level: :watch)
+
+      other_projects.each do |other_project|
+        create(:notification_setting, project: other_project, user: watcher, level: :watch)
+      end
+    end
+
+    it 'avoids N+1 queries' do
+      create_watcher
+
+      service.build_new_note_recipients(note)
+
+      control_count = ActiveRecord::QueryRecorder.new do
+        service.build_new_note_recipients(note)
+      end
+
+      create_watcher
+
+      expect { service.build_new_note_recipients(note) }.not_to exceed_query_limit(control_count)
+    end
+  end
+end