Correctly preserve LFS objects in design or wiki repositories

The gitlab:cleanup:orphan_lfs_file_references rake task forgets to take
account of design and wiki repositories, both of which contain LFS
objects. Add them to the list of repositories we get OIDs from.

app/models/lfs_objects_project.rb

@@ -19,6 +19,7 @@ class LfsObjectsProject < ApplicationRecord
   }
 
   scope :project_id_in, ->(ids) { where(project_id: ids) }
+  scope :lfs_object_in, -> (lfs_objects) { where(lfs_object: lfs_objects) }
 
   private
 

app/models/project.rb

@@ -2702,9 +2702,11 @@ class Project < ApplicationRecord
   end
 
   def oids(objects, oids: [])
-    collection = oids.any? ? objects.where(oid: oids) : objects
+    objects = objects.where(oid: oids) if oids.any?
 
-    collection.pluck(:oid)
+    [].tap do |out|
+      objects.each_batch { |relation| out.concat(relation.pluck(:oid)) }
+    end
   end
 end
 

changelogs/unreleased/244848-fix-lfs-cleanup-data-loss-bug.yml

+---
+title: Correctly preserve LFS objects in design or wiki repositories
+merge_request: 41352
+author:
+type: fixed

lib/gitlab/cleanup/orphan_lfs_file_references.rb

@@ -25,7 +25,7 @@ module Gitlab
       private
 
       def remove_orphan_references
-        invalid_references = project.lfs_objects_projects.where(lfs_object: orphan_objects) # rubocop:disable CodeReuse/ActiveRecord
+        invalid_references = project.lfs_objects_projects.lfs_object_in(orphan_objects)
 
         if dry_run
           log_info("Found invalid references: #{invalid_references.count}")
@@ -41,26 +41,22 @@ module Gitlab
         end
       end
 
-      def lfs_oids_from_repository
-        project.repository.gitaly_blob_client.get_all_lfs_pointers.map(&:lfs_oid)
-      end
-
-      def orphan_oids
-        lfs_oids_from_database - lfs_oids_from_repository
-      end
+      def orphan_objects
+        # Get these first so racing with a git push can't remove any LFS objects
+        oids = project.lfs_objects_oids
 
-      def lfs_oids_from_database
-        oids = []
+        repos = [
+          project.repository,
+          project.design_repository,
+          project.wiki.repository
+        ].select(&:exists?)
 
-        project.lfs_objects.each_batch do |relation|
-          oids += relation.pluck(:oid) # rubocop:disable CodeReuse/ActiveRecord
+        repos.flat_map do |repo|
+          oids -= repo.gitaly_blob_client.get_all_lfs_pointers.map(&:lfs_oid)
         end
 
-        oids
-      end
-
-      def orphan_objects
-        LfsObject.where(oid: orphan_oids) # rubocop:disable CodeReuse/ActiveRecord
+        # The remaining OIDs are not used by any repository, so are orphans
+        LfsObject.for_oids(oids)
       end
 
       def log_info(msg)

spec/lib/gitlab/cleanup/orphan_lfs_file_references_spec.rb

@@ -9,6 +9,8 @@ RSpec.describe Gitlab::Cleanup::OrphanLfsFileReferences do
 
   let!(:invalid_reference) { create(:lfs_objects_project, project: project, lfs_object: lfs_object) }
 
+  subject(:service) { described_class.new(project, logger: null_logger, dry_run: dry_run) }
+
   before do
     allow(null_logger).to receive(:info)
 
@@ -21,25 +23,66 @@ RSpec.describe Gitlab::Cleanup::OrphanLfsFileReferences do
   end
 
   context 'dry run' do
+    let(:dry_run) { true }
+
     it 'prints messages and does not delete references' do
       expect(null_logger).to receive(:info).with("[DRY RUN] Looking for orphan LFS files for project #{project.name_with_namespace}")
       expect(null_logger).to receive(:info).with("[DRY RUN] Found invalid references: 1")
 
-      expect { described_class.new(project, logger: null_logger).run! }
-        .not_to change { project.lfs_objects.count }
+      expect { service.run! }.not_to change { project.lfs_objects.count }
     end
   end
 
   context 'regular run' do
+    let(:dry_run) { false }
+
     it 'prints messages and deletes invalid reference' do
       expect(null_logger).to receive(:info).with("Looking for orphan LFS files for project #{project.name_with_namespace}")
       expect(null_logger).to receive(:info).with("Removed invalid references: 1")
       expect(ProjectCacheWorker).to receive(:perform_async).with(project.id, [], [:lfs_objects_size])
 
-      expect { described_class.new(project, logger: null_logger, dry_run: false).run! }
-        .to change { project.lfs_objects.count }.from(2).to(1)
+      expect { service.run! }.to change { project.lfs_objects.count }.from(2).to(1)
 
       expect(LfsObjectsProject.exists?(invalid_reference.id)).to be_falsey
     end
+
+    context 'LFS object is in design repository' do
+      before do
+        expect(project.design_repository).to receive(:exists?).and_return(true)
+
+        stub_lfs_pointers(project.design_repository, lfs_object.oid)
+      end
+
+      it 'is not removed' do
+        expect { service.run! }.not_to change { project.lfs_objects.count }
+      end
+    end
+
+    context 'LFS object is in wiki repository' do
+      before do
+        expect(project.wiki.repository).to receive(:exists?).and_return(true)
+
+        stub_lfs_pointers(project.wiki.repository, lfs_object.oid)
+      end
+
+      it 'is not removed' do
+        expect { service.run! }.not_to change { project.lfs_objects.count }
+      end
+    end
+  end
+
+  context 'LFS for project snippets' do
+    let(:snippet) { create(:project_snippet) }
+
+    it 'is disabled' do
+      # Support project snippets here before enabling LFS for them
+      expect(snippet.repository.lfs_enabled?).to be_falsy
+    end
+  end
+
+  def stub_lfs_pointers(repo, *oids)
+    expect(repo.gitaly_blob_client)
+      .to receive(:get_all_lfs_pointers)
+      .and_return(oids.map { |oid| OpenStruct.new(lfs_oid: oid) })
   end
 end