From bd818902272f44dc6a1f25c04697b4d93b705c99 Mon Sep 17 00:00:00 2001
From: Luke Duncalfe <lduncalfe@eml.cc>
Date: Mon, 8 Jul 2024 10:37:32 +1200
Subject: [PATCH 1/2] Allow Import to check placeholder loading progress

https://gitlab.com/gitlab-org/gitlab/-/merge_requests/156704 introduced
the batch loading of import placeholder references from Redis to
PostgreSQL.

This change allows an import to check if there are values in the Redis
set.

The worker that finalises an import will include
`Import::WaitsForPlaceholderReferences`, and specify a `import_source`,
and `import_uid` that must match the values passed to
`Import::PlaceholderReferences::PushService`.

Before finishing the import the worker can call `#pending_references?`
and choose to requeue.

Here's an example of checking the progress, and retrying the finalising
import worker again later (we may as well queue another
`LoadPlaceholderContributionsWorker`):

```ruby
class MyWorker
  include ApplicationWorker
  include ::Import::WaitsForPlaceholderReferences

  REQUEUE_DELAY = 10.seconds

  import_source :example_source
  import_uid -> { project.id }

  def perform(project_id)
    @project = Project.find(project_id)

    if pending_references?
      info(
        message: 'Re-enqueuing, waiting for placeholder references',
        pending_references_count: pending_references_count
      )

      Gitlab::Import::LoadPlaceholderContributionsWorker
        .perform_async(:example_source, project.id)

      return self.class.perform_in(REQUEUE_DELAY, project.id)
    end

    # Do import finalization
  end

  private

  attr_reader :project
end
```

https://gitlab.com/gitlab-org/gitlab/-/issues/467511
---
 app/services/import/placeholder_references.rb | 11 ++++
 .../placeholder_references/base_service.rb    |  2 +-
 .../waits_for_placeholder_references.rb       | 42 +++++++++++++
 lib/gitlab/cache/import/caching.rb            | 11 ++++
 .../load_service_spec.rb                      |  4 +-
 .../import/placeholder_references_spec.rb     | 13 ++++
 .../waits_for_placeholder_references_spec.rb  | 59 +++++++++++++++++++
 7 files changed, 138 insertions(+), 4 deletions(-)
 create mode 100644 app/services/import/placeholder_references.rb
 create mode 100644 app/workers/concerns/import/waits_for_placeholder_references.rb
 create mode 100644 spec/services/import/placeholder_references_spec.rb
 create mode 100644 spec/workers/concerns/import/waits_for_placeholder_references_spec.rb

diff --git a/app/services/import/placeholder_references.rb b/app/services/import/placeholder_references.rb
new file mode 100644
index 00000000000000..4328ff045bf3de
--- /dev/null
+++ b/app/services/import/placeholder_references.rb
@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+module Import
+  module PlaceholderReferences
+    class << self
+      def cache_key(import_source, import_uid)
+        [:'placeholder-references', import_source, import_uid].join(':')
+      end
+    end
+  end
+end
diff --git a/app/services/import/placeholder_references/base_service.rb b/app/services/import/placeholder_references/base_service.rb
index 2aa083d796a1eb..54e0f43212dbc9 100644
--- a/app/services/import/placeholder_references/base_service.rb
+++ b/app/services/import/placeholder_references/base_service.rb
@@ -19,7 +19,7 @@ def cache
       end
 
       def cache_key
-        @cache_key ||= [:'placeholder-reference', import_source, import_uid].join(':')
+        @cache_key ||= PlaceholderReferences.cache_key(import_source, import_uid)
       end
 
       def logger
diff --git a/app/workers/concerns/import/waits_for_placeholder_references.rb b/app/workers/concerns/import/waits_for_placeholder_references.rb
new file mode 100644
index 00000000000000..abcb665bbdf9df
--- /dev/null
+++ b/app/workers/concerns/import/waits_for_placeholder_references.rb
@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module Import
+  module WaitsForPlaceholderReferences
+    extend ActiveSupport::Concern
+    include Gitlab::ClassAttributes
+
+    def pending_references?
+      pending_references_count > 0
+    end
+
+    private
+
+    def pending_references_count
+      Gitlab::Cache::Import::Caching.set_count(pending_references_cache_key)
+    end
+
+    def pending_references_cache_key
+      uid = instance_exec(&self.class.get_import_uid)
+
+      Import::PlaceholderReferences.cache_key(self.class.get_import_source, uid)
+    end
+
+    class_methods do
+      def import_source(source)
+        set_class_attribute(:import_source, source)
+      end
+
+      def import_uid(uid)
+        set_class_attribute(:import_uid, uid)
+      end
+
+      def get_import_source
+        get_class_attribute(:import_source)
+      end
+
+      def get_import_uid
+        get_class_attribute(:import_uid)
+      end
+    end
+  end
+end
diff --git a/lib/gitlab/cache/import/caching.rb b/lib/gitlab/cache/import/caching.rb
index 7692af6875f19c..c091d9909428e8 100644
--- a/lib/gitlab/cache/import/caching.rb
+++ b/lib/gitlab/cache/import/caching.rb
@@ -143,6 +143,17 @@ def self.set_includes?(raw_key, value)
           end
         end
 
+        # Returns the number of values in the set.
+        #
+        # raw_key - The key of the set to check.
+        def self.set_count(raw_key)
+          key = cache_key_for(raw_key)
+
+          with_redis do |redis|
+            redis.scard(key)
+          end
+        end
+
         # Returns the values of the given set.
         #
         # raw_key - The key of the set to check.
diff --git a/spec/services/import/placeholder_references/load_service_spec.rb b/spec/services/import/placeholder_references/load_service_spec.rb
index b12646aa4aebf7..fd2df058daf0c7 100644
--- a/spec/services/import/placeholder_references/load_service_spec.rb
+++ b/spec/services/import/placeholder_references/load_service_spec.rb
@@ -265,9 +265,7 @@ def expect_log_message(type, **params)
     end
 
     def set_members_count
-      Gitlab::Cache::Import::Caching.with_redis do |redis|
-        redis.scard(Gitlab::Cache::Import::Caching.cache_key_for(cache_key))
-      end
+      Gitlab::Cache::Import::Caching.set_count(cache_key)
     end
 
     def cache_key
diff --git a/spec/services/import/placeholder_references_spec.rb b/spec/services/import/placeholder_references_spec.rb
new file mode 100644
index 00000000000000..78e10e81e20d8d
--- /dev/null
+++ b/spec/services/import/placeholder_references_spec.rb
@@ -0,0 +1,13 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Import::PlaceholderReferences, feature_category: :importers do
+  describe '.cache_key' do
+    it 'builds the cache key correctly' do
+      cache_key = described_class.cache_key(:test_source, 123)
+
+      expect(cache_key).to eq('placeholder-references:test_source:123')
+    end
+  end
+end
diff --git a/spec/workers/concerns/import/waits_for_placeholder_references_spec.rb b/spec/workers/concerns/import/waits_for_placeholder_references_spec.rb
new file mode 100644
index 00000000000000..9e138a28fc25d9
--- /dev/null
+++ b/spec/workers/concerns/import/waits_for_placeholder_references_spec.rb
@@ -0,0 +1,59 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Import::WaitsForPlaceholderReferences, feature_category: :importers do
+  let_it_be(:worker_class) do
+    Class.new do
+      def self.name
+        'MyWorker'
+      end
+
+      include ApplicationWorker
+      include Import::WaitsForPlaceholderReferences
+
+      import_source :test_source
+      import_uid -> { @uid }
+
+      def perform
+        @uid = 123
+
+        pending_references?
+      end
+    end
+  end
+
+  it 'sets the import source' do
+    expect(worker_class.get_import_source).to eq(:test_source)
+  end
+
+  it 'sets the import uid' do
+    expect(worker_class.get_import_uid).to be_kind_of(Proc)
+  end
+
+  describe '#pending_references?', :clean_gitlab_redis_shared_state do
+    subject(:pending_references) { worker_class.new.perform }
+
+    it 'returns false if there are no pending references' do
+      expect(pending_references).to eq(false)
+    end
+
+    it 'returns true if there is a pending reference' do
+      add_reference(:test_source, 123)
+
+      expect(pending_references).to eq(true)
+    end
+
+    it 'returns false if there is a pending reference for another import' do
+      add_reference(:another_import_source, 123)
+      add_reference(:test_source, 234)
+
+      expect(pending_references).to eq(false)
+    end
+
+    def add_reference(import_source, import_uid)
+      cache_key = Import::PlaceholderReferences.cache_key(import_source, import_uid)
+      Gitlab::Cache::Import::Caching.set_add(cache_key, :a_reference)
+    end
+  end
+end
-- 
GitLab


From f3c8d067610a13714e9e637a43f59326ead57c5d Mon Sep 17 00:00:00 2001
From: Luke Duncalfe <lduncalfe@eml.cc>
Date: Tue, 13 Aug 2024 13:05:10 +1200
Subject: [PATCH 2/2] Refactor using a Store

https://gitlab.com/gitlab-org/gitlab/-/merge_requests/158536#note_2023171699
---
 .../source_user_placeholder_reference.rb      |  6 +-
 app/services/import/placeholder_references.rb | 11 ---
 .../placeholder_references/base_service.rb    | 10 +--
 .../placeholder_references/load_service.rb    |  4 +-
 .../placeholder_references/push_service.rb    |  6 +-
 .../waits_for_placeholder_references.rb       | 42 ----------
 .../placeholder_reference_alias_resolver.rb   | 36 ---------
 .../placeholder_references/alias_resolver.rb  | 38 +++++++++
 lib/import/placeholder_references/store.rb    | 50 ++++++++++++
 spec/lib/gitlab/cache/import/caching_spec.rb  | 13 ++++
 .../alias_resolver_spec.rb}                   |  2 +-
 .../placeholder_references/store_spec.rb      | 78 +++++++++++++++++++
 .../source_user_placeholder_reference_spec.rb |  8 +-
 .../load_service_spec.rb                      | 13 ++--
 .../push_service_spec.rb                      |  8 +-
 .../import/placeholder_references_spec.rb     | 13 ----
 .../waits_for_placeholder_references_spec.rb  | 59 --------------
 17 files changed, 202 insertions(+), 195 deletions(-)
 delete mode 100644 app/services/import/placeholder_references.rb
 delete mode 100644 app/workers/concerns/import/waits_for_placeholder_references.rb
 delete mode 100644 lib/import/placeholder_reference_alias_resolver.rb
 create mode 100644 lib/import/placeholder_references/alias_resolver.rb
 create mode 100644 lib/import/placeholder_references/store.rb
 rename spec/lib/import/{placeholder_reference_alias_resolver_spec.rb => placeholder_references/alias_resolver_spec.rb} (97%)
 create mode 100644 spec/lib/import/placeholder_references/store_spec.rb
 delete mode 100644 spec/services/import/placeholder_references_spec.rb
 delete mode 100644 spec/workers/concerns/import/waits_for_placeholder_references_spec.rb

diff --git a/app/models/import/source_user_placeholder_reference.rb b/app/models/import/source_user_placeholder_reference.rb
index ab27dc7ff5e55b..9338d47e680e8c 100644
--- a/app/models/import/source_user_placeholder_reference.rb
+++ b/app/models/import/source_user_placeholder_reference.rb
@@ -42,16 +42,16 @@ class SourceUserPlaceholderReference < ApplicationRecord
     SerializationError = Class.new(StandardError)
 
     def aliased_model
-      Import::PlaceholderReferenceAliasResolver.aliased_model(model)
+      PlaceholderReferences::AliasResolver.aliased_model(model)
     end
 
     def aliased_user_reference_column
-      Import::PlaceholderReferenceAliasResolver.aliased_column(model, user_reference_column)
+      PlaceholderReferences::AliasResolver.aliased_column(model, user_reference_column)
     end
 
     def aliased_composite_key
       composite_key.transform_keys do |key|
-        Import::PlaceholderReferenceAliasResolver.aliased_column(model, key)
+        PlaceholderReferences::AliasResolver.aliased_column(model, key)
       end
     end
 
diff --git a/app/services/import/placeholder_references.rb b/app/services/import/placeholder_references.rb
deleted file mode 100644
index 4328ff045bf3de..00000000000000
--- a/app/services/import/placeholder_references.rb
+++ /dev/null
@@ -1,11 +0,0 @@
-# frozen_string_literal: true
-
-module Import
-  module PlaceholderReferences
-    class << self
-      def cache_key(import_source, import_uid)
-        [:'placeholder-references', import_source, import_uid].join(':')
-      end
-    end
-  end
-end
diff --git a/app/services/import/placeholder_references/base_service.rb b/app/services/import/placeholder_references/base_service.rb
index 54e0f43212dbc9..681298f1d99f89 100644
--- a/app/services/import/placeholder_references/base_service.rb
+++ b/app/services/import/placeholder_references/base_service.rb
@@ -12,14 +12,10 @@ def initialize(import_source:, import_uid:)
 
       private
 
-      attr_reader :import_source, :import_uid, :reference
+      attr_reader :import_source, :import_uid
 
-      def cache
-        Gitlab::Cache::Import::Caching
-      end
-
-      def cache_key
-        @cache_key ||= PlaceholderReferences.cache_key(import_source, import_uid)
+      def store
+        @store ||= PlaceholderReferences::Store.new(import_source: import_source, import_uid: import_uid)
       end
 
       def logger
diff --git a/app/services/import/placeholder_references/load_service.rb b/app/services/import/placeholder_references/load_service.rb
index c18a155a62216e..6a382c033e10cc 100644
--- a/app/services/import/placeholder_references/load_service.rb
+++ b/app/services/import/placeholder_references/load_service.rb
@@ -37,7 +37,7 @@ def execute
       attr_accessor :error_count, :processed_count
 
       def next_batch
-        cache.limited_values_from_set(cache_key, limit: BATCH_LIMIT)
+        store.get(BATCH_LIMIT)
       end
 
       def load!(batch)
@@ -79,7 +79,7 @@ def clear_batch!(batch)
 
         self.processed_count += processed_count
 
-        cache.set_remove(cache_key, batch)
+        store.remove(batch)
       end
 
       def log_error(item, exception)
diff --git a/app/services/import/placeholder_references/push_service.rb b/app/services/import/placeholder_references/push_service.rb
index 468d74c52a996b..3d0de3e2c56856 100644
--- a/app/services/import/placeholder_references/push_service.rb
+++ b/app/services/import/placeholder_references/push_service.rb
@@ -48,7 +48,7 @@ def execute
 
         serialized_reference = reference.to_serialized
 
-        cache.set_add(cache_key, serialized_reference, timeout: cache_ttl)
+        store.add(serialized_reference)
 
         success(serialized_reference: serialized_reference)
       end
@@ -57,10 +57,6 @@ def execute
 
       attr_reader :reference
 
-      def cache_ttl
-        Gitlab::Cache::Import::Caching::TIMEOUT
-      end
-
       def track_error(reference)
         Gitlab::ErrorTracking.track_and_raise_for_dev_exception(
           InvalidReferenceError.new('Invalid placeholder user reference'),
diff --git a/app/workers/concerns/import/waits_for_placeholder_references.rb b/app/workers/concerns/import/waits_for_placeholder_references.rb
deleted file mode 100644
index abcb665bbdf9df..00000000000000
--- a/app/workers/concerns/import/waits_for_placeholder_references.rb
+++ /dev/null
@@ -1,42 +0,0 @@
-# frozen_string_literal: true
-
-module Import
-  module WaitsForPlaceholderReferences
-    extend ActiveSupport::Concern
-    include Gitlab::ClassAttributes
-
-    def pending_references?
-      pending_references_count > 0
-    end
-
-    private
-
-    def pending_references_count
-      Gitlab::Cache::Import::Caching.set_count(pending_references_cache_key)
-    end
-
-    def pending_references_cache_key
-      uid = instance_exec(&self.class.get_import_uid)
-
-      Import::PlaceholderReferences.cache_key(self.class.get_import_source, uid)
-    end
-
-    class_methods do
-      def import_source(source)
-        set_class_attribute(:import_source, source)
-      end
-
-      def import_uid(uid)
-        set_class_attribute(:import_uid, uid)
-      end
-
-      def get_import_source
-        get_class_attribute(:import_source)
-      end
-
-      def get_import_uid
-        get_class_attribute(:import_uid)
-      end
-    end
-  end
-end
diff --git a/lib/import/placeholder_reference_alias_resolver.rb b/lib/import/placeholder_reference_alias_resolver.rb
deleted file mode 100644
index fb5715826c7a17..00000000000000
--- a/lib/import/placeholder_reference_alias_resolver.rb
+++ /dev/null
@@ -1,36 +0,0 @@
-# frozen_string_literal: true
-
-module Import
-  module PlaceholderReferenceAliasResolver
-    MissingAlias = Class.new(StandardError)
-
-    ALIASES = {
-      "Note" => {
-        model: Note,
-        columns: {
-          "author_id" => "author_id"
-        }
-      }
-    }.freeze
-
-    def self.aliased_model(model)
-      aliased_model = ALIASES.dig(model, :model)
-      return aliased_model if aliased_model.present?
-
-      message = "ALIASES must be extended to include #{model}"
-      Gitlab::ErrorTracking.track_and_raise_for_dev_exception(MissingAlias.new(message))
-
-      model.safe_constantize
-    end
-
-    def self.aliased_column(model, column)
-      aliased_column = ALIASES.dig(model, :columns, column)
-      return aliased_column if aliased_column.present?
-
-      message = "ALIASES must be extended to include #{model} #{column}"
-      Gitlab::ErrorTracking.track_and_raise_for_dev_exception(MissingAlias.new(message))
-
-      column
-    end
-  end
-end
diff --git a/lib/import/placeholder_references/alias_resolver.rb b/lib/import/placeholder_references/alias_resolver.rb
new file mode 100644
index 00000000000000..18a8539fca248e
--- /dev/null
+++ b/lib/import/placeholder_references/alias_resolver.rb
@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Import
+  module PlaceholderReferences
+    module AliasResolver
+      MissingAlias = Class.new(StandardError)
+
+      ALIASES = {
+        "Note" => {
+          model: Note,
+          columns: {
+            "author_id" => "author_id"
+          }
+        }
+      }.freeze
+
+      def self.aliased_model(model)
+        aliased_model = ALIASES.dig(model, :model)
+        return aliased_model if aliased_model.present?
+
+        message = "ALIASES must be extended to include #{model}"
+        Gitlab::ErrorTracking.track_and_raise_for_dev_exception(MissingAlias.new(message))
+
+        model.safe_constantize
+      end
+
+      def self.aliased_column(model, column)
+        aliased_column = ALIASES.dig(model, :columns, column)
+        return aliased_column if aliased_column.present?
+
+        message = "ALIASES must be extended to include #{model} #{column}"
+        Gitlab::ErrorTracking.track_and_raise_for_dev_exception(MissingAlias.new(message))
+
+        column
+      end
+    end
+  end
+end
diff --git a/lib/import/placeholder_references/store.rb b/lib/import/placeholder_references/store.rb
new file mode 100644
index 00000000000000..b3d07c736b44e8
--- /dev/null
+++ b/lib/import/placeholder_references/store.rb
@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module Import
+  module PlaceholderReferences
+    class Store
+      CACHE_TTL = 1.day
+
+      def initialize(import_source:, import_uid:)
+        @import_source = import_source
+        @import_uid = import_uid
+      end
+
+      def add(serialized_reference)
+        cache.set_add(cache_key, serialized_reference, timeout: CACHE_TTL)
+      end
+
+      def get(limit = 1)
+        cache.limited_values_from_set(cache_key, limit: limit)
+      end
+
+      def remove(values)
+        cache.set_remove(cache_key, values)
+      end
+
+      def count
+        cache.set_count(cache_key)
+      end
+
+      def empty?
+        count == 0
+      end
+
+      def any?
+        !empty?
+      end
+
+      private
+
+      attr_reader :import_source, :import_uid
+
+      def cache
+        Gitlab::Cache::Import::Caching
+      end
+
+      def cache_key
+        [:'placeholder-references', import_source, import_uid].join(':')
+      end
+    end
+  end
+end
diff --git a/spec/lib/gitlab/cache/import/caching_spec.rb b/spec/lib/gitlab/cache/import/caching_spec.rb
index c84721eb7cb70a..6c813450a645ca 100644
--- a/spec/lib/gitlab/cache/import/caching_spec.rb
+++ b/spec/lib/gitlab/cache/import/caching_spec.rb
@@ -193,6 +193,19 @@
     end
   end
 
+  describe '.set_count' do
+    it 'returns 0 when the set does not exist' do
+      expect(described_class.set_count('foo')).to eq(0)
+    end
+
+    it 'returns count of set' do
+      described_class.set_add('foo', 10)
+      described_class.set_add('foo', 20)
+
+      expect(described_class.set_count('foo')).to eq(2)
+    end
+  end
+
   describe '.hash_add' do
     it 'adds a value to a hash' do
       described_class.hash_add('foo', 1, 1)
diff --git a/spec/lib/import/placeholder_reference_alias_resolver_spec.rb b/spec/lib/import/placeholder_references/alias_resolver_spec.rb
similarity index 97%
rename from spec/lib/import/placeholder_reference_alias_resolver_spec.rb
rename to spec/lib/import/placeholder_references/alias_resolver_spec.rb
index 1126cd657c315d..6ed3383654b1e5 100644
--- a/spec/lib/import/placeholder_reference_alias_resolver_spec.rb
+++ b/spec/lib/import/placeholder_references/alias_resolver_spec.rb
@@ -2,7 +2,7 @@
 
 require "spec_helper"
 
-RSpec.describe Import::PlaceholderReferenceAliasResolver, feature_category: :importers do
+RSpec.describe Import::PlaceholderReferences::AliasResolver, feature_category: :importers do
   describe "ALIASES" do
     it "points to real columns" do
       def failure_message(model, column)
diff --git a/spec/lib/import/placeholder_references/store_spec.rb b/spec/lib/import/placeholder_references/store_spec.rb
new file mode 100644
index 00000000000000..5fed3c2dbcae6e
--- /dev/null
+++ b/spec/lib/import/placeholder_references/store_spec.rb
@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+RSpec.describe Import::PlaceholderReferences::Store, :clean_gitlab_redis_shared_state, feature_category: :importers do
+  subject(:store) { described_class.new(import_source: 'source', import_uid: 'uid') }
+
+  let(:cache_key) { [:'placeholder-references', 'source', 'uid'].join(':') }
+
+  describe '#add' do
+    it 'adds to the set' do
+      store.add('foo')
+
+      expect(cache.values_from_set(cache_key)).to eq(['foo'])
+    end
+  end
+
+  describe '#get' do
+    before do
+      cache.set_add(cache_key, 'foo')
+      cache.set_add(cache_key, 'bar')
+      cache.set_add(cache_key, 'baz')
+    end
+
+    it 'returns a member in the set' do
+      expect(store.get.size).to eq(1)
+      expect(store.get.first).to be_in(%w[foo bar baz])
+    end
+
+    it 'accepts an argument to return more members' do
+      expect(store.get(2).size).to eq(2)
+    end
+  end
+
+  describe '#remove' do
+    it 'removes members from the set' do
+      cache.set_add(cache_key, 'foo')
+      cache.set_add(cache_key, 'bar')
+      cache.set_add(cache_key, 'baz')
+
+      store.remove(%w[foo baz])
+
+      expect(cache.values_from_set(cache_key)).to eq(['bar'])
+    end
+  end
+
+  describe '#count' do
+    it 'returns the count of members in the set' do
+      cache.set_add(cache_key, 'foo')
+      cache.set_add(cache_key, 'bar')
+
+      expect(store.count).to eq(2)
+    end
+  end
+
+  describe '#empty?' do
+    it 'returns true if the set is empty' do
+      expect(store.empty?).to eq(true)
+    end
+
+    it 'returns false if the set is not empty' do
+      cache.set_add(cache_key, 'foo')
+
+      expect(store.empty?).to eq(false)
+    end
+  end
+
+  describe '#any?' do
+    it 'returns the inverse of #empty?' do
+      expect(store).to receive(:empty?).and_return(true)
+      expect(store.any?).to eq(false)
+    end
+  end
+
+  def cache
+    Gitlab::Cache::Import::Caching
+  end
+end
diff --git a/spec/models/import/source_user_placeholder_reference_spec.rb b/spec/models/import/source_user_placeholder_reference_spec.rb
index f4f000eef479ac..a23e4d1a75b2e0 100644
--- a/spec/models/import/source_user_placeholder_reference_spec.rb
+++ b/spec/models/import/source_user_placeholder_reference_spec.rb
@@ -136,7 +136,7 @@ def validation_errors(...)
       end
 
       before do
-        allow(Import::PlaceholderReferenceAliasResolver).to receive(:aliased_model).and_return(Note)
+        allow(Import::PlaceholderReferences::AliasResolver).to receive(:aliased_model).and_return(Note)
       end
 
       it "uses the new model" do
@@ -158,7 +158,7 @@ def validation_errors(...)
 
     context "when the column name has changed" do
       before do
-        allow(Import::PlaceholderReferenceAliasResolver).to receive(:aliased_column).and_return("user_id")
+        allow(Import::PlaceholderReferences::AliasResolver).to receive(:aliased_column).and_return("user_id")
       end
 
       it "uses the new column" do
@@ -177,8 +177,8 @@ def validation_errors(...)
     end
 
     before do
-      allow(Import::PlaceholderReferenceAliasResolver).to receive(:aliased_column).and_call_original
-      allow(Import::PlaceholderReferenceAliasResolver).to receive(:aliased_column)
+      allow(Import::PlaceholderReferences::AliasResolver).to receive(:aliased_column).and_call_original
+      allow(Import::PlaceholderReferences::AliasResolver).to receive(:aliased_column)
         .with("Note", "old_id").and_return("new_id")
     end
 
diff --git a/spec/services/import/placeholder_references/load_service_spec.rb b/spec/services/import/placeholder_references/load_service_spec.rb
index fd2df058daf0c7..bfe57554711c5d 100644
--- a/spec/services/import/placeholder_references/load_service_spec.rb
+++ b/spec/services/import/placeholder_references/load_service_spec.rb
@@ -120,7 +120,7 @@
         push(valid_reference_2)
         # We can't use `#push` because that validates,
         # so push directly to the set.
-        Gitlab::Cache::Import::Caching.set_add(cache_key, invalid_reference.values.to_json)
+        store.add(invalid_reference.values.to_json)
         push(valid_reference_3)
       end
 
@@ -255,7 +255,8 @@
 
     def push(reference)
       serialized = Import::SourceUserPlaceholderReference.new(reference).to_serialized
-      Gitlab::Cache::Import::Caching.set_add(cache_key, serialized)
+
+      store.add(serialized)
     end
 
     def expect_log_message(type, **params)
@@ -265,14 +266,14 @@ def expect_log_message(type, **params)
     end
 
     def set_members_count
-      Gitlab::Cache::Import::Caching.set_count(cache_key)
+      store.count
     end
 
-    def cache_key
-      Import::PlaceholderReferences::BaseService.new(
+    def store
+      Import::PlaceholderReferences::Store.new(
         import_source: import_source,
         import_uid: import_uid
-      ).send(:cache_key)
+      )
     end
   end
 end
diff --git a/spec/services/import/placeholder_references/push_service_spec.rb b/spec/services/import/placeholder_references/push_service_spec.rb
index b9ddb2acd56a1c..89e3d66a9745b6 100644
--- a/spec/services/import/placeholder_references/push_service_spec.rb
+++ b/spec/services/import/placeholder_references/push_service_spec.rb
@@ -151,13 +151,9 @@
   end
 
   def set
-    Gitlab::Cache::Import::Caching.values_from_set(cache_key)
-  end
-
-  def cache_key
-    Import::PlaceholderReferences::BaseService.new(
+    Import::PlaceholderReferences::Store.new(
       import_source: import_source,
       import_uid: import_uid
-    ).send(:cache_key)
+    ).get
   end
 end
diff --git a/spec/services/import/placeholder_references_spec.rb b/spec/services/import/placeholder_references_spec.rb
deleted file mode 100644
index 78e10e81e20d8d..00000000000000
--- a/spec/services/import/placeholder_references_spec.rb
+++ /dev/null
@@ -1,13 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Import::PlaceholderReferences, feature_category: :importers do
-  describe '.cache_key' do
-    it 'builds the cache key correctly' do
-      cache_key = described_class.cache_key(:test_source, 123)
-
-      expect(cache_key).to eq('placeholder-references:test_source:123')
-    end
-  end
-end
diff --git a/spec/workers/concerns/import/waits_for_placeholder_references_spec.rb b/spec/workers/concerns/import/waits_for_placeholder_references_spec.rb
deleted file mode 100644
index 9e138a28fc25d9..00000000000000
--- a/spec/workers/concerns/import/waits_for_placeholder_references_spec.rb
+++ /dev/null
@@ -1,59 +0,0 @@
-# frozen_string_literal: true
-
-require 'spec_helper'
-
-RSpec.describe Import::WaitsForPlaceholderReferences, feature_category: :importers do
-  let_it_be(:worker_class) do
-    Class.new do
-      def self.name
-        'MyWorker'
-      end
-
-      include ApplicationWorker
-      include Import::WaitsForPlaceholderReferences
-
-      import_source :test_source
-      import_uid -> { @uid }
-
-      def perform
-        @uid = 123
-
-        pending_references?
-      end
-    end
-  end
-
-  it 'sets the import source' do
-    expect(worker_class.get_import_source).to eq(:test_source)
-  end
-
-  it 'sets the import uid' do
-    expect(worker_class.get_import_uid).to be_kind_of(Proc)
-  end
-
-  describe '#pending_references?', :clean_gitlab_redis_shared_state do
-    subject(:pending_references) { worker_class.new.perform }
-
-    it 'returns false if there are no pending references' do
-      expect(pending_references).to eq(false)
-    end
-
-    it 'returns true if there is a pending reference' do
-      add_reference(:test_source, 123)
-
-      expect(pending_references).to eq(true)
-    end
-
-    it 'returns false if there is a pending reference for another import' do
-      add_reference(:another_import_source, 123)
-      add_reference(:test_source, 234)
-
-      expect(pending_references).to eq(false)
-    end
-
-    def add_reference(import_source, import_uid)
-      cache_key = Import::PlaceholderReferences.cache_key(import_source, import_uid)
-      Gitlab::Cache::Import::Caching.set_add(cache_key, :a_reference)
-    end
-  end
-end
-- 
GitLab