From b12aecd8400230ed09c2bca44771db44a2bc6129 Mon Sep 17 00:00:00 2001
From: Lin Jen-Shin <jen-shin@gitlab.com>
Date: Sat, 4 Oct 2025 00:22:09 +0800
Subject: [PATCH] Ensure feature flag cells_unique_claims is consistent across
 Current

---
 .rubocop_todo/gitlab/feature_flag_without_actor.yml | 1 -
 app/models/concerns/cells/transaction_callback.rb   | 8 ++++++++
 app/models/concerns/cells/unique.rb                 | 4 ++--
 app/models/current.rb                               | 7 +++++++
 spec/models/concerns/cells/unique_spec.rb           | 1 +
 5 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/.rubocop_todo/gitlab/feature_flag_without_actor.yml b/.rubocop_todo/gitlab/feature_flag_without_actor.yml
index 57245a6e1fbd53..4bc8e1c724ba78 100644
--- a/.rubocop_todo/gitlab/feature_flag_without_actor.yml
+++ b/.rubocop_todo/gitlab/feature_flag_without_actor.yml
@@ -21,7 +21,6 @@ Gitlab/FeatureFlagWithoutActor:
     - 'app/models/clusters/instance.rb'
     - 'app/models/concerns/counter_attribute.rb'
     - 'app/models/concerns/reset_on_column_errors.rb'
-    - 'app/models/concerns/cells/unique.rb'
     - 'app/models/concerns/web_hooks/auto_disabling.rb'
     - 'app/models/merge_request.rb'
     - 'app/models/namespace.rb'
diff --git a/app/models/concerns/cells/transaction_callback.rb b/app/models/concerns/cells/transaction_callback.rb
index 69047adbeac274..736f01cdf2326b 100644
--- a/app/models/concerns/cells/transaction_callback.rb
+++ b/app/models/concerns/cells/transaction_callback.rb
@@ -14,16 +14,22 @@ def self.visited_stack
     end
 
     def begin_isolated_db_transaction(...)
+      return super unless Current.feature_cells_unique_claims
+
       TransactionCallback.visited_stack << {}
       super
     end
 
     def begin_db_transaction
+      return super unless Current.feature_cells_unique_claims
+
       TransactionCallback.visited_stack << {}
       super
     end
 
     def commit_db_transaction
+      return super unless Current.feature_cells_unique_claims
+
       # Don't pop the visited stack yet, and only do that after it's
       # successfully committed, because errors can happen while claiming,
       # and when that happens, we'll roll it back, and we need to look at
@@ -36,6 +42,8 @@ def commit_db_transaction
     end
 
     def rollback_db_transaction
+      return super unless Current.feature_cells_unique_claims
+
       Unique.rollback do
         super
 
diff --git a/app/models/concerns/cells/unique.rb b/app/models/concerns/cells/unique.rb
index db8c3c84f39b7d..ceb44e268fea7d 100644
--- a/app/models/concerns/cells/unique.rb
+++ b/app/models/concerns/cells/unique.rb
@@ -62,14 +62,14 @@ def self.rollback
 
     def inspect_save_changes
       return unless Gitlab.config.cell.enabled
-      return unless Feature.enabled?(:cells_unique_claims)
+      return unless Current.feature_cells_unique_claims
 
       InspectChanges.new(self).inspect_changes
     end
 
     def inspect_destroy_changes
       return unless Gitlab.config.cell.enabled
-      return unless Feature.enabled?(:cells_unique_claims)
+      return unless Current.feature_cells_unique_claims
 
       InspectChanges.new(self, destroy: true).inspect_changes
     end
diff --git a/app/models/current.rb b/app/models/current.rb
index 6747b08526f7a6..a4f78841b0c8ff 100644
--- a/app/models/current.rb
+++ b/app/models/current.rb
@@ -18,6 +18,9 @@ def message
   attribute :organization, :organization_assigned
   attribute :token_info
 
+  attribute :feature_cells_unique_claims
+  resets :reset_feature_cells_unique_claims
+
   def organization=(value)
     # We want to explicitly allow only one organization assignment per thread
     # This fits the request/response cycle, but of course for rake tasks/background jobs that use the same thread,
@@ -45,6 +48,10 @@ def organization
     super
   end
 
+  def reset_feature_cells_unique_claims
+    self.feature_cells_unique_claims = Feature.enabled?(:cells_unique_claims) # rubocop:disable Gitlab/FeatureFlagWithoutActor -- We can't easily tie this to an actor
+  end
+
   private
 
   # Do not allow to reset this
diff --git a/spec/models/concerns/cells/unique_spec.rb b/spec/models/concerns/cells/unique_spec.rb
index d4121c29301125..392bcd492a0d0c 100644
--- a/spec/models/concerns/cells/unique_spec.rb
+++ b/spec/models/concerns/cells/unique_spec.rb
@@ -88,6 +88,7 @@
       context 'when feature flag is disabled' do
         before do
           stub_feature_flags(cells_unique_claims: false)
+          Current.reset
         end
 
         it 'does not call inspect_save_changes on save' do
-- 
GitLab