Calculate storage statistics for dependency proxy

Calculates namespace statistic storage size for Dependency Proxy assets
at both the namespace and root namespace storage statistics level

Changelog: added

app/models/concerns/update_namespace_statistics.rb

+# frozen_string_literal: true
+
+# This module provides helpers for updating `NamespaceStatistics` with `after_save` and
+# `after_destroy` hooks.
+#
+# Models including this module must respond to and return a `namespace`
+#
+# Example:
+#
+# class DependencyProxy::Manifest
+#   include UpdateNamespaceStatistics
+#
+#   belongs_to :group
+#   alias_attribute :namespace, :group
+#
+#   update_namespace_statistics namespace_statistics_name: :dependency_proxy_size
+# end
+module UpdateNamespaceStatistics
+  extend ActiveSupport::Concern
+  include AfterCommitQueue
+
+  class_methods do
+    attr_reader :namespace_statistics_name, :statistic_attribute
+
+    # Configure the model to update `namespace_statistics_name` on NamespaceStatistics,
+    # when `statistic_attribute` changes
+    #
+    # - namespace_statistics_name: A column of `NamespaceStatistics` to update
+    # - statistic_attribute: An attribute of the current model, default to `size`
+    def update_namespace_statistics(namespace_statistics_name:, statistic_attribute: :size)
+      @namespace_statistics_name = namespace_statistics_name
+      @statistic_attribute = statistic_attribute
+
+      after_save(:schedule_namespace_statistics_refresh, if: :update_namespace_statistics?)
+      after_destroy(:schedule_namespace_statistics_refresh)
+    end
+
+    private :update_namespace_statistics
+  end
+
+  included do
+    private
+
+    def update_namespace_statistics?
+      saved_change_to_attribute?(self.class.statistic_attribute)
+    end
+
+    def schedule_namespace_statistics_refresh
+      run_after_commit do
+        Groups::UpdateStatisticsWorker.perform_async(namespace.id, [self.class.namespace_statistics_name])
+      end
+    end
+  end
+end

app/models/dependency_proxy/blob.rb

@@ -5,8 +5,10 @@ class DependencyProxy::Blob < ApplicationRecord
   include TtlExpirable
   include Packages::Destructible
   include EachBatch
+  include UpdateNamespaceStatistics
 
   belongs_to :group
+  alias_attribute :namespace, :group
 
   MAX_FILE_SIZE = 5.gigabytes.freeze
 
@@ -17,6 +19,7 @@ class DependencyProxy::Blob < ApplicationRecord
   scope :with_files_stored_locally, -> { where(file_store: ::DependencyProxy::FileUploader::Store::LOCAL) }
 
   mount_file_store_uploader DependencyProxy::FileUploader
+  update_namespace_statistics namespace_statistics_name: :dependency_proxy_size
 
   def self.total_size
     sum(:size)

app/models/dependency_proxy/manifest.rb

@@ -5,8 +5,10 @@ class DependencyProxy::Manifest < ApplicationRecord
   include TtlExpirable
   include Packages::Destructible
   include EachBatch
+  include UpdateNamespaceStatistics
 
   belongs_to :group
+  alias_attribute :namespace, :group
 
   MAX_FILE_SIZE = 10.megabytes.freeze
   DIGEST_HEADER = 'Docker-Content-Digest'
@@ -20,6 +22,7 @@ class DependencyProxy::Manifest < ApplicationRecord
   scope :with_files_stored_locally, -> { where(file_store: ::DependencyProxy::FileUploader::Store::LOCAL) }
 
   mount_file_store_uploader DependencyProxy::FileUploader
+  update_namespace_statistics namespace_statistics_name: :dependency_proxy_size
 
   def self.find_by_file_name_or_digest(file_name:, digest:)
     find_by(file_name: file_name) || find_by(digest: digest)

spec/controllers/groups/dependency_proxy_for_containers_controller_spec.rb

@@ -170,6 +170,14 @@ RSpec.describe Groups::DependencyProxyForContainersController do
     end
   end
 
+  shared_examples 'namespace statistics refresh' do
+    it 'updates namespace statistics' do
+      expect(Groups::UpdateStatisticsWorker).to receive(:perform_async)
+
+      subject
+    end
+  end
+
   before do
     allow(Gitlab.config.dependency_proxy)
       .to receive(:enabled).and_return(true)
@@ -403,13 +411,15 @@ RSpec.describe Groups::DependencyProxyForContainersController do
     context 'with a valid user' do
       before do
         group.add_guest(user)
-
-        expect_next_found_instance_of(Group) do |instance|
-          expect(instance).to receive_message_chain(:dependency_proxy_blobs, :create!)
-        end
       end
 
       it_behaves_like 'a package tracking event', described_class.name, 'pull_blob'
+
+      it 'creates a blob' do
+        expect { subject }.to change { group.dependency_proxy_blobs.count }.by(1)
+      end
+
+      it_behaves_like 'namespace statistics refresh'
     end
   end
 
@@ -473,6 +483,8 @@ RSpec.describe Groups::DependencyProxyForContainersController do
           expect(manifest.digest).to eq(digest)
           expect(manifest.file_name).to eq(file_name)
         end
+
+        it_behaves_like 'namespace statistics refresh'
       end
 
       context 'with existing stale manifest' do
@@ -483,6 +495,8 @@ RSpec.describe Groups::DependencyProxyForContainersController do
           expect { subject }.to change { group.dependency_proxy_manifests.count }.by(0)
             .and change { manifest.reload.digest }.from(old_digest).to(digest)
         end
+
+        it_behaves_like 'namespace statistics refresh'
       end
     end
   end

spec/models/dependency_proxy/blob_spec.rb

@@ -5,6 +5,10 @@ RSpec.describe DependencyProxy::Blob, type: :model do
   it_behaves_like 'ttl_expirable'
   it_behaves_like 'destructible', factory: :dependency_proxy_blob
 
+  it_behaves_like 'updates namespace statistics' do
+    let(:statistic_source) { build(:dependency_proxy_blob, size: 10) }
+  end
+
   describe 'relationships' do
     it { is_expected.to belong_to(:group) }
   end

spec/models/dependency_proxy/manifest_spec.rb

@@ -5,6 +5,10 @@ RSpec.describe DependencyProxy::Manifest, type: :model do
   it_behaves_like 'ttl_expirable'
   it_behaves_like 'destructible', factory: :dependency_proxy_manifest
 
+  it_behaves_like 'updates namespace statistics' do
+    let(:statistic_source) { build(:dependency_proxy_manifest, size: 10) }
+  end
+
   describe 'relationships' do
     it { is_expected.to belong_to(:group) }
   end

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

+# frozen_string_literal: true
+
+RSpec.shared_examples 'updates namespace statistics' do
+  let(:namespace_statistics_name) { described_class.namespace_statistics_name }
+  let(:statistic_attribute) { described_class.statistic_attribute }
+
+  context 'when creating' do
+    before do
+      statistic_source.send("#{statistic_attribute}=", 10)
+    end
+
+    it 'schedules a statistic refresh' do
+      expect(Groups::UpdateStatisticsWorker)
+        .to receive(:perform_async)
+
+      statistic_source.save!
+    end
+  end
+
+  context 'when updating' do
+    before do
+      statistic_source.save!
+
+      expect(statistic_source).to be_persisted
+    end
+
+    context 'when the statistic attribute has not changed' do
+      it 'does not schedule a statistic refresh' do
+        expect(Groups::UpdateStatisticsWorker)
+          .not_to receive(:perform_async)
+
+        statistic_source.update!(file_name: 'new-file-name.txt')
+      end
+    end
+
+    context 'when the statistic attribute has changed' do
+      it 'schedules a statistic refresh' do
+        expect(Groups::UpdateStatisticsWorker)
+          .to receive(:perform_async)
+
+        statistic_source.update!(statistic_attribute => 20)
+      end
+    end
+  end
+
+  context 'when deleting' do
+    it 'schedules a statistic refresh' do
+      expect(Groups::UpdateStatisticsWorker)
+        .to receive(:perform_async)
+
+      statistic_source.destroy!
+    end
+  end
+end