Merge branch '206946-registry-is-ssot-for-uploads' into 'master'

Geo - Make registry table SSOT for Uploads

See merge request gitlab-org/gitlab!25482

ee/app/finders/geo/attachment_registry_finder.rb

@@ -34,64 +34,119 @@ module Geo
       Upload
     end
 
+    # Returns untracked IDs as well as tracked IDs that are unused.
+    #
+    # Untracked IDs are model IDs that are supposed to be synced but don't yet
+    # have a registry entry.
+    #
+    # Unused tracked IDs are model IDs that are not supposed to be synced but
+    # already have a registry entry. For example:
+    #
+    #   - orphaned registries
+    #   - records that became excluded from selective sync
+    #   - records that are in object storage, and `sync_object_storage` became
+    #     disabled
+    #
+    # We compute both sets in this method to reduce the number of DB queries
+    # performed.
+    #
+    # @return [Array] the first element is an Array of untracked IDs, and the second element is an Array of tracked IDs that are unused
+    def find_registry_differences(range)
+      source = attachments(fdw: false).where(id: range).pluck(::Upload.arel_table[:id], ::Upload.arel_table[:uploader]) # rubocop:disable CodeReuse/ActiveRecord
+      tracked = Geo::UploadRegistry.where(file_id: range).pluck(:file_id, :file_type) # rubocop:disable CodeReuse/ActiveRecord
+
+      untracked = source - tracked
+      unused_tracked = tracked - source
+
+      [untracked, unused_tracked]
+    end
+
+    # Returns Geo::UploadRegistry records that have never been synced.
+    #
+    # Does not care about selective sync, because it considers the Registry
+    # table to be the single source of truth. The contract is that other
+    # processes need to ensure that the table only contains records that should
+    # be synced.
+    #
+    # Any registries that have ever been synced that currently need to be
+    # resynced will be handled by other find methods (like
+    # #find_retryable_failed_registries)
+    #
+    # You can pass a list with `except_ids:` so you can exclude items you
+    # already scheduled but haven't finished and aren't persisted to the database yet
+    #
+    # @param [Integer] batch_size used to limit the results returned
+    # @param [Array<Integer>] except_ids ids that will be ignored from the query
+    # rubocop:disable CodeReuse/ActiveRecord
+    def find_never_synced_registries(batch_size:, except_ids: [])
+      Geo::UploadRegistry
+        .never
+        .model_id_not_in(except_ids)
+        .limit(batch_size)
+    end
+    # rubocop:enable CodeReuse/ActiveRecord
+
+    # Deprecated in favor of the process using
+    # #find_registry_differences and #find_never_synced_registries
+    #
     # Find limited amount of non replicated attachments.
     #
-    # You can pass a list with `except_file_ids:` so you can exclude items you
+    # You can pass a list with `except_ids:` so you can exclude items you
     # already scheduled but haven't finished and aren't persisted to the database yet
     #
     # TODO: Alternative here is to use some sort of window function with a cursor instead
     #       of simply limiting the query and passing a list of items we don't want
     #
     # @param [Integer] batch_size used to limit the results returned
-    # @param [Array<Integer>] except_file_ids ids that will be ignored from the query
+    # @param [Array<Integer>] except_ids ids that will be ignored from the query
     # rubocop: disable CodeReuse/ActiveRecord
-    def find_unsynced(batch_size:, except_file_ids: [])
+    def find_unsynced(batch_size:, except_ids: [])
       attachments
         .missing_registry
-        .id_not_in(except_file_ids)
+        .id_not_in(except_ids)
         .limit(batch_size)
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
     # rubocop: disable CodeReuse/ActiveRecord
-    def find_migrated_local(batch_size:, except_file_ids: [])
+    def find_migrated_local(batch_size:, except_ids: [])
       all_attachments
         .inner_join_registry
         .with_files_stored_remotely
-        .id_not_in(except_file_ids)
+        .id_not_in(except_ids)
         .limit(batch_size)
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
     # rubocop: disable CodeReuse/ActiveRecord
-    def find_retryable_failed_registries(batch_size:, except_file_ids: [])
+    def find_retryable_failed_registries(batch_size:, except_ids: [])
       Geo::UploadRegistry
         .failed
         .retry_due
-        .file_id_not_in(except_file_ids)
+        .model_id_not_in(except_ids)
         .limit(batch_size)
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
     # rubocop: disable CodeReuse/ActiveRecord
-    def find_retryable_synced_missing_on_primary_registries(batch_size:, except_file_ids: [])
+    def find_retryable_synced_missing_on_primary_registries(batch_size:, except_ids: [])
       Geo::UploadRegistry
         .synced
         .missing_on_primary
         .retry_due
-        .file_id_not_in(except_file_ids)
+        .model_id_not_in(except_ids)
         .limit(batch_size)
     end
     # rubocop: enable CodeReuse/ActiveRecord
 
     private
 
-    def attachments
-      local_storage_only? ? all_attachments.with_files_stored_locally : all_attachments
+    def attachments(fdw: true)
+      local_storage_only?(fdw: fdw) ? all_attachments(fdw: fdw).with_files_stored_locally : all_attachments(fdw: fdw)
     end
 
-    def all_attachments
-      current_node.attachments
+    def all_attachments(fdw: true)
+      current_node(fdw: fdw).attachments
     end
 
     def registries_for_attachments

ee/app/models/geo/upload_registry.rb

@@ -3,6 +3,9 @@
 class Geo::UploadRegistry < Geo::BaseRegistry
   include Geo::Syncable
 
+  MODEL_CLASS = ::Upload
+  MODEL_FOREIGN_KEY = :file_id
+
   self.table_name = 'file_registry'
 
   belongs_to :upload, foreign_key: :file_id
@@ -11,16 +14,27 @@ class Geo::UploadRegistry < Geo::BaseRegistry
   scope :fresh, -> { order(created_at: :desc) }
   scope :never, -> { where(success: false, retry_count: nil) }
 
-  def self.file_id_in(ids)
-    where(file_id: ids)
+  def self.registry_consistency_worker_enabled?
+    Feature.enabled?(:geo_file_registry_ssot_sync)
+  end
+
+  def self.finder_class
+    ::Geo::AttachmentRegistryFinder
   end
 
-  def self.file_id_not_in(ids)
-    where.not(file_id: ids)
+  # If false, RegistryConsistencyService will frequently check the end of the
+  # table to quickly handle new replicables.
+  def self.has_create_events?
+    false
   end
 
-  def self.pluck_file_key
-    where(nil).pluck(:file_id)
+  # TODO: Investigate replacing this with bulk insert (there was an obstacle).
+  #       https://gitlab.com/gitlab-org/gitlab/issues/197310
+  def self.insert_for_model_ids(attrs)
+    attrs.map do |file_id, file_type|
+      registry = create(file_id: file_id, file_type: file_type)
+      registry.id
+    end.compact
   end
 
   def self.with_search(query)

ee/app/workers/geo/file_download_dispatch_worker/attachment_job_finder.rb

@@ -3,12 +3,22 @@
 module Geo
   class FileDownloadDispatchWorker # rubocop:disable Scalability/IdempotentWorker
     class AttachmentJobFinder < JobFinder # rubocop:disable Scalability/IdempotentWorker
-      EXCEPT_RESOURCE_IDS_KEY = :except_file_ids
+      EXCEPT_RESOURCE_IDS_KEY = :except_ids
 
       def registry_finder
         @registry_finder ||= Geo::AttachmentRegistryFinder.new(current_node_id: Gitlab::Geo.current_node.id)
       end
 
+      def find_unsynced_jobs(batch_size:)
+        if Geo::UploadRegistry.registry_consistency_worker_enabled?
+          convert_registry_relation_to_job_args(
+            registry_finder.find_never_synced_registries(find_batch_params(batch_size))
+          )
+        else
+          super
+        end
+      end
+
       private
 
       # Why do we need a different `file_type` for each Uploader? Why not just use 'upload'?

ee/app/workers/geo/file_download_dispatch_worker/lfs_object_job_finder.rb

@@ -12,7 +12,7 @@ module Geo
       end
 
       def find_unsynced_jobs(batch_size:)
-        if Feature.enabled?(:geo_lfs_registry_ssot_sync)
+        if Geo::LfsObjectRegistry.registry_consistency_worker_enabled?
           convert_registry_relation_to_job_args(
             registry_finder.find_never_synced_registries(find_batch_params(batch_size))
           )

ee/app/workers/geo/migrated_local_files_clean_up_worker.rb

@@ -58,7 +58,7 @@ module Geo
     def find_migrated_local_attachments_ids(batch_size:)
       return [] unless attachments_object_store_enabled?
 
-      attachments_finder.find_migrated_local(batch_size: batch_size, except_file_ids: scheduled_file_ids(Gitlab::Geo::Replication::USER_UPLOADS_OBJECT_TYPES))
+      attachments_finder.find_migrated_local(batch_size: batch_size, except_ids: scheduled_file_ids(Gitlab::Geo::Replication::USER_UPLOADS_OBJECT_TYPES))
                         .pluck(:uploader, :id)
                         .map { |uploader, id| [uploader.sub(/Uploader\z/, '').underscore, id] }
     end

ee/app/workers/geo/secondary/registry_consistency_worker.rb

@@ -16,7 +16,7 @@ module Geo
       feature_category :geo_replication
 
       # This is probably not the best place to "register" replicables for this functionality
-      REGISTRY_CLASSES = [Geo::JobArtifactRegistry, Geo::LfsObjectRegistry].freeze
+      REGISTRY_CLASSES = [Geo::JobArtifactRegistry, Geo::LfsObjectRegistry, Geo::UploadRegistry].freeze
       BATCH_SIZE = 1000
 
       # @return [Boolean] true if at least 1 registry was created, else false

ee/spec/factories/geo/upload_registry.rb

@@ -19,6 +19,11 @@ FactoryBot.define do
       retry_count { 1 }
     end
 
+    trait :never_synced do
+      success { false }
+      retry_count { nil }
+    end
+
     trait :with_file do
       after(:build, :stub) do |registry, _|
         file =

ee/spec/finders/geo/attachment_registry_finder_spec.rb

@@ -53,7 +53,7 @@ describe Geo::AttachmentRegistryFinder, :geo, :geo_fdw do
         end
 
         it 'returns attachments without an entry on the tracking database, excluding from exception list' do
-          attachments = subject.find_unsynced(batch_size: 10, except_file_ids: [upload_issuable_synced_nested_project.id])
+          attachments = subject.find_unsynced(batch_size: 10, except_ids: [upload_issuable_synced_nested_project.id])
 
           expect(attachments).to match_ids(upload_unsynced_project, upload_synced_project, upload_personal_snippet,
                                            upload_remote_unsynced_project, upload_remote_synced_group)
@@ -64,7 +64,7 @@ describe Geo::AttachmentRegistryFinder, :geo, :geo_fdw do
         let(:secondary) { create(:geo_node, :local_storage_only) }
 
         it 'returns local attachments only' do
-          attachments = subject.find_unsynced(batch_size: 10, except_file_ids: [upload_synced_project.id])
+          attachments = subject.find_unsynced(batch_size: 10, except_ids: [upload_synced_project.id])
 
           expect(attachments).to match_ids(upload_issuable_synced_nested_project, upload_unsynced_project,
                                            upload_personal_snippet)
@@ -75,7 +75,7 @@ describe Geo::AttachmentRegistryFinder, :geo, :geo_fdw do
         let(:secondary) { create(:geo_node, selective_sync_type: 'namespaces', namespaces: [synced_group]) }
 
         it 'returns attachments without an entry on the tracking database, excluding from exception list' do
-          attachments = subject.find_unsynced(batch_size: 10, except_file_ids: [upload_synced_project.id])
+          attachments = subject.find_unsynced(batch_size: 10, except_ids: [upload_synced_project.id])
 
           expect(attachments).to match_ids(upload_issuable_synced_nested_project, upload_personal_snippet,
                                            upload_remote_synced_group)
@@ -101,7 +101,7 @@ describe Geo::AttachmentRegistryFinder, :geo, :geo_fdw do
       end
 
       it 'returns attachments stored remotely and successfully synced locally' do
-        attachments = subject.find_migrated_local(batch_size: 100, except_file_ids: [upload_remote_unsynced_project.id])
+        attachments = subject.find_migrated_local(batch_size: 100, except_ids: [upload_remote_unsynced_project.id])
 
         expect(attachments).to match_ids(upload_remote_synced_project)
       end

ee/spec/workers/geo/file_download_dispatch_worker_spec.rb

@@ -62,6 +62,120 @@ describe Geo::FileDownloadDispatchWorker, :geo, :geo_fdw do
   context 'with attachments (Upload records)' do
     let(:upload) { create(:upload) }
 
+    context 'with geo_file_registry_ssot_sync feature enabled' do
+      before do
+        stub_feature_flags(geo_file_registry_ssot_sync: true)
+      end
+
+      it 'performs Geo::FileDownloadWorker for unsynced attachments' do
+        create(:geo_upload_registry, :avatar, :never_synced, file_id: upload.id)
+
+        expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', upload.id)
+
+        subject.perform
+      end
+
+      it 'performs Geo::FileDownloadWorker for failed-sync attachments' do
+        create(:geo_upload_registry, :avatar, :failed, file_id: upload.id, bytes: 0)
+
+        expect(Geo::FileDownloadWorker).to receive(:perform_async)
+          .with('avatar', upload.id).once.and_return(spy)
+
+        subject.perform
+      end
+
+      it 'does not perform Geo::FileDownloadWorker for synced attachments' do
+        create(:geo_upload_registry, :avatar, file_id: upload.id, bytes: 1234)
+
+        expect(Geo::FileDownloadWorker).not_to receive(:perform_async)
+
+        subject.perform
+      end
+
+      it 'does not perform Geo::FileDownloadWorker for synced attachments even with 0 bytes downloaded' do
+        create(:geo_upload_registry, :avatar, file_id: upload.id, bytes: 0)
+
+        expect(Geo::FileDownloadWorker).not_to receive(:perform_async)
+
+        subject.perform
+      end
+
+      context 'with a failed file' do
+        let(:failed_registry) { create(:geo_upload_registry, :avatar, :failed, file_id: 999) }
+
+        it 'does not stall backfill' do
+          unsynced_registry = create(:geo_upload_registry, :avatar, :with_file, :never_synced)
+
+          stub_const('Geo::Scheduler::SchedulerWorker::DB_RETRIEVE_BATCH_SIZE', 1)
+
+          expect(Geo::FileDownloadWorker).not_to receive(:perform_async).with('avatar', failed_registry.file_id)
+          expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', unsynced_registry.file_id)
+
+          subject.perform
+        end
+
+        it 'retries failed files' do
+          expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', failed_registry.file_id)
+
+          subject.perform
+        end
+
+        it 'does not retry failed files when retry_at is tomorrow' do
+          failed_registry = create(:geo_upload_registry, :avatar, :failed, file_id: 999, retry_at: Date.tomorrow)
+
+          expect(Geo::FileDownloadWorker).not_to receive(:perform_async).with('avatar', failed_registry.file_id)
+
+          subject.perform
+        end
+
+        it 'retries failed files when retry_at is in the past' do
+          failed_registry = create(:geo_upload_registry, :avatar, :failed, file_id: 999, retry_at: Date.yesterday)
+
+          expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', failed_registry.file_id)
+
+          subject.perform
+        end
+      end
+
+      context 'with Upload files missing on the primary that are marked as synced' do
+        let(:synced_upload_with_file_missing_on_primary) { create(:upload) }
+
+        before do
+          Geo::UploadRegistry.create!(file_type: :avatar, file_id: synced_upload_with_file_missing_on_primary.id, bytes: 1234, success: true, missing_on_primary: true)
+        end
+
+        it 'retries the files if there is spare capacity' do
+          expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', synced_upload_with_file_missing_on_primary.id)
+
+          subject.perform
+        end
+
+        it 'does not retry those files if there is no spare capacity' do
+          unsynced_registry = create(:geo_upload_registry, :avatar, :with_file, :never_synced)
+          expect(subject).to receive(:db_retrieve_batch_size).and_return(1).twice
+
+          expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', unsynced_registry.file_id)
+
+          subject.perform
+        end
+
+        it 'does not retry those files if they are already scheduled' do
+          unsynced_registry = create(:geo_upload_registry, :avatar, :with_file, :never_synced)
+
+          scheduled_jobs = [{ type: 'avatar', id: synced_upload_with_file_missing_on_primary.id, job_id: 'foo' }]
+          expect(subject).to receive(:scheduled_jobs).and_return(scheduled_jobs).at_least(1)
+          expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', unsynced_registry.file_id)
+
+          subject.perform
+        end
+      end
+    end
+
+    context 'with geo_file_registry_ssot_sync feature disabled' do
+      before do
+        stub_feature_flags(geo_file_registry_ssot_sync: false)
+      end
+
       it 'performs Geo::FileDownloadWorker for unsynced attachments' do
         expect(Geo::FileDownloadWorker).to receive(:perform_async).with('avatar', upload.id)
 
@@ -163,6 +277,7 @@ describe Geo::FileDownloadDispatchWorker, :geo, :geo_fdw do
         end
       end
     end
+  end
 
   context 'with LFS objects' do
     let!(:lfs_object_local_store) { create(:lfs_object, :with_file) }
@@ -534,13 +649,13 @@ describe Geo::FileDownloadDispatchWorker, :geo, :geo_fdw do
     result_object = double(:result, success: true, bytes_downloaded: 100, primary_missing_file: false)
     allow_any_instance_of(::Gitlab::Geo::Replication::BaseTransfer).to receive(:download_from_primary).and_return(result_object)
 
-    avatar = fixture_file_upload('spec/fixtures/dk.png')
     create_list(:geo_lfs_object_registry, 2, :with_lfs_object, :never_synced)
-    create_list(:user, 2, avatar: avatar)
-    create_list(:note, 2, :with_attachment)
-    create(:upload, :personal_snippet_upload)
+    create_list(:geo_upload_registry, 2, :avatar, :with_file, :never_synced)
+    create_list(:geo_upload_registry, 2, :attachment, :with_file, :never_synced)
+    create(:geo_upload_registry, :favicon, :with_file, :never_synced)
+    create(:geo_upload_registry, :import_export, :with_file, :never_synced)
+    create(:geo_upload_registry, :personal_file, :with_file, :never_synced)
     create(:geo_job_artifact_registry, :with_artifact, :never_synced)
-    create(:appearance, logo: avatar, header_logo: avatar)
 
     expect(Geo::FileDownloadWorker).to receive(:perform_async).exactly(10).times.and_call_original
     # For 10 downloads, we expect four database reloads:
@@ -602,6 +717,11 @@ describe Geo::FileDownloadDispatchWorker, :geo, :geo_fdw do
       end
     end
 
+    context 'with geo_file_registry_ssot_sync feature disabled' do
+      before do
+        stub_feature_flags(geo_file_registry_ssot_sync: false)
+      end
+
       it 'does not perform Geo::FileDownloadWorker for upload objects that do not belong to selected namespaces to replicate' do
         avatar = fixture_file_upload('spec/fixtures/dk.png')
         avatar_in_synced_group = create(:upload, model: synced_group, path: avatar)
@@ -618,4 +738,5 @@ describe Geo::FileDownloadDispatchWorker, :geo, :geo_fdw do
         subject.perform
       end
     end
+  end
 end

ee/spec/workers/geo/secondary/registry_consistency_worker_spec.rb

@@ -25,8 +25,6 @@ describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
   end
 
   describe '#perform' do
-    subject { described_class.new }
-
     before do
       allow(subject).to receive(:sleep) # faster tests
     end
@@ -80,14 +78,17 @@ describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
     it 'creates missing registries for each registry class' do
       lfs_object = create(:lfs_object)
       job_artifact = create(:ci_job_artifact)
+      upload = create(:upload)
 
       expect(Geo::LfsObjectRegistry.where(lfs_object_id: lfs_object.id).count).to eq(0)
       expect(Geo::JobArtifactRegistry.where(artifact_id: job_artifact.id).count).to eq(0)
+      expect(Geo::UploadRegistry.where(file_id: upload.id).count).to eq(0)
 
       subject.perform
 
       expect(Geo::LfsObjectRegistry.where(lfs_object_id: lfs_object.id).count).to eq(1)
       expect(Geo::JobArtifactRegistry.where(artifact_id: job_artifact.id).count).to eq(1)
+      expect(Geo::UploadRegistry.where(file_id: upload.id).count).to eq(1)
     end
 
     context 'when geo_lfs_registry_ssot_sync is disabled' do
@@ -103,6 +104,7 @@ describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
 
       it 'does not execute RegistryConsistencyService for LFS objects' do
         allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::JobArtifactRegistry, batch_size: 1000).and_call_original
+        allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::UploadRegistry, batch_size: 1000).and_call_original
 
         expect(Geo::RegistryConsistencyService).not_to receive(:new).with(Geo::LfsObjectRegistry, batch_size: 1000)
 
@@ -123,6 +125,7 @@ describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
 
       it 'does not execute RegistryConsistencyService for Job Artifacts' do
         allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::LfsObjectRegistry, batch_size: 1000).and_call_original
+        allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::UploadRegistry, batch_size: 1000).and_call_original
 
         expect(Geo::RegistryConsistencyService).not_to receive(:new).with(Geo::JobArtifactRegistry, batch_size: 1000)
 
@@ -130,6 +133,27 @@ describe Geo::Secondary::RegistryConsistencyWorker, :geo, :geo_fdw do
       end
     end
 
+    context 'when geo_file_registry_ssot_sync is disabled' do
+      let_it_be(:upload) { create(:upload) }
+
+      before do
+        stub_feature_flags(geo_file_registry_ssot_sync: false)
+      end
+
+      it 'returns false' do
+        expect(subject.perform).to be_falsey
+      end
+
+      it 'does not execute RegistryConsistencyService for Uploads' do
+        allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::JobArtifactRegistry, batch_size: 1000).and_call_original
+        allow(Geo::RegistryConsistencyService).to receive(:new).with(Geo::LfsObjectRegistry, batch_size: 1000).and_call_original
+
+        expect(Geo::RegistryConsistencyService).not_to receive(:new).with(Geo::UploadRegistry, batch_size: 1000)
+
+        subject.perform
+      end
+    end
+
     context 'when the current Geo node is disabled or primary' do
       before do
         stub_primary_node