active_storage_quota 0.1.0

data/lib/active_storage_quota/errors.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # Base class for every error raised by this gem. Rescue this to catch them all.
+  class Error < StandardError; end
+
+  # Raised when the gem is configured with a value it cannot use, for example a
+  # negative +reservation_ttl+ or an unusable +limit:+ option.
+  class ConfigurationError < Error; end
+
+  # Raised by +reserve_storage!+ when the requested bytes do not fit within the
+  # limit that was resolved for that reservation.
+  #
+  # Carries the numbers involved so callers can render a message without
+  # re-querying:
+  #
+  #   rescue ActiveStorageQuota::QuotaExceeded => e
+  #     render json: { error: e.message, remaining: e.remaining_bytes }, status: :payload_too_large
+  #   end
+  class QuotaExceeded < Error
+    attr_reader :owner, :scope_name, :requested_bytes, :limit_bytes, :used_bytes, :reserved_bytes
+
+    def initialize(owner:, scope_name:, requested_bytes:, limit_bytes:, used_bytes:, reserved_bytes:)
+      @owner           = owner
+      @scope_name      = scope_name
+      @requested_bytes = requested_bytes
+      @limit_bytes     = limit_bytes
+      @used_bytes      = used_bytes
+      @reserved_bytes  = reserved_bytes
+
+      super(build_message)
+    end
+
+    # Bytes still available under the resolved limit. Never negative.
+    def remaining_bytes
+      [ limit_bytes - used_bytes - reserved_bytes, 0 ].max
+    end
+
+    private
+      def build_message
+        "storage quota exceeded for #{owner.class.name}##{owner.id} " \
+          "(scope: #{scope_name}): requested #{requested_bytes} bytes, " \
+          "#{remaining_bytes} of #{limit_bytes} available " \
+          "(#{used_bytes} used, #{reserved_bytes} reserved)"
+      end
+  end
+
+  # Raised when a reservation cannot make the requested state transition, for
+  # example committing one that was already released.
+  class InvalidReservation < Error; end
+
+  # Raised when a model declares +storage_quota_owner+ but the owner resolves to
+  # nil for a record being attached to.
+  #
+  # Declaring the macro opts the model into quota accounting, so an unresolvable
+  # owner is a bug rather than a state to tolerate: without this, creating a
+  # record with no parent, attaching to it, and assigning the parent afterwards
+  # would store bytes nobody is charged for.
+  class MissingQuotaOwner < Error
+    attr_reader :record, :owner_source
+
+    def initialize(message = nil, record: nil, owner_source: nil)
+      @record = record
+      @owner_source = owner_source
+
+      super(message || "#{record.class.name}##{record.id} declares " \
+                       "storage_quota_owner #{owner_source.inspect}, which resolved to nil")
+    end
+  end
+
+  # Raised when backfill is asked to apply while some owners still have a real
+  # quota limit.
+  #
+  # Backfill writes charges without writing +used_bytes+, so counters understate
+  # the ledger until counter reconciliation runs. Enforcing against them in the
+  # meantime would admit uploads that should have been refused.
+  class EnforcementActive < Error
+    attr_reader :owners
+
+    def initialize(owners:)
+      @owners = owners
+
+      super("#{owners.size} owner(s) still have an active quota limit, and " \
+            "backfill leaves used_bytes incomplete until counter reconciliation " \
+            "runs. Let limits resolve to nil during adoption, or pass " \
+            "allow_active_enforcement: true to proceed anyway.")
+    end
+  end
+
+  # Raised when charge bookkeeping repeatedly loses a race to concurrent
+  # attach/detach activity on the same account and blob.
+  #
+  # Each attempt is correct in isolation; this means the row was created and
+  # destroyed underneath us more times than the bounded retry allows. On the
+  # attachment-create path it fails closed, rolling the attachment back.
+  class ConcurrentAccountingConflict < Error; end
+
+  # Raised when a reservation's bytes cannot be moved because the account no
+  # longer holds them: the reservation says N bytes are reserved, but
+  # +reserved_bytes+ is below N.
+  #
+  # This means the counters have drifted from the reservations that explain
+  # them -- a bug, or an external write. The transaction is rolled back rather
+  # than repaired: silently adjusting counters here would hide the drift and
+  # make the real cause unfindable. Reconciliation is a separate, deliberate
+  # operation.
+  class CounterInvariantViolation < Error; end
+
+  # Raised when committing a reservation whose +expires_at+ has passed. This is
+  # deterministic and does not depend on whether the expiry sweeper has run.
+  #
+  # A subclass of InvalidReservation so callers that treat every unusable
+  # reservation the same way can rescue just the parent.
+  class ReservationExpired < InvalidReservation; end
+end

data/lib/active_storage_quota/owner.rb

@@ -0,0 +1,154 @@
+# frozen_string_literal: true
+
+require "active_support/concern"
+
+module ActiveStorageQuota
+  # The model-side API, mixed into any Active Record class that declares
+  # +has_storage_quota+.
+  #
+  # Every reader here is ADVISORY. They describe the quota as it stood when
+  # they ran, which is exactly what a progress bar or an admin screen needs,
+  # and exactly what an admission check must not rely on:
+  #
+  #   # Racy. Two requests can both read the same free space and both proceed.
+  #   company.documents.attach(file) if company.storage_remaining >= file.size
+  #
+  # Admitting an upload is the job of the reservation API, which performs the
+  # check and the claim as one atomic statement.
+  module Owner
+    extend ActiveSupport::Concern
+
+    # Extended onto ActiveRecord::Base so every model can declare a quota.
+    module Macro
+      # Declares this model as a storage quota owner.
+      #
+      #   class Company < ApplicationRecord
+      #     has_storage_quota limit: 5.gigabytes
+      #   end
+      #
+      # +limit+ accepts an Integer number of bytes, a Symbol naming a method on
+      # the owner, a callable taking the owner or evaluated against it, or nil
+      # for unlimited. See ActiveStorageQuota::Definition.
+      #
+      # Subclasses may redeclare it; because the definition is held in a
+      # class_attribute, doing so does not affect the parent class.
+      def has_storage_quota(limit:)
+        include ActiveStorageQuota::Owner
+
+        self.storage_quota_definition = ActiveStorageQuota::Definition.new(limit: limit)
+      end
+    end
+
+    included do
+      # class_attribute, not a plain ivar: a subclass that redeclares its quota
+      # gets its own value instead of overwriting its parent's.
+      class_attribute :storage_quota_definition, instance_accessor: false
+
+      has_many :storage_quota_accounts,
+               as: :owner,
+               class_name: "ActiveStorageQuota::Account",
+               dependent: :destroy
+    end
+
+    # This owner's limit in bytes, or nil when unlimited.
+    def storage_limit
+      self.class.storage_quota_definition.resolve(self)
+    end
+
+    # Committed storage in bytes. Zero when no account row exists yet.
+    def storage_usage
+      storage_quota_account&.used_bytes || 0
+    end
+
+    # Storage claimed by in-flight uploads, in bytes. Zero when no account row
+    # exists yet.
+    def storage_reserved
+      storage_quota_account&.reserved_bytes || 0
+    end
+
+    # Bytes still available, floored at zero so a lowered limit reports no
+    # space rather than a negative number. Returns nil when unlimited.
+    #
+    # Advisory: see the note on this module. A value here is not permission to
+    # upload.
+    def storage_remaining
+      limit = storage_limit
+      return nil if limit.nil?
+
+      used, reserved = storage_quota_counters
+      [ limit - used - reserved, 0 ].max
+    end
+
+    # Whether this owner is currently over its limit. False when unlimited.
+    #
+    # This is a state reader, not an admission check: it answers "are we over
+    # budget right now", not "will the next upload fit". A limit lowered below
+    # existing usage makes this true, which is a normal state rather than an
+    # error.
+    def storage_quota_exceeded?
+      limit = storage_limit
+      return false if limit.nil?
+
+      used, reserved = storage_quota_counters
+      used + reserved > limit
+    end
+
+    # Claims +byte_size+ bytes against this owner's quota and returns the
+    # Reservation holding them.
+    #
+    #   reservation = company.reserve_storage!(file.size)
+    #   begin
+    #     company.documents.attach(file)
+    #     reservation.commit!
+    #   rescue StandardError
+    #     reservation.release!
+    #     raise
+    #   end
+    #
+    # This is the only race-safe way to admit an upload: the capacity test and
+    # the claim are a single statement, so concurrent callers across any number
+    # of processes cannot both be admitted into the same free space.
+    #
+    # Raises ActiveStorageQuota::QuotaExceeded if the bytes do not fit, and
+    # ArgumentError unless +byte_size+ is an exact non-negative Integer.
+    #
+    # The limit is resolved once, here, and enforced atomically against that
+    # value.
+    # Pass +blob:+ when holding capacity for a specific blob, as a direct-upload
+    # endpoint does. The attachment accounting then finds and consumes this hold
+    # instead of claiming capacity a second time.
+    def reserve_storage!(byte_size, blob: nil)
+      ActiveStorageQuota::Account.reserve!(
+        owner: self,
+        byte_size: byte_size,
+        limit_bytes: storage_limit,
+        blob_id: blob&.id
+      )
+    end
+
+    # The Account for +scope_name+, or nil when none exists.
+    #
+    # Reads never create rows, so an owner that has never stored anything costs
+    # nothing. When the association has already been loaded -- by
+    # Company.includes(:storage_quota_accounts), say -- the row is found in
+    # memory rather than with another query.
+    def storage_quota_account(scope_name = ActiveStorageQuota::DEFAULT_SCOPE)
+      scope_name = scope_name.to_s
+
+      if storage_quota_accounts.loaded?
+        storage_quota_accounts.find { |account| account.scope_name == scope_name }
+      else
+        storage_quota_accounts.find_by(scope_name: scope_name)
+      end
+    end
+
+    private
+      # Both counters from a single account lookup, so the readers that need
+      # both do not issue two queries for one answer.
+      def storage_quota_counters(scope_name = ActiveStorageQuota::DEFAULT_SCOPE)
+        account = storage_quota_account(scope_name)
+
+        [ account&.used_bytes || 0, account&.reserved_bytes || 0 ]
+      end
+  end
+end

data/lib/active_storage_quota/owner_resolver.rb

@@ -0,0 +1,402 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # Maps Active Storage attachments to the owners, and then the quota accounts,
+  # that should be charged for them.
+  #
+  # Internal. Shared by auditing, reconciliation and backfill so that reflection
+  # rules are decided in exactly one place; not part of the public API.
+  #
+  # Two paths. Where a model declares its payer as a Symbol naming a plain
+  # +belongs_to+, the record-to-owner step is a SQL join and no attachment or
+  # record is instantiated. Everywhere else, records are loaded in batches and
+  # the declared source is resolved in Ruby.
+  #
+  # The fast path is deliberately hard to qualify for. A +belongs_to+ can
+  # customise its primary key, its foreign key, its type column and its scope,
+  # and "the column is called company_id" does not mean it holds a Company's
+  # primary key. Anything that makes the mapping less than certain falls back,
+  # because a wrong mapping here would report healthy accounts as broken and,
+  # in reconciliation, could drive a destructive repair.
+  #
+  # Records are always looked up +unscoped+. A default_scope says what the host
+  # application wants to list, not whether a row exists: a soft-deleted or
+  # archived record still owns its attachment rows, and those files are still
+  # stored. Honouring the scope here made healthy charges look unexpected, let
+  # budgeted reconciliation delete them, and made backfill skip them. Only the
+  # attached record is unscoped -- its payer is still resolved through the
+  # declared association, default scope and all, and the SQL join honours that
+  # too: a foreign key only counts when it names a payer the association could
+  # load.
+  class OwnerResolver # :nodoc:
+    def initialize(result:, record_findings: true)
+      @result = result
+      @record_findings = record_findings
+      @fast_paths = {}
+      @managed_subclasses = {}
+      @unresolvable_types = {}
+      @reported_types = Set.new
+    end
+
+    # Every record_type in the data that is, or contains, quota-managed models.
+    # => { record_type => klass }
+    def managed_record_types
+      ActiveStorage::Attachment.distinct.pluck(:record_type).each_with_object({}) do |type, map|
+        klass = resolve_record_type(type)
+        map[type] = klass if klass
+      end
+    end
+
+    # => { [owner_type, owner_id, blob_id] => attachment_count }
+    #
+    # Stops short of quota accounts on purpose: backfill runs before most
+    # accounts exist and needs owner identities in order to create them.
+    def owner_pairs_for_blobs(blob_ids)
+      pairs = Hash.new(0)
+
+      rows_by_type(blob_ids).each do |record_type, rows|
+        klass = resolve_record_type(record_type)
+        next if klass.nil?
+
+        owner_keys = owner_keys_for(klass, rows.map { |row| row[:record_id] }.uniq)
+
+        rows.each do |row|
+          owner_key = owner_keys[row[:record_id]]
+          next if owner_key.nil?
+
+          pairs[[ owner_key.first, owner_key.last, row[:blob_id].to_s ]] += 1
+        end
+      end
+
+      pairs
+    end
+
+    # => { [quota_account_id, blob_id] => attachment_count }
+    #
+    # Only counts attachments whose owner is quota-managed and already has an
+    # account; an owner without one is reported, never created.
+    def pairs_for_blobs(blob_ids)
+      owner_pairs = owner_pairs_for_blobs(blob_ids)
+      return {} if owner_pairs.empty?
+
+      accounts = accounts_by_owner(owner_pairs.keys.map { |type, id, _blob| [ type, id ] }.uniq)
+
+      owner_pairs.each_with_object(Hash.new(0)) do |((owner_type, owner_id, blob_id), count), map|
+        account_id = accounts[[ owner_type, owner_id ]]
+
+        if account_id.nil?
+          report_missing_account(owner_type, owner_id)
+          next
+        end
+
+        map[[ account_id, blob_id ]] += count
+      end
+    end
+
+    # Records of +klass+ whose declared payer is +owner+. Uses the SQL join when
+    # the reflection qualifies, and resolves in bounded batches otherwise.
+    def record_ids_for_owner(klass, owner)
+      fast_path = fast_path_for(klass)
+
+      if fast_path && fast_path[:owner_type] == owner.class.polymorphic_name
+        klass.unscoped.where(fast_path[:foreign_key] => owner.id).pluck(klass.primary_key)
+      else
+        slow_record_ids_for_owner(klass, owner)
+      end
+    end
+
+    private
+      attr_reader :result
+
+      def record_findings?
+        @record_findings
+      end
+
+      def rows_by_type(blob_ids)
+        ActiveStorage::Attachment
+          .where(blob_id: blob_ids)
+          .pluck(:record_type, :record_id, :blob_id)
+          .group_by(&:first)
+          .transform_values do |rows|
+            rows.map { |(_type, record_id, blob_id)| { record_id: record_id, blob_id: blob_id } }
+          end
+      end
+
+      # Keeps rows whose foreign key names a payer that record.<association>
+      # would actually return: the payer model's own scope applies, so a payer
+      # it hides -- or one deleted outright, leaving the key dangling -- resolves
+      # to nothing, exactly as it does when records are loaded. The payer is
+      # never looked up unscoped.
+      def with_available_payer(scope, fast_path)
+        payer = fast_path[:owner_class]
+        scope.where(fast_path[:foreign_key] => payer.select(payer.primary_key))
+      end
+
+      # Rails stores polymorphic_name, so safe_constantize is the right inverse.
+      # A type that no longer resolves is a finding, not a crash -- renamed and
+      # deleted models are exactly what an audit exists to find.
+      def resolve_record_type(record_type)
+        klass = record_type.safe_constantize
+
+        if klass.nil?
+          report_once(record_type) do
+            result.record(:unresolvable_record_type, record_type: record_type)
+          end
+          return nil
+        end
+
+        return klass if quota_managed?(klass)
+
+        # The base class declares no payer, but single table inheritance means
+        # record_type holds the base name while the rows may belong to
+        # subclasses that do. Asking the data which subclasses exist avoids both
+        # eager-loading the host application and silently dropping their
+        # attachments on the floor.
+        return klass if managed_subclasses?(klass)
+
+        # Framework-owned records such as a blob's preview_image live here. Not
+        # declaring a payer is the normal case, not a defect.
+        result.ignore_record_type(record_type)
+        nil
+      end
+
+      def managed_subclasses?(klass)
+        @managed_subclasses.fetch(klass.name) do
+          @managed_subclasses[klass.name] = compute_managed_subclasses(klass)
+        end
+      end
+
+      def compute_managed_subclasses(klass)
+        return false unless sti_enabled?(klass)
+
+        # Loaded subclasses answer for free; otherwise ask the table.
+        return true if klass.descendants.any? { |subclass| quota_managed?(subclass) }
+
+        subclass_names(klass).any? do |name|
+          subclass = name.safe_constantize
+
+          if subclass.nil?
+            report_once("sti:#{name}") do
+              result.record(:unresolvable_record_type, record_type: name)
+            end
+            next false
+          end
+
+          quota_managed?(subclass)
+        end
+      end
+
+      def subclass_names(klass)
+        klass.unscoped.distinct.pluck(klass.inheritance_column).compact
+      rescue StandardError
+        []
+      end
+
+      def sti_enabled?(klass)
+        klass.respond_to?(:inheritance_column) &&
+          klass.respond_to?(:column_names) &&
+          klass.column_names.include?(klass.inheritance_column)
+      rescue StandardError
+        false
+      end
+
+      def quota_managed?(klass)
+        klass.respond_to?(:storage_quota_owner_source) &&
+          klass.storage_quota_owner_source.present?
+      end
+
+      def report_once(key)
+        return unless record_findings?
+        return unless @reported_types.add?(key)
+
+        yield
+      end
+
+      def report_missing_account(owner_type, owner_id)
+        return unless record_findings?
+
+        report_once("account:#{owner_type}:#{owner_id}") do
+          result.record(:missing_quota_account, owner_type: owner_type, owner_id: owner_id)
+        end
+      end
+
+      # => { record_id => [owner_type, owner_id_string] }
+      def owner_keys_for(klass, record_ids)
+        fast_path = fast_path_for(klass)
+
+        if fast_path
+          resolve_by_join(klass, fast_path, record_ids)
+        else
+          resolve_by_loading(klass, record_ids)
+        end
+      end
+
+      def resolve_by_join(klass, fast_path, record_ids)
+        with_available_payer(klass.unscoped.where(klass.primary_key => record_ids), fast_path)
+          .pluck(klass.primary_key, fast_path[:foreign_key])
+          .each_with_object({}) do |(record_id, owner_id), map|
+            next if owner_id.nil?   # nil owners are reported below
+
+            map[record_id] = [ fast_path[:owner_type], owner_id.to_s ]
+          end
+          .tap { |map| report_unowned(klass, record_ids - map.keys) }
+      end
+
+      # Loading through the base class instantiates each record as its real
+      # class, so an STI subclass that declares a different payer resolves
+      # through *its own* source. Reading the source off the base class here
+      # would silently attribute every subclass to the base's owner.
+      def resolve_by_loading(klass, record_ids)
+        map = {}
+
+        loadable_scope(klass, record_ids).find_each(batch_size: 500) do |record|
+          source = source_for(record)
+          # An STI subclass with no payer, under a base that has one. Nothing to
+          # charge, nothing wrong.
+          next if source.nil?
+
+          owner = source.resolve_safely(record)
+
+          if owner.nil?
+            report_unresolved_owner(record, source)
+            next
+          end
+
+          map[record.public_send(klass.primary_key)] =
+            [ owner.class.polymorphic_name, owner.id.to_s ]
+        end
+
+        map
+      end
+
+      def slow_record_ids_for_owner(klass, owner)
+        ids = []
+
+        loadable_scope(klass, nil).find_each(batch_size: 500) do |record|
+          source = source_for(record)
+          next if source.nil?
+
+          ids << record.public_send(klass.primary_key) if source.resolve_safely(record) == owner
+        end
+
+        ids
+      end
+
+      # A row whose STI type no longer constantizes raises SubclassNotFound on
+      # instantiation, which would abort a whole run over one bad row. Excluding
+      # those rows explicitly keeps things going and reports them.
+      #
+      # A nil record_ids means every record of this class.
+      def loadable_scope(klass, record_ids)
+        scope = record_ids.nil? ? klass.unscoped : klass.unscoped.where(klass.primary_key => record_ids)
+        unresolvable = unresolvable_types(klass)
+        return scope if unresolvable.empty?
+
+        scope.where.not(klass.inheritance_column => unresolvable)
+      end
+
+      def unresolvable_types(klass)
+        return [] unless sti_enabled?(klass)
+
+        @unresolvable_types.fetch(klass.name) do
+          names = subclass_names(klass).reject { |name| name.safe_constantize }
+          names.each do |name|
+            report_once("sti:#{name}") do
+              result.record(:unresolvable_record_type, record_type: name)
+            end
+          end
+
+          @unresolvable_types[klass.name] = names
+        end
+      end
+
+      def source_for(record)
+        return nil unless record.class.respond_to?(:storage_quota_owner_source)
+
+        record.class.storage_quota_owner_source
+      end
+
+      # A record whose declared owner is nil. On the join path this is a null
+      # foreign key; on the loading path resolve_safely returned nothing.
+      def report_unowned(klass, missing_ids)
+        return unless record_findings?
+
+        missing_ids.each do |record_id|
+          result.record(:missing_quota_owner,
+                        record_type: klass.polymorphic_name, record_id: record_id.to_s,
+                        owner_source: klass.storage_quota_owner_source.inspect)
+        end
+      end
+
+      def report_unresolved_owner(record, source)
+        return unless record_findings?
+
+        category = raw_owner_present?(record, source) ? :invalid_quota_owner : :missing_quota_owner
+
+        result.record(category,
+                      record_type: record.class.polymorphic_name, record_id: record.id.to_s,
+                      owner_source: source.inspect)
+      end
+
+      # resolve_safely collapses "no owner" and "owner is not a quota owner"
+      # into nil. Asking again, without the quota check, separates them.
+      def raw_owner_present?(record, source)
+        raw = source.source
+        value = raw.is_a?(Symbol) ? record.public_send(raw) : raw.call(record)
+        !value.nil?
+      rescue StandardError
+        false
+      end
+
+      def accounts_by_owner(owner_keys)
+        scope = ActiveStorageQuota::Account.where(scope_name: ActiveStorageQuota::DEFAULT_SCOPE)
+        conditions = owner_keys.group_by(&:first).map do |owner_type, keys|
+          scope.where(owner_type: owner_type, owner_id: keys.map(&:last))
+        end
+        return {} if conditions.empty?
+
+        conditions
+          .flat_map { |relation| relation.pluck(:owner_type, :owner_id, :id) }
+          .each_with_object({}) { |(type, id, account_id), map| map[[ type, id ]] = account_id }
+      end
+
+      # Nil unless a SQL join is *certainly* equivalent to resolving the declared
+      # source in Ruby.
+      def fast_path_for(klass)
+        @fast_paths.fetch(klass.name) { @fast_paths[klass.name] = compute_fast_path(klass) }
+      end
+
+      def compute_fast_path(klass)
+        return nil unless quota_managed?(klass)
+
+        source = klass.storage_quota_owner_source.source
+        return nil unless source.is_a?(Symbol)
+
+        # Single table inheritance: record_type holds the base class name, so the
+        # rows of one type may belong to subclasses that declare different
+        # payers. Establishing otherwise would mean loading every subclass,
+        # which this deliberately does not do.
+        return nil if klass.column_names.include?(klass.inheritance_column)
+
+        reflection = klass.reflect_on_association(source)
+        return nil unless reflection&.macro == :belongs_to
+        return nil if reflection.polymorphic?
+        return nil if reflection.scope
+
+        owner_klass = reflection.klass
+        return nil unless owner_klass.respond_to?(:storage_quota_definition)
+
+        # belongs_to :company, primary_key: :external_id means the foreign key
+        # holds an external_id, not a Company id. Joining on it would key the
+        # quota account to the wrong value.
+        return nil unless reflection.association_primary_key.to_s == owner_klass.primary_key.to_s
+
+        { foreign_key: reflection.foreign_key.to_s, owner_type: owner_klass.polymorphic_name,
+          owner_class: owner_klass }
+      rescue NameError
+        # class_name points at a constant that no longer exists. Fall back rather
+        # than guess; the ledger audit will report the attachments.
+        nil
+      end
+  end
+end