Resolve "Allow deduplication for UserRefreshFromReplicaWorker, while still...

Resolve "Allow deduplication for UserRefreshFromReplicaWorker, while still reading from the replica database."

app/workers/all_queues.yml

@@ -37,7 +37,7 @@
   :urgency: :low
   :resource_boundary: :unknown
   :weight: 1
-  :idempotent:
+  :idempotent: true
   :tags: []
 - :name: authorized_project_update:authorized_project_update_user_refresh_over_user_range
   :worker_name: AuthorizedProjectUpdate::UserRefreshOverUserRangeWorker

app/workers/authorized_project_update/user_refresh_from_replica_worker.rb

 # frozen_string_literal: true
 
 module AuthorizedProjectUpdate
-  class UserRefreshFromReplicaWorker # rubocop:disable Scalability/IdempotentWorker
+  class UserRefreshFromReplicaWorker
     include ApplicationWorker
 
     sidekiq_options retry: 3
@@ -9,31 +9,36 @@ module AuthorizedProjectUpdate
     urgency :low
     queue_namespace :authorized_project_update
 
-    # This job will not be deduplicated since it is marked with
-    # `data_consistency :delayed` and not `idempotent!`
-    # See https://gitlab.com/gitlab-org/gitlab/-/issues/325291
+    idempotent!
     deduplicate :until_executing, including_scheduled: true
 
-    data_consistency :delayed
-
     def perform(user_id)
-      user = User.find_by_id(user_id)
-      return unless user
-
       if Feature.enabled?(:user_refresh_from_replica_worker_uses_replica_db)
-        enqueue_project_authorizations_refresh(user) if project_authorizations_needs_refresh?(user)
+        use_replica_if_available do
+          user = User.find_by_id(user_id)
+
+          if user && project_authorizations_needs_refresh?(user)
+            enqueue_project_authorizations_refresh(user)
+          end
+        end
       else
-        use_primary_database
+        user = User.find_by_id(user_id)
+        return unless user
+
         user.refresh_authorized_projects(source: self.class.name)
       end
     end
 
     private
 
-    def use_primary_database
-      if ::Gitlab::Database::LoadBalancing.enable?
-        ::Gitlab::Database::LoadBalancing::Session.current.use_primary!
-      end
+    # We use this approach instead of specifying `data_consistency :delayed` because these jobs
+    # are enqueued in large numbers, and using `data_consistency :delayed`
+    # does not allow us to deduplicate these jobs.
+    # https://gitlab.com/gitlab-org/gitlab/-/issues/325291
+    def use_replica_if_available(&block)
+      return yield unless ::Gitlab::Database::LoadBalancing.enable?
+
+      ::Gitlab::Database::LoadBalancing::Session.current.use_replicas_for_read_queries(&block)
     end
 
     def project_authorizations_needs_refresh?(user)

spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb

@@ -12,9 +12,9 @@ RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
     expect(described_class.get_urgency).to eq(:low)
   end
 
-  it_behaves_like 'worker with data consistency',
-                  described_class,
-                  data_consistency: :delayed
+  it_behaves_like 'an idempotent worker' do
+    let(:job_args) { user.id }
+  end
 
   describe '#perform' do
     it 'checks if a project_authorization refresh is needed for the user' do
@@ -44,22 +44,21 @@ RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
       end
     end
 
-    context 'when the feature flag `user_refresh_from_replica_worker_uses_replica_db` is disabled' do
+    context 'with load balancing enabled' do
       before do
-        stub_feature_flags(user_refresh_from_replica_worker_uses_replica_db: false)
+        allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
       end
 
-      context 'when load balancing is enabled' do
-        before do
-          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
-        end
+      it 'reads from the replica database' do
+        expect(Gitlab::Database::LoadBalancing::Session.current).to receive(:use_replicas_for_read_queries).and_call_original
 
-        it 'reads from the primary database' do
-          expect(Gitlab::Database::LoadBalancing::Session.current)
-            .to receive(:use_primary!)
+        execute_worker
+      end
+    end
 
-          execute_worker
-        end
+    context 'when the feature flag `user_refresh_from_replica_worker_uses_replica_db` is disabled' do
+      before do
+        stub_feature_flags(user_refresh_from_replica_worker_uses_replica_db: false)
       end
 
       it 'calls Users::RefreshAuthorizedProjectsService' do