Merge branch '47408-migrateuploadsworker-is-doing-n-1-queries-on-migration' into 'master'

Resolve "`MigrateUploadsWorker` is doing N+1 queries on migration"

Closes #47408

See merge request gitlab-org/gitlab-ce!19547

app/uploaders/file_uploader.rb

@@ -65,10 +65,10 @@ class FileUploader < GitlabUploader
     SecureRandom.hex
   end
 
-  def upload_paths(filename)
+  def upload_paths(identifier)
     [
-      File.join(secret, filename),
-      File.join(base_dir(Store::REMOTE), secret, filename)
+      File.join(secret, identifier),
+      File.join(base_dir(Store::REMOTE), secret, identifier)
     ]
   end
 

app/uploaders/object_storage.rb

@@ -10,6 +10,17 @@ module ObjectStorage
   UnknownStoreError = Class.new(StandardError)
   ObjectStorageUnavailable = Class.new(StandardError)
 
+  class ExclusiveLeaseTaken < StandardError
+    def initialize(lease_key)
+      @lease_key = lease_key
+    end
+
+    def message
+      *lease_key_group, _ = *@lease_key.split(":")
+      "Exclusive lease for #{lease_key_group.join(':')} is already taken."
+    end
+  end
+
   TMP_UPLOAD_PATH = 'tmp/uploads'.freeze
 
   module Store
@@ -29,7 +40,7 @@ module ObjectStorage
       end
 
       def retrieve_from_store!(identifier)
-        paths = store_dirs.map { |store, path| File.join(path, identifier) }
+        paths = upload_paths(identifier)
 
         unless current_upload_satisfies?(paths, model)
           # the upload we already have isn't right, find the correct one
@@ -261,7 +272,7 @@ module ObjectStorage
     end
 
     def delete_migrated_file(migrated_file)
-      migrated_file.delete if exists?
+      migrated_file.delete
     end
 
     def exists?
@@ -279,6 +290,13 @@ module ObjectStorage
       }
     end
 
+    # Returns all the possible paths for an upload.
+    # the `upload.path` is a lookup parameter, and it may change
+    # depending on the `store` param.
+    def upload_paths(identifier)
+      store_dirs.map { |store, path| File.join(path, identifier) }
+    end
+
     def cache!(new_file = sanitized_file)
       # We intercept ::UploadedFile which might be stored on remote storage
       # We use that for "accelerated" uploads, where we store result on remote storage
@@ -369,12 +387,13 @@ module ObjectStorage
     end
 
     def with_exclusive_lease
-      uuid = Gitlab::ExclusiveLease.new(exclusive_lease_key, timeout: 1.hour.to_i).try_obtain
-      raise 'exclusive lease already taken' unless uuid
+      lease_key = exclusive_lease_key
+      uuid = Gitlab::ExclusiveLease.new(lease_key, timeout: 1.hour.to_i).try_obtain
+      raise ExclusiveLeaseTaken.new(lease_key) unless uuid
 
       yield uuid
     ensure
-      Gitlab::ExclusiveLease.cancel(exclusive_lease_key, uuid)
+      Gitlab::ExclusiveLease.cancel(lease_key, uuid)
     end
 
     #

app/uploaders/records_uploads.rb

@@ -22,7 +22,7 @@ module RecordsUploads
 
       Upload.transaction do
         uploads.where(path: upload_path).delete_all
-        upload.destroy! if upload
+        upload.delete if upload
 
         self.upload = build_upload.tap(&:save!)
       end

changelogs/unreleased/47408-migrateuploadsworker-is-doing-n-1-queries-on-migration.yml

+---
+title: Optimize the upload migration proces
+merge_request: 15947
+author:
+type: fixed

spec/support/shared_examples/uploaders/object_storage_shared_examples.rb

@@ -85,13 +85,13 @@ shared_examples "migrates" do |to_store:, from_store: nil|
     it 'does not execute migrate!' do
       expect(subject).not_to receive(:unsafe_migrate!)
 
-      expect { migrate(to) }.to raise_error('exclusive lease already taken')
+      expect { migrate(to) }.to raise_error(ObjectStorage::ExclusiveLeaseTaken)
     end
 
     it 'does not execute use_file' do
       expect(subject).not_to receive(:unsafe_use_file)
 
-      expect { subject.use_file }.to raise_error('exclusive lease already taken')
+      expect { subject.use_file }.to raise_error(ObjectStorage::ExclusiveLeaseTaken)
     end
 
     after do

spec/uploaders/object_storage_spec.rb

@@ -321,7 +321,7 @@ describe ObjectStorage do
       when_file_is_in_use do
         expect(uploader).not_to receive(:unsafe_migrate!)
 
-        expect { uploader.migrate!(described_class::Store::REMOTE) }.to raise_error('exclusive lease already taken')
+        expect { uploader.migrate!(described_class::Store::REMOTE) }.to raise_error(ObjectStorage::ExclusiveLeaseTaken)
       end
     end
 
@@ -329,7 +329,7 @@ describe ObjectStorage do
       when_file_is_in_use do
         expect(uploader).not_to receive(:unsafe_use_file)
 
-        expect { uploader.use_file }.to raise_error('exclusive lease already taken')
+        expect { uploader.use_file }.to raise_error(ObjectStorage::ExclusiveLeaseTaken)
       end
     end
   end

spec/uploaders/workers/object_storage/migrate_uploads_worker_spec.rb

@@ -11,6 +11,12 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
   let(:uploads) { Upload.all }
   let(:to_store) { ObjectStorage::Store::REMOTE }
 
+  def perform(uploads)
+    described_class.new.perform(uploads.ids, model_class.to_s, mounted_as, to_store)
+  rescue ObjectStorage::MigrateUploadsWorker::Report::MigrationFailures
+    # swallow
+  end
+
   shared_examples "uploads migration worker" do
     describe '.enqueue!' do
       def enqueue!
@@ -69,12 +75,6 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
     end
 
     describe '#perform' do
-      def perform
-        described_class.new.perform(uploads.ids, model_class.to_s, mounted_as, to_store)
-      rescue ObjectStorage::MigrateUploadsWorker::Report::MigrationFailures
-        # swallow
-      end
-
       shared_examples 'outputs correctly' do |success: 0, failures: 0|
         total = success + failures
 
@@ -82,7 +82,7 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
           it 'outputs the reports' do
             expect(Rails.logger).to receive(:info).with(%r{Migrated #{success}/#{total} files})
 
-            perform
+            perform(uploads)
           end
         end
 
@@ -90,7 +90,7 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
           it 'outputs upload failures' do
             expect(Rails.logger).to receive(:warn).with(/Error .* I am a teapot/)
 
-            perform
+            perform(uploads)
           end
         end
       end
@@ -98,7 +98,7 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
       it_behaves_like 'outputs correctly', success: 10
 
       it 'migrates files' do
-        perform
+        perform(uploads)
 
         expect(Upload.where(store: ObjectStorage::Store::LOCAL).count).to eq(0)
       end
@@ -123,6 +123,17 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
     end
 
     it_behaves_like "uploads migration worker"
+
+    describe "limits N+1 queries" do
+      it "to N*5" do
+        query_count = ActiveRecord::QueryRecorder.new { perform(uploads) }
+
+        more_projects = create_list(:project, 3, :with_avatar)
+
+        expected_queries_per_migration = 5 * more_projects.count
+        expect { perform(Upload.all) }.not_to exceed_query_limit(query_count).with_threshold(expected_queries_per_migration)
+      end
+    end
   end
 
   context "for FileUploader" do
@@ -130,15 +141,29 @@ describe ObjectStorage::MigrateUploadsWorker, :sidekiq do
     let(:secret) { SecureRandom.hex }
     let(:mounted_as) { nil }
 
+    def upload_file(project)
+      uploader = FileUploader.new(project)
+      uploader.store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
+    end
+
     before do
       stub_uploads_object_storage(FileUploader)
 
-      projects.map do |project|
-        uploader = FileUploader.new(project)
-        uploader.store!(fixture_file_upload('spec/fixtures/doc_sample.txt'))
-      end
+      projects.map(&method(:upload_file))
     end
 
     it_behaves_like "uploads migration worker"
+
+    describe "limits N+1 queries" do
+      it "to N*5" do
+        query_count = ActiveRecord::QueryRecorder.new { perform(uploads) }
+
+        more_projects = create_list(:project, 3)
+        more_projects.map(&method(:upload_file))
+
+        expected_queries_per_migration = 5 * more_projects.count
+        expect { perform(Upload.all) }.not_to exceed_query_limit(query_count).with_threshold(expected_queries_per_migration)
+      end
+    end
   end
 end