Allow UserRefreshFromReplicaWorker to read from replica

This change allows
UserRefreshFromReplicaWorker to read
from replica

app/workers/all_queues.yml

@@ -37,7 +37,7 @@
   :urgency: :low
   :resource_boundary: :unknown
   :weight: 1
-  :idempotent: true
+  :idempotent:
   :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 < ::AuthorizedProjectsWorker
+  class UserRefreshFromReplicaWorker # rubocop:disable Scalability/IdempotentWorker
+    include ApplicationWorker
+
+    sidekiq_options retry: 3
     feature_category :authentication_and_authorization
     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
     deduplicate :until_executing, including_scheduled: true
 
-    idempotent!
+    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)
+      else
+        use_primary_database
+        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
+    end
+
+    def project_authorizations_needs_refresh?(user)
+      AuthorizedProjectUpdate::FindRecordsDueForRefreshService.new(user).needs_refresh?
+    end
+
+    def enqueue_project_authorizations_refresh(user)
+      with_context(user: user, related_class: current_caller_id) do
+        AuthorizedProjectUpdate::UserRefreshWithLowUrgencyWorker.perform_async(user.id)
+      end
+    end
 
-    # This worker will start reading data from the replica database soon
-    # Issue: https://gitlab.com/gitlab-org/gitlab/-/issues/333219
+    # We use this so that we can obtain the details of the original caller
+    # in the enqueued `AuthorizedProjectUpdate::UserRefreshWithLowUrgencyWorker` job.
+    def current_caller_id
+      Gitlab::ApplicationContext.current_context_attribute('meta.caller_id').presence
+    end
   end
 end

config/feature_flags/development/user_refresh_from_replica_worker_uses_replica_db.yml

+---
+name: user_refresh_from_replica_worker_uses_replica_db
+introduced_by_url: https://gitlab.com/gitlab-org/gitlab/-/merge_requests/64276
+rollout_issue_url: https://gitlab.com/gitlab-org/gitlab/-/issues/334766
+milestone: '14.1'
+type: development
+group: group::access
+default_enabled: false

spec/workers/authorized_project_update/user_refresh_from_replica_worker_spec.rb

@@ -3,9 +3,97 @@
 require 'spec_helper'
 
 RSpec.describe AuthorizedProjectUpdate::UserRefreshFromReplicaWorker do
+  let_it_be(:project) { create(:project) }
+  let_it_be(:user) { project.namespace.owner }
+
+  let(:execute_worker) { subject.perform(user.id) }
+
   it 'is labeled as low urgency' do
     expect(described_class.get_urgency).to eq(:low)
   end
 
-  it_behaves_like "refreshes user's project authorizations"
+  it_behaves_like 'worker with data consistency',
+                  described_class,
+                  data_consistency: :delayed
+
+  describe '#perform' do
+    it 'checks if a project_authorization refresh is needed for the user' do
+      expect(AuthorizedProjectUpdate::FindRecordsDueForRefreshService).to(
+        receive(:new).with(user).and_call_original)
+
+      execute_worker
+    end
+
+    context 'when there are project authorization records due for either removal or addition for a specific user' do
+      before do
+        user.project_authorizations.delete_all
+      end
+
+      it 'enqueues a new project authorization update job for the user' do
+        expect(AuthorizedProjectUpdate::UserRefreshWithLowUrgencyWorker).to receive(:perform_async).with(user.id)
+
+        execute_worker
+      end
+
+      context 'setting `meta.caller_id` as `meta.related_class` in the context of the newly enqueued `UserRefreshWithLowUrgencyWorker` job' do
+        context 'when the `UserRefreshFromReplicaWorker` job has a `caller_id` set' do
+          it 'sets the same `caller_id` as `related_class`' do
+            expect(AuthorizedProjectUpdate::UserRefreshWithLowUrgencyWorker).to receive(:perform_async).with(user.id) do
+              expect(Gitlab::ApplicationContext.current).to include('meta.related_class' => 'Foo')
+            end
+
+            Gitlab::ApplicationContext.with_context(caller_id: 'Foo') do
+              execute_worker
+            end
+          end
+        end
+
+        context 'when the `UserRefreshFromReplicaWorker` job does not have a `caller_id` set' do
+          it 'does not set the value of `related_class`' do
+            expect(AuthorizedProjectUpdate::UserRefreshWithLowUrgencyWorker).to receive(:perform_async).with(user.id) do
+              expect(Gitlab::ApplicationContext.current).not_to include('meta.related_class')
+            end
+
+            execute_worker
+          end
+        end
+      end
+    end
+
+    context 'when there are no additions or removals to be made to project authorizations for a specific user' do
+      it 'does not enqueue a new project authorization update job for the user' do
+        expect(AuthorizedProjectUpdate::UserRefreshWithLowUrgencyWorker).not_to receive(:perform_async)
+
+        execute_worker
+      end
+    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
+
+      context 'when load balancing is enabled' do
+        before do
+          allow(Gitlab::Database::LoadBalancing).to receive(:enable?).and_return(true)
+        end
+
+        it 'reads from the primary database' do
+          expect(Gitlab::Database::LoadBalancing::Session.current)
+            .to receive(:use_primary!)
+
+          execute_worker
+        end
+      end
+
+      it 'calls Users::RefreshAuthorizedProjectsService' do
+        source = 'AuthorizedProjectUpdate::UserRefreshFromReplicaWorker'
+        expect_next_instance_of(Users::RefreshAuthorizedProjectsService, user, { source: source }) do |service|
+          expect(service).to receive(:execute)
+        end
+
+        execute_worker
+      end
+    end
+  end
 end