Merge branch 'jc-pick-weighted-repository' into 'master'

Pick repository storage based on weight

Closes #220852

See merge request gitlab-org/gitlab!34095

app/models/application_setting_implementation.rb

@@ -297,10 +297,21 @@ module ApplicationSettingImplementation
     performance_bar_allowed_group_id.present?
   end
 
-  # Choose one of the available repository storage options. Currently all have
-  # equal weighting.
+  def normalized_repository_storage_weights
+    strong_memoize(:normalized_repository_storage_weights) do
+      weights_total = repository_storages_weighted.values.reduce(:+)
+
+      repository_storages_weighted.transform_values do |w|
+        next w if weights_total == 0
+
+        w.to_f / weights_total
+      end
+    end
+  end
+
+  # Choose one of the available repository storage options based on a normalized weighted probability.
   def pick_repository_storage
-    repository_storages.sample
+    normalized_repository_storage_weights.max_by { |_, weight| rand**(1.0 / weight) }.first
   end
 
   def runners_registration_token

changelogs/unreleased/jc-pick-weighted-repository.yml

+---
+title: Pick repository storage based on weight
+merge_request: 34095
+author:
+type: changed

spec/support/shared_examples/models/application_setting_shared_examples.rb

@@ -288,11 +288,37 @@ RSpec.shared_examples 'application settings examples' do
   end
 
   describe '#pick_repository_storage' do
-    it 'uses Array#sample to pick a random storage' do
-      array = double('array', sample: 'random')
-      expect(setting).to receive(:repository_storages).and_return(array)
+    before do
+      allow(setting).to receive(:repository_storages_weighted).and_return({ 'default' => 20, 'backup' => 80 })
+    end
+
+    it 'chooses repository based on weight' do
+      picked_storages = { 'default' => 0.0, 'backup' => 0.0 }
+      10_000.times { picked_storages[setting.pick_repository_storage] += 1 }
+
+      expect(((picked_storages['default'] / 10_000) * 100).round.to_i).to be_between(19, 21)
+      expect(((picked_storages['backup'] / 10_000) * 100).round.to_i).to be_between(79, 81)
+    end
+  end
+
+  describe '#normalized_repository_storage_weights' do
+    using RSpec::Parameterized::TableSyntax
 
-      expect(setting.pick_repository_storage).to eq('random')
+    where(:storages, :normalized) do
+      { 'default' => 0, 'backup' => 100 }   | { 'default' => 0.0, 'backup' => 1.0 }
+      { 'default' => 100, 'backup' => 100 } | { 'default' => 0.5, 'backup' => 0.5 }
+      { 'default' => 20, 'backup' => 80 }   | { 'default' => 0.2, 'backup' => 0.8 }
+      { 'default' => 0, 'backup' => 0 }     | { 'default' => 0.0, 'backup' => 0.0 }
+    end
+
+    with_them do
+      before do
+        allow(setting).to receive(:repository_storages_weighted).and_return(storages)
+      end
+
+      it 'normalizes storage weights' do
+        expect(setting.normalized_repository_storage_weights).to eq(normalized)
+      end
     end
   end
 

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

@@ -148,18 +148,18 @@ RSpec.shared_examples 'model with repository' do
       expect(subject).to eq('picked')
     end
 
-    it 'picks from the latest available storage', :request_store do
+    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 = %w(picked)
+      settings.repository_storages_weighted = { 'picked' => 100, 'default' => 0 }
       settings.save!
 
-      expect(Gitlab::CurrentSettings.repository_storages).to eq(%w(default))
+      expect(Gitlab::CurrentSettings.repository_storages_weighted).to eq({ 'default' => 100 })
       expect(subject).to eq('picked')
-      expect(Gitlab::CurrentSettings.repository_storages).to eq(%w(picked))
+      expect(Gitlab::CurrentSettings.repository_storages_weighted).to eq({ 'default' => 0, 'picked' => 100 })
     end
   end
 end