Fix an N+1 in avatar URLsfix-avatars-n-plus-one

This is tricky: the query was being run in
`ObjectStorage::Extension::RecordsUploads#retrieve_from_store!`, but we can't
just add batch loading there, because the `#upload=` method there would use the
result immediately, making the batch only have one item.

Instead, we can pre-emptively add an item to the batch whenever an avatarable
object is initialized, and then reuse that batch item in
`#retrieve_from_store!`. However, this also has problems:

1. There is a lot of logic in `Avatarable#retrieve_upload_from_batch`.
2. Some of that logic constructs a 'fake' model for the batch key. This should
   be fine, because of ActiveRecord's override of `#==`, but it relies on that
   staying the same.

8 files changed, 85 insertions, 3 deletions

diff --git a/app/models/concerns/avatarable.rb b/app/models/concerns/avatarable.rb
index 13246a774e3..f6315a22ce7 100644
--- a/app/models/concerns/avatarable.rb
+++ b/app/models/concerns/avatarable.rb
@@ -4,11 +4,14 @@ module Avatarable
   included do
     prepend ShadowMethods
     include ObjectStorage::BackgroundMove
+    include Gitlab::Utils::StrongMemoize
 
     validate :avatar_type, if: ->(user) { user.avatar.present? && user.avatar_changed? }
     validates :avatar, file_size: { maximum: 200.kilobytes.to_i }
 
     mount_uploader :avatar, AvatarUploader
+
+    after_initialize :add_avatar_to_batch
   end
 
   module ShadowMethods
@@ -18,6 +21,17 @@ module Avatarable
 
       avatar_path(only_path: args.fetch(:only_path, true)) || super
     end
+
+    def retrieve_upload(identifier, paths)
+      upload = retrieve_upload_from_batch(identifier)
+
+      # This fallback is needed when deleting an upload, because we may have
+      # already been removed from the DB. We have to check an explicit `#nil?`
+      # because it's a BatchLoader instance.
+      upload = super if upload.nil?
+
+      upload
+    end
   end
 
   def avatar_type
@@ -52,4 +66,36 @@ module Avatarable
 
     url_base + avatar.local_url
   end
+
+  # Path that is persisted in the tracking Upload model. Used to fetch the
+  # upload from the model.
+  def upload_paths(identifier)
+    avatar_mounter.blank_uploader.store_dirs.map { |store, path| File.join(path, identifier) }
+  end
+
+  private
+
+  def retrieve_upload_from_batch(identifier)
+    BatchLoader.for(identifier: identifier, model: self).batch(key: self.class) do |upload_params, loader, args|
+      paths = upload_params.flat_map do |params|
+        params[:model].upload_paths(params[:identifier])
+      end
+
+      Upload.where(uploader: AvatarUploader, path: paths).find_each do |upload|
+        model = args[:key].instantiate('id' => upload.model_id)
+
+        loader.call({ model: model, identifier: File.basename(upload.path) }, upload)
+      end
+    end
+  end
+
+  def add_avatar_to_batch
+    return unless avatar_mounter
+
+    avatar_mounter.read_identifiers.each(&method(:retrieve_upload_from_batch))
+  end
+
+  def avatar_mounter
+    strong_memoize(:avatar_mounter) { _mounter(:avatar) }
+  end
 end
diff --git a/app/models/concerns/with_uploads.rb b/app/models/concerns/with_uploads.rb
index e7cfffb775b..4245d083a49 100644
--- a/app/models/concerns/with_uploads.rb
+++ b/app/models/concerns/with_uploads.rb
@@ -36,4 +36,8 @@ module WithUploads
       upload.destroy
     end
   end
+
+  def retrieve_upload(_identifier, paths)
+    uploads.find_by(path: paths)
+  end
 end
diff --git a/app/models/note.rb b/app/models/note.rb
index 02f7a9b1e4f..41c04ae0571 100644
--- a/app/models/note.rb
+++ b/app/models/note.rb
@@ -435,6 +435,10 @@ class Note < ActiveRecord::Base
     super.merge(noteable: noteable)
   end
 
+  def retrieve_upload(_identifier, paths)
+    Upload.find_by(model: self, path: paths)
+  end
+
   private
 
   def keep_around_commit
diff --git a/app/models/personal_snippet.rb b/app/models/personal_snippet.rb
index 82c1c4de3a0..355624fd552 100644
--- a/app/models/personal_snippet.rb
+++ b/app/models/personal_snippet.rb
@@ -1,2 +1,3 @@
 class PersonalSnippet < Snippet
+  include WithUploads
 end
diff --git a/app/uploaders/object_storage.rb b/app/uploaders/object_storage.rb
index 5bdca26a584..810f73af485 100644
--- a/app/uploaders/object_storage.rb
+++ b/app/uploaders/object_storage.rb
@@ -35,7 +35,7 @@ module ObjectStorage
 
         unless current_upload_satisfies?(paths, model)
           # the upload we already have isn't right, find the correct one
-          self.upload = uploads.find_by(model: model, path: paths)
+          self.upload = model&.retrieve_upload(identifier, paths)
         end
 
         super
@@ -48,7 +48,7 @@ module ObjectStorage
       end
 
       def upload=(upload)
-        return unless upload
+        return if upload.nil?
 
         self.object_store = upload.store
         super
diff --git a/changelogs/unreleased/fix-avatars-n-plus-one.yml b/changelogs/unreleased/fix-avatars-n-plus-one.yml
new file mode 100644
index 00000000000..c5b42071f2b
--- /dev/null
+++ b/changelogs/unreleased/fix-avatars-n-plus-one.yml
@@ -0,0 +1,5 @@
+---
+title: Fix an N+1 when loading user avatars
+merge_request:
+author:
+type: performance
diff --git a/spec/uploaders/object_storage_spec.rb b/spec/uploaders/object_storage_spec.rb
index 2dd0925a8e6..14d0ef8629a 100644
--- a/spec/uploaders/object_storage_spec.rb
+++ b/spec/uploaders/object_storage_spec.rb
@@ -659,4 +659,24 @@ describe ObjectStorage do
       end
     end
   end
+
+  describe '#retrieve_from_store!' do
+    [:group, :project, :user].each do |model|
+      context "for #{model}s" do
+        let(:models) { create_list(model, 10, :with_avatar).map(&:reload) }
+        let(:avatars) { models.map(&:avatar) }
+
+        it 'batches fetching uploads from the database' do
+          models
+
+          expect { avatars }.not_to exceed_query_limit(1)
+        end
+
+        it 'fetches the correct upload for each' do
+          expect(avatars.map(&:url).uniq).to eq(avatars.map(&:url))
+          expect(avatars.map(&:upload).uniq).to eq(avatars.map(&:upload))
+        end
+      end
+    end
+  end
 end
diff --git a/spec/uploaders/workers/object_storage/background_move_worker_spec.rb b/spec/uploaders/workers/object_storage/background_move_worker_spec.rb
index b34f427fd8a..acb7006db42 100644
--- a/spec/uploaders/workers/object_storage/background_move_worker_spec.rb
+++ b/spec/uploaders/workers/object_storage/background_move_worker_spec.rb
@@ -125,8 +125,10 @@ describe ObjectStorage::BackgroundMoveWorker do
 
           it "migrates file to remote storage" do
             perform
+            project.reload
+            BatchLoader::Executor.clear_current
 
-            expect(project.reload.avatar.file_storage?).to be_falsey
+            expect(project.avatar.file_storage?).to be_falsey
           end
         end