From 4f7fa32af92f47dea621b0b53af5cb863665a6c1 Mon Sep 17 00:00:00 2001
From: agius <andrew@atevans.com>
Date: Mon, 9 Dec 2024 20:26:34 -0800
Subject: [PATCH] Per-cell organization for internal bot users

Creates a per-cell organization for internal bot users. Currently, all
internal bot users belong to the default organization, but we will want
them to be present for each cell individually at minimum.
---
 lib/users/internal.rb           | 35 ++++++++++++++++++++++++++++++++-
 spec/lib/users/internal_spec.rb |  4 ++--
 2 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/lib/users/internal.rb b/lib/users/internal.rb
index ff432f20460f60..66daa82e7136c5 100644
--- a/lib/users/internal.rb
+++ b/lib/users/internal.rb
@@ -162,13 +162,46 @@ def create_unique_internal(scope, username, email_pattern, &creation_block)
         )
 
         # https://gitlab.com/gitlab-org/gitlab/-/issues/442780
-        user.assign_personal_namespace(::Organizations::Organization.first)
+        user.assign_personal_namespace(internal_org)
 
         Users::UpdateService.new(user, user: user).execute(validate: false)
         user
       ensure
         Gitlab::ExclusiveLease.cancel(lease_key, uuid)
       end
+
+      def internal_org
+        org_name = "#{Gitlab.config.cell.name} Bot Users"
+        org_path = org_name.underscore.tr(' ', '_')
+        org = Organizations::Organization.find_by_path(org_path)
+
+        return org if org.present?
+
+        lease_key = "organization:unique_internal:#{org_path}"
+        lease = Gitlab::ExclusiveLease.new(lease_key, timeout: 1.minute.to_i)
+
+        uuid = lease.try_obtain
+        until uuid.present?
+          # Keep trying until we obtain the lease. To prevent hammering Redis too
+          # much we'll wait for a bit between retries.
+          sleep(1)
+          uuid = lease.try_obtain
+        end
+
+        # Recheck if the org is already present. One might have been
+        # added between the time we last checked (first line of this method)
+        # and the time we acquired the lock.
+        existing_org = Organizations::Organization.find_by_path(org_path)
+        return existing_org if existing_org.present?
+
+        Organizations::Organization.create!(
+          name: org_name,
+          path: org_path,
+          visibility_level: 'internal'
+        )
+      ensure
+        Gitlab::ExclusiveLease.cancel(lease_key, uuid)
+      end
     end
   end
 end
diff --git a/spec/lib/users/internal_spec.rb b/spec/lib/users/internal_spec.rb
index 39e1a15694e306..948f07af6c6079 100644
--- a/spec/lib/users/internal_spec.rb
+++ b/spec/lib/users/internal_spec.rb
@@ -3,7 +3,7 @@
 require 'spec_helper'
 
 RSpec.describe Users::Internal, feature_category: :user_profile do
-  let_it_be(:default_organization) { create(:organization, :default) }
+  let_it_be(:bot_users_organization) { create(:organization, name: 'Cell 1 Bot Users', path: 'cell_1_bot_users') }
 
   shared_examples 'bot users' do |bot_type, username, email|
     it 'creates the user if it does not exist' do
@@ -16,7 +16,7 @@
       bot_user = described_class.public_send(bot_type)
 
       expect(bot_user.namespace.route).to be_present
-      expect(bot_user.namespace.organization).to eq(default_organization)
+      expect(bot_user.namespace.organization).to eq(bot_users_organization)
     end
 
     it 'does not create a new user if it already exists' do
-- 
GitLab