Add storage enforcement with excess storage logic

Add storage enforcement based on purchased excess storage. This logic is
behind a feature flag.

app/models/snippet.rb

@@ -282,7 +282,9 @@ class Snippet < ApplicationRecord
     strong_memoize(:repository_size_checker) do
       ::Gitlab::RepositorySizeChecker.new(
         current_size_proc: -> { repository.size.megabytes },
-        limit: Gitlab::CurrentSettings.snippet_size_limit
+        limit: Gitlab::CurrentSettings.snippet_size_limit,
+        total_repository_size_excess: nil,
+        additional_purchased_storage: nil
       )
     end
   end

ee/app/models/ee/project.rb

@@ -30,7 +30,7 @@ module EE
       after_update :remove_mirror_repository_reference,
         if: ->(project) { project.mirror? && project.import_url_updated? }
 
-      belongs_to :mirror_user, foreign_key: 'mirror_user_id', class_name: 'User'
+      belongs_to :mirror_user, class_name: 'User'
       belongs_to :deleting_user, foreign_key: 'marked_for_deletion_by_user_id', class_name: 'User'
 
       has_one :repository_state, class_name: 'ProjectRepositoryState', inverse_of: :project
@@ -527,6 +527,8 @@ module EE
         ::Gitlab::RepositorySizeChecker.new(
           current_size_proc: -> { statistics.total_repository_size },
           limit: actual_size_limit,
+          total_repository_size_excess: namespace.total_repository_size_excess,
+          additional_purchased_storage: namespace.additional_purchased_storage_size.megabytes,
           enabled: License.feature_available?(:repository_size_limit)
         )
       end

ee/app/models/ee/snippet.rb

@@ -3,10 +3,23 @@
 module EE
   module Snippet
     extend ActiveSupport::Concern
+    extend ::Gitlab::Utils::Override
 
     prepended do
       include Elastic::SnippetsSearch
       include UsageStatistics
     end
+
+    override :repository_size_checker
+    def repository_size_checker
+      strong_memoize(:repository_size_checker) do
+        ::Gitlab::RepositorySizeChecker.new(
+          current_size_proc: -> { repository.size.megabytes },
+          limit: ::Gitlab::CurrentSettings.snippet_size_limit,
+          total_repository_size_excess: project&.namespace&.total_repository_size_excess,
+          additional_purchased_storage: project&.namespace&.additional_purchased_storage_size&.megabytes
+        )
+      end
+    end
   end
 end

ee/lib/gitlab/repository_size_checker.rb

+# frozen_string_literal: true
+
+module Gitlab
+  module RepositorySizeChecker
+    extend ActiveSupport::Concern
+    extend ::Gitlab::Utils::Override
+
+    override :above_size_limit?
+    def above_size_limit?
+      return false unless enabled?
+      return false if additional_repo_storage_available? && total_repository_size_excess <= additional_purchased_storage
+
+      super
+    end
+
+    override :exceeded_size
+    # @param change_size [int] in bytes
+    def exceeded_size(change_size = 0)
+      exceeded_size = super
+      exceeded_size -= remaining_additional_purchased_storage if additional_repo_storage_available?
+      exceeded_size
+    end
+
+    private
+
+    def additional_repo_storage_available?
+      return false unless Gitlab::CurrentSettings.automatic_purchased_storage_allocation?
+
+      Feature.enabled?(:additional_repo_storage_by_namespace)
+    end
+
+    def remaining_additional_purchased_storage
+      additional_purchased_storage - total_repository_size_excess
+    end
+  end
+end

ee/spec/lib/gitlab/repository_size_checker_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Gitlab::RepositorySizeChecker do
+  let(:current_size) { 0 }
+  let(:limit) { 50 }
+  let(:total_repository_size_excess) { 0 }
+  let(:additional_purchased_storage) { 0 }
+  let(:enabled) { true }
+  let(:gitlab_setting_enabled) { true }
+
+  subject do
+    described_class.new(
+      current_size_proc: -> { current_size },
+      limit: limit,
+      total_repository_size_excess: total_repository_size_excess,
+      additional_purchased_storage: additional_purchased_storage,
+      enabled: enabled
+    )
+  end
+
+  before do
+    allow(Gitlab::CurrentSettings).to receive(:automatic_purchased_storage_allocation?).and_return(gitlab_setting_enabled)
+  end
+
+  describe '#above_size_limit?' do
+    context 'when Gitlab app setting for automatic purchased storage allocation is not enabled' do
+      let(:gitlab_setting_enabled) { false }
+
+      include_examples 'checker size above limit'
+      include_examples 'checker size not over limit'
+    end
+
+    context 'with feature flag :additional_repo_storage_by_namespace enabled' do
+      context 'when there is available excess storage' do
+        it 'returns false' do
+          expect(subject.above_size_limit?).to eq(false)
+        end
+      end
+
+      context 'when size is above the limit and there is no exccess storage' do
+        let(:current_size) { 100 }
+        let(:total_repository_size_excess) { 20 }
+        let(:additional_purchased_storage) { 10 }
+
+        it 'returns true' do
+          expect(subject.above_size_limit?).to eq(true)
+        end
+      end
+
+      it 'returns false when not over the limit' do
+        expect(subject.above_size_limit?).to eq(false)
+      end
+    end
+
+    context 'with feature flag :additional_repo_storage_by_namespace disabled' do
+      before do
+        stub_feature_flags(additional_repo_storage_by_namespace: false)
+      end
+
+      include_examples 'checker size above limit'
+      include_examples 'checker size not over limit'
+    end
+  end
+
+  describe '#exceeded_size' do
+    context 'with feature flag :additional_repo_storage_by_namespace enabled' do
+      context 'when Gitlab app setting for automatic purchased storage allocation is not enabled' do
+        let(:gitlab_setting_enabled) { false }
+
+        include_examples 'checker size exceeded'
+      end
+
+      context 'when current size + total repository size excess are below or equal to the limit + additional purchased storage' do
+        let(:current_size) { 50 }
+        let(:total_repository_size_excess) { 10 }
+        let(:additional_purchased_storage) { 10 }
+
+        it 'returns zero' do
+          expect(subject.exceeded_size).to eq(0)
+        end
+      end
+
+      context 'when current size + total repository size excess are over the limit + additional purchased storage' do
+        let(:current_size) { 51 }
+        let(:total_repository_size_excess) { 10 }
+        let(:additional_purchased_storage) { 10 }
+
+        it 'returns 1' do
+          expect(subject.exceeded_size).to eq(1)
+        end
+      end
+
+      context 'when change size will be over the limit' do
+        let(:current_size) { 50 }
+
+        it 'returns 1' do
+          expect(subject.exceeded_size(1)).to eq(1)
+        end
+      end
+
+      context 'when change size will not be over the limit' do
+        let(:current_size) { 49 }
+
+        it 'returns zero' do
+          expect(subject.exceeded_size(1)).to eq(0)
+        end
+      end
+    end
+
+    context 'with feature flag :additional_repo_storage_by_namespace disabled' do
+      before do
+        stub_feature_flags(additional_repo_storage_by_namespace: false)
+      end
+
+      include_examples 'checker size exceeded'
+    end
+  end
+end

ee/spec/models/project_spec.rb

@@ -2235,6 +2235,22 @@ RSpec.describe Project do
       end
     end
 
+    describe '#total_repository_size_excess' do
+      it 'returns the total repository size excess of the namespace' do
+        allow(project.namespace).to receive(:total_repository_size_excess).and_return(50)
+
+        expect(checker.total_repository_size_excess).to eq(50)
+      end
+    end
+
+    describe '#additional_purchased_storage' do
+      it 'returns the additional purchased storage size of the namespace' do
+        allow(project.namespace).to receive(:additional_purchased_storage_size).and_return(100)
+
+        expect(checker.additional_purchased_storage).to eq(100.megabytes)
+      end
+    end
+
     describe '#enabled?' do
       it 'returns true when not equal to zero' do
         project.repository_size_limit = 1

ee/spec/models/snippet_spec.rb

+# frozen_string_literal: true
+
+require 'spec_helper'
+
+RSpec.describe Snippet do
+  describe '#repository_size_checker' do
+    let(:checker) { subject.repository_size_checker }
+    let(:current_size) { 60 }
+
+    before do
+      allow(subject.repository).to receive(:size).and_return(current_size)
+    end
+
+    context 'when snippet belongs to a project' do
+      subject { build(:project_snippet, project: project) }
+
+      let(:namespace) { build(:namespace, additional_purchased_storage_size: 50) }
+      let(:project) { build(:project, namespace: namespace) }
+      let(:total_repository_size_excess) { 100 }
+      let(:additional_purchased_storage) { 50.megabytes }
+
+      before do
+        allow(namespace).to receive(:total_repository_size_excess).and_return(total_repository_size_excess)
+      end
+
+      include_examples 'size checker for snippet'
+    end
+
+    context 'when snippet without a project' do
+      let(:total_repository_size_excess) { 0 }
+      let(:additional_purchased_storage) { 0 }
+
+      include_examples 'size checker for snippet'
+    end
+  end
+end

lib/gitlab/repository_size_checker.rb

@@ -3,12 +3,14 @@
 module Gitlab
   # Centralized class for repository size related calculations.
   class RepositorySizeChecker
-    attr_reader :limit
+    attr_reader :limit, :total_repository_size_excess, :additional_purchased_storage
 
     # @param current_size_proc [Proc] returns repository size in bytes
-    def initialize(current_size_proc:, limit:, enabled: true)
+    def initialize(current_size_proc:, limit:, total_repository_size_excess:, additional_purchased_storage:, enabled: true)
       @current_size_proc = current_size_proc
       @limit = limit
+      @total_repository_size_excess = total_repository_size_excess.to_i
+      @additional_purchased_storage = additional_purchased_storage.to_i
       @enabled = enabled && limit != 0
     end
 

lib/gitlab/repository_size_error_message.rb

@@ -4,7 +4,7 @@ module Gitlab
   class RepositorySizeErrorMessage
     include ActiveSupport::NumberHelper
 
-    delegate :current_size, :limit, :exceeded_size, to: :@checker
+    delegate :current_size, :limit, :total_repository_size_excess, :additional_purchased_storage, :exceeded_size, to: :@checker
 
     # @param checher [RepositorySizeChecker]
     def initialize(checker)

spec/lib/gitlab/repository_size_checker_spec.rb

@@ -11,6 +11,8 @@ RSpec.describe Gitlab::RepositorySizeChecker do
     described_class.new(
       current_size_proc: -> { current_size },
       limit: limit,
+      total_repository_size_excess: 0,
+      additional_purchased_storage: 0,
       enabled: enabled
     )
   end
@@ -44,50 +46,11 @@ RSpec.describe Gitlab::RepositorySizeChecker do
   end
 
   describe '#above_size_limit?' do
-    context 'when size is above the limit' do
-      let(:current_size) { 100 }
-
-      it 'returns true' do
-        expect(subject.above_size_limit?).to be_truthy
-      end
-    end
-
-    it 'returns false when not over the limit' do
-      expect(subject.above_size_limit?).to be_falsey
-    end
+    include_examples 'checker size above limit'
+    include_examples 'checker size not over limit'
   end
 
   describe '#exceeded_size' do
-    context 'when current size is below or equal to the limit' do
-      let(:current_size) { 50 }
-
-      it 'returns zero' do
-        expect(subject.exceeded_size).to eq(0)
-      end
-    end
-
-    context 'when current size is over the limit' do
-      let(:current_size) { 51 }
-
-      it 'returns zero' do
-        expect(subject.exceeded_size).to eq(1)
-      end
-    end
-
-    context 'when change size will be over the limit' do
-      let(:current_size) { 50 }
-
-      it 'returns zero' do
-        expect(subject.exceeded_size(1)).to eq(1)
-      end
-    end
-
-    context 'when change size will not be over the limit' do
-      let(:current_size) { 49 }
-
-      it 'returns zero' do
-        expect(subject.exceeded_size(1)).to eq(0)
-      end
-    end
+    include_examples 'checker size exceeded'
   end
 end

spec/lib/gitlab/repository_size_error_message_spec.rb

@@ -6,6 +6,8 @@ RSpec.describe Gitlab::RepositorySizeErrorMessage do
   let(:checker) do
     Gitlab::RepositorySizeChecker.new(
       current_size_proc: -> { 15.megabytes },
+      total_repository_size_excess: 0,
+      additional_purchased_storage: 0,
       limit: 10.megabytes
     )
   end

spec/models/snippet_spec.rb

@@ -666,16 +666,14 @@ RSpec.describe Snippet do
 
     let(:checker) { subject.repository_size_checker }
     let(:current_size) { 60 }
+    let(:total_repository_size_excess) { 0 }
+    let(:additional_purchased_storage) { 0 }
 
     before do
       allow(subject.repository).to receive(:size).and_return(current_size)
     end
 
-    it 'sets up size checker', :aggregate_failures do
-      expect(checker.current_size).to eq(current_size.megabytes)
-      expect(checker.limit).to eq(Gitlab::CurrentSettings.snippet_size_limit)
-      expect(checker.enabled?).to be_truthy
-    end
+    include_examples 'size checker for snippet'
   end
 
   describe '#can_cache_field?' do

spec/support/shared_examples/lib/gitlab/repository_size_checker_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'checker size above limit' do
+  context 'when size is above the limit' do
+    let(:current_size) { 100 }
+
+    it 'returns true' do
+      expect(subject.above_size_limit?).to eq(true)
+    end
+  end
+end
+
+RSpec.shared_examples 'checker size not over limit' do
+  it 'returns false when not over the limit' do
+    expect(subject.above_size_limit?).to eq(false)
+  end
+end
+
+RSpec.shared_examples 'checker size exceeded' do
+  context 'when current size is below or equal to the limit' do
+    let(:current_size) { 50 }
+
+    it 'returns zero' do
+      expect(subject.exceeded_size).to eq(0)
+    end
+  end
+
+  context 'when current size is over the limit' do
+    let(:current_size) { 51 }
+
+    it 'returns zero' do
+      expect(subject.exceeded_size).to eq(1)
+    end
+  end
+
+  context 'when change size will be over the limit' do
+    let(:current_size) { 50 }
+
+    it 'returns zero' do
+      expect(subject.exceeded_size(1)).to eq(1)
+    end
+  end
+
+  context 'when change size will not be over the limit' do
+    let(:current_size) { 49 }
+
+    it 'returns zero' do
+      expect(subject.exceeded_size(1)).to eq(0)
+    end
+  end
+end

spec/support/shared_examples/models/snippet_shared_examples.rb

+# frozen_string_literal: true
+
+RSpec.shared_examples 'size checker for snippet' do |action|
+  it 'sets up size checker', :aggregate_failures do
+    expect(checker.current_size).to eq(current_size.megabytes)
+    expect(checker.limit).to eq(Gitlab::CurrentSettings.snippet_size_limit)
+    expect(checker.total_repository_size_excess).to eq(total_repository_size_excess)
+    expect(checker.additional_purchased_storage).to eq(additional_purchased_storage)
+    expect(checker.enabled?).to eq(true)
+  end
+end