Remove pick_repository_storage from containers

In this commit remove the pick_repository_storage method
from repository containers and, instead, we created its own
method in the `Repository` class.

app/models/concerns/has_repository.rb

@@ -15,15 +15,6 @@ module HasRepository
 
   delegate :base_dir, :disk_path, to: :storage
 
-  class_methods do
-    def pick_repository_storage
-      # We need to ensure application settings are fresh when we pick
-      # a repository storage to use.
-      Gitlab::CurrentSettings.expire_current_application_settings
-      Gitlab::CurrentSettings.pick_repository_storage
-    end
-  end
-
   def valid_repo?
     repository.exists?
   rescue

app/models/concerns/repository_storage_movable.rb

@@ -20,7 +20,7 @@ module RepositoryStorageMovable
     validate :container_repository_writable, on: :create
 
     default_value_for(:destination_storage_name, allows_nil: false) do
-      pick_repository_storage
+      Repository.pick_storage_shard
     end
 
     state_machine initial: :initial do
@@ -82,16 +82,6 @@ module RepositoryStorageMovable
     end
   end
 
-  class_methods do
-    private
-
-    def pick_repository_storage
-      container_klass = reflect_on_association(:container).class_name.constantize
-
-      container_klass.pick_repository_storage
-    end
-  end
-
   # Projects, snippets, and group wikis has different db structure. In projects,
   # we need to update some columns in this step, but we don't with the other resources.
   #

app/models/project.rb

@@ -75,7 +75,7 @@ class Project < ApplicationRecord
   default_value_for :resolve_outdated_diff_discussions, false
   default_value_for :container_registry_enabled, gitlab_config_features.container_registry
   default_value_for(:repository_storage) do
-    pick_repository_storage
+    Repository.pick_storage_shard
   end
 
   default_value_for(:shared_runners_enabled) { Gitlab::CurrentSettings.shared_runners_enabled }

app/models/repository.rb

@@ -1143,6 +1143,13 @@ class Repository
     end
   end
 
+  # Choose one of the available repository storage options based on a normalized weighted probability.
+  # We should always use the latest settings, to avoid picking a deleted shard.
+  def self.pick_storage_shard(expire: true)
+    Gitlab::CurrentSettings.expire_current_application_settings if expire
+    Gitlab::CurrentSettings.pick_repository_storage
+  end
+
   private
 
   # TODO Genericize finder, later split this on finders by Ref or Oid

app/models/snippet.rb

@@ -317,7 +317,7 @@ class Snippet < ApplicationRecord
   end
 
   def repository_storage
-    snippet_repository&.shard_name || self.class.pick_repository_storage
+    snippet_repository&.shard_name || Repository.pick_storage_shard
   end
 
   # Repositories are created by default with the `master` branch.

ee/app/models/group_wiki.rb

@@ -24,7 +24,7 @@ class GroupWiki < Wiki
   override :repository_storage
   def repository_storage
     strong_memoize(:repository_storage) do
-      container.group_wiki_repository&.shard_name || self.class.pick_repository_storage
+      container.group_wiki_repository&.shard_name || Repository.pick_storage_shard
     end
   end
 

ee/spec/models/group_wiki_spec.rb

@@ -84,8 +84,8 @@ RSpec.describe GroupWiki do
           let(:shards) { (1..).each }
 
           before do
-            # Force pick_repository_storage to always return a different value
-            allow(Gitlab::CurrentSettings).to receive(:pick_repository_storage) { "storage-#{shards.next}" }
+            # Force pick_storage_shard to always return a different value
+            allow(Repository).to receive(:pick_storage_shard) { "storage-#{shards.next}" }
           end
 
           it 'always returns the same shard when called repeatedly' do

spec/models/project_spec.rb

@@ -1551,10 +1551,7 @@ RSpec.describe Project, factory_default: :keep do
     let(:project) { build(:project) }
 
     it 'picks storage from ApplicationSetting' do
-      expect_next_instance_of(ApplicationSetting) do |instance|
-        expect(instance).to receive(:pick_repository_storage).and_return('picked')
-      end
-      expect(described_class).to receive(:pick_repository_storage).and_call_original
+      expect(Repository).to receive(:pick_storage_shard).and_return('picked')
 
       expect(project.repository_storage).to eq('picked')
     end

spec/models/repository_spec.rb

@@ -3049,4 +3049,51 @@ RSpec.describe Repository do
       end
     end
   end
+
+  describe '.pick_storage_shard', :request_store do
+    before do
+      storages = {
+        'default' => Gitlab::GitalyClient::StorageSettings.new('path' => 'tmp/tests/repositories'),
+        'picked'  => Gitlab::GitalyClient::StorageSettings.new('path' => 'tmp/tests/repositories')
+      }
+
+      allow(Gitlab.config.repositories).to receive(:storages).and_return(storages)
+      stub_env('IN_MEMORY_APPLICATION_SETTINGS', 'false')
+      Gitlab::CurrentSettings.current_application_settings
+
+      update_storages({ 'picked' => 0, 'default' => 100 })
+    end
+
+    context 'when expire is false' do
+      it 'does not expire existing repository storage value' do
+        previous_storage = described_class.pick_storage_shard
+        expect(previous_storage).to eq('default')
+        expect(Gitlab::CurrentSettings).not_to receive(:expire_current_application_settings)
+
+        update_storages({ 'picked' => 100, 'default' => 0 })
+
+        new_storage = described_class.pick_storage_shard(expire: false)
+        expect(new_storage).to eq(previous_storage)
+      end
+    end
+
+    context 'when expire is true' do
+      it 'expires existing repository storage value' do
+        previous_storage = described_class.pick_storage_shard
+        expect(previous_storage).to eq('default')
+        expect(Gitlab::CurrentSettings).to receive(:expire_current_application_settings).and_call_original
+
+        update_storages({ 'picked' => 100, 'default' => 0 })
+
+        new_storage = described_class.pick_storage_shard(expire: true)
+        expect(new_storage).to eq('picked')
+      end
+    end
+
+    def update_storages(storage_hash)
+      settings = ApplicationSetting.last
+      settings.repository_storages_weighted = storage_hash
+      settings.save!
+    end
+  end
 end

spec/models/snippet_spec.rb

@@ -630,14 +630,10 @@ RSpec.describe Snippet do
     subject { snippet.repository_storage }
 
     before do
-      expect_next_instance_of(ApplicationSetting) do |instance|
-        expect(instance).to receive(:pick_repository_storage).and_return('picked')
-      end
+      expect(Repository).to receive(:pick_storage_shard).and_return('picked')
     end
 
     it 'returns repository storage from ApplicationSetting' do
-      expect(described_class).to receive(:pick_repository_storage).and_call_original
-
       expect(subject).to eq 'picked'
     end
 

spec/support/shared_examples/models/concerns/has_repository_shared_examples.rb

@@ -152,36 +152,4 @@ RSpec.shared_examples 'model with repository' do
     it { is_expected.to respond_to(:disk_path) }
     it { is_expected.to respond_to(:gitlab_shell) }
   end
-
-  describe '.pick_repository_storage' do
-    subject { described_class.pick_repository_storage }
-
-    before do
-      storages = {
-        'default' => Gitlab::GitalyClient::StorageSettings.new('path' => 'tmp/tests/repositories'),
-        'picked'  => Gitlab::GitalyClient::StorageSettings.new('path' => 'tmp/tests/repositories')
-      }
-      allow(Gitlab.config.repositories).to receive(:storages).and_return(storages)
-    end
-
-    it 'picks storage from ApplicationSetting' do
-      expect(Gitlab::CurrentSettings).to receive(:pick_repository_storage).and_return('picked')
-
-      expect(subject).to eq('picked')
-    end
-
-    it 'picks from the available storages based on weight', :request_store do
-      stub_env('IN_MEMORY_APPLICATION_SETTINGS', 'false')
-      Gitlab::CurrentSettings.expire_current_application_settings
-      Gitlab::CurrentSettings.current_application_settings
-
-      settings = ApplicationSetting.last
-      settings.repository_storages_weighted = { 'picked' => 100, 'default' => 0 }
-      settings.save!
-
-      expect(Gitlab::CurrentSettings.repository_storages_weighted).to eq({ 'default' => 100 })
-      expect(subject).to eq('picked')
-      expect(Gitlab::CurrentSettings.repository_storages_weighted).to eq({ 'default' => 0, 'picked' => 100 })
-    end
-  end
 end

spec/support/shared_examples/models/concerns/repository_storage_movable_shared_examples.rb

@@ -45,8 +45,8 @@ RSpec.shared_examples 'handles repository moves' do
     context 'destination_storage_name' do
       subject { build(repository_storage_factory_key) }
 
-      it 'picks storage from ApplicationSetting' do
-        expect(Gitlab::CurrentSettings).to receive(:pick_repository_storage).and_return('picked').at_least(:once)
+      it 'can pick new storage' do
+        expect(Repository).to receive(:pick_storage_shard).and_return('picked').at_least(:once)
 
         expect(subject.destination_storage_name).to eq('picked')
       end