active_storage_quota 0.1.0

data/lib/active_storage_quota/owner_source.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # Resolves the quota owner declared by +storage_quota_owner+.
+  #
+  #   storage_quota_owner :company            # Symbol: association or method
+  #   storage_quota_owner ->(record) { ... }  # callable taking the record
+  #   storage_quota_owner -> { company }      # callable evaluated on the record
+  #
+  # A Symbol naming a +belongs_to+ is worth preferring: it is the form a future
+  # set-based backfill can express as a SQL join, where a callable can only be
+  # resolved one record at a time.
+  class OwnerSource
+    SUPPORTED_ARITIES = [ 0, 1 ].freeze
+
+    attr_reader :source
+
+    def initialize(source)
+      @source = source
+      validate_declaration!
+    end
+
+    # The quota owner for +record+.
+    #
+    # Raises MissingQuotaOwner when it resolves to nil, and ConfigurationError
+    # when the result is not a quota owner.
+    def resolve(record)
+      owner = call_source(record)
+
+      raise MissingQuotaOwner.new(record: record, owner_source: source) if owner.nil?
+
+      unless quota_owner?(owner)
+        raise ConfigurationError,
+              "#{record.class.name} declares storage_quota_owner #{source.inspect}, " \
+              "which returned #{owner.class.name}; that model does not declare " \
+              "has_storage_quota"
+      end
+
+      owner
+    end
+
+    # The quota owner for +record+, or nil for any reason at all.
+    #
+    # Used where accounting must not raise: destruction bookkeeping, which fails
+    # open, and auditing, which has to survive malformed historical data in
+    # order to report it.
+    def resolve_safely(record)
+      owner = call_source(record)
+      owner if quota_owner?(owner)
+    rescue StandardError
+      nil
+    end
+
+    def inspect
+      source.inspect
+    end
+
+    private
+      def call_source(record)
+        case source
+        when Symbol then record.public_send(source)
+        else
+          source.arity.zero? ? record.instance_exec(&source) : source.call(record)
+        end
+      end
+
+      def quota_owner?(owner)
+        !owner.nil? && owner.respond_to?(:storage_quota_account)
+      end
+
+      def validate_declaration!
+        return if source.is_a?(Symbol)
+
+        unless source.respond_to?(:call)
+          raise ConfigurationError,
+                "storage_quota_owner must be a Symbol naming a method on the " \
+                "record, or a callable; got #{source.inspect} (#{source.class})"
+        end
+
+        return if SUPPORTED_ARITIES.include?(source.arity)
+
+        raise ConfigurationError,
+              "storage_quota_owner: a callable must take no arguments " \
+              "(evaluated against the record) or one argument (the record); " \
+              "got arity #{source.arity}"
+      end
+  end
+end

data/lib/active_storage_quota/reconciliation/budget.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  module Reconciliation
+    # How much destructive repair one reconciliation invocation is permitted.
+    #
+    # There is no default budget on purpose. Deleting a charge lowers a
+    # customer's recorded usage, and an over-generous default would let a bug in
+    # owner resolution silently hand out storage. With no limits supplied,
+    # unexpected charges are reported and left alone.
+    #
+    # The asymmetry with additive repair is deliberate: over-charging announces
+    # itself when the customer hits their quota and is trivially reversible,
+    # while under-charging is silent and may never be noticed.
+    class Budget
+      attr_reader :max_charge_deletions, :max_removed_fraction
+
+      def initialize(max_charge_deletions: nil, max_removed_fraction: nil)
+        @max_charge_deletions = validate_deletions(max_charge_deletions)
+        @max_removed_fraction = validate_fraction(max_removed_fraction)
+      end
+
+      # Destructive repair happens only when explicitly budgeted.
+      def destructive?
+        !max_charge_deletions.nil? || !max_removed_fraction.nil?
+      end
+
+      # Whether removing this much would break the budget. Used twice: once on
+      # the whole account before any deletion starts, and again before each
+      # individual deletion, because the account is live and moves underneath us.
+      def exceeded?(deletions:, bytes_removed:, used_bytes:)
+        return true if max_charge_deletions && deletions > max_charge_deletions
+
+        if max_removed_fraction && used_bytes.to_i.positive?
+          return true if bytes_removed > used_bytes * max_removed_fraction
+        end
+
+        false
+      end
+
+      def to_h
+        { max_charge_deletions: max_charge_deletions, max_removed_fraction: max_removed_fraction }
+      end
+
+      private
+        def validate_deletions(value)
+          return nil if value.nil?
+
+          unless value.is_a?(Integer) && value.positive?
+            raise ArgumentError,
+                  "max_charge_deletions must be a positive Integer, got #{value.inspect}"
+          end
+
+          value
+        end
+
+        def validate_fraction(value)
+          return nil if value.nil?
+
+          unless value.is_a?(Numeric) && value > 0 && value <= 1
+            raise ArgumentError,
+                  "max_removed_fraction must be a number greater than 0 and at most 1, " \
+                  "got #{value.inspect}"
+          end
+
+          value
+        end
+    end
+  end
+end

data/lib/active_storage_quota/reconciliation/counter_reconciler.rb

@@ -0,0 +1,78 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  module Reconciliation
+    # Level 1: repairs Account#used_bytes from the charges it aggregates.
+    #
+    # Cheap and safe to run often. Touches no Active Storage table, so it works
+    # regardless of how Active Storage is configured or whether it is loaded.
+    #
+    # The account row is locked FOR UPDATE before the sum is taken, which is
+    # what stops this race:
+    #
+    #   reconciler reads SUM = 100
+    #   an upload commits, used_bytes += 20
+    #   reconciler writes 100          <- the 20 would be lost
+    #
+    # Charge.record! and Charge.discard! both reach used_bytes through an UPDATE
+    # of that same row, so they block behind the lock and apply on top of the
+    # value written here rather than being overwritten by it.
+    #
+    # Locks only the account. Charges are read, never locked, so this can never
+    # deadlock against ledger reconciliation's charge-then-account order.
+    class CounterReconciler
+      def initialize(result:, dry_run:)
+        @result = result
+        @dry_run = dry_run
+      end
+
+      def run(accounts)
+        accounts.each { |account| reconcile(account) }
+      end
+
+      private
+        attr_reader :result
+
+        def reconcile(account)
+          result.increment(:accounts_checked)
+
+          if @dry_run
+            plan(account)
+          else
+            apply(account)
+          end
+        end
+
+        def plan(account)
+          stored = account.reload.used_bytes
+          expected = expected_used_bytes(account.id)
+          return if stored == expected
+
+          result.increment(:counters_corrected)
+          result.record(:planned_change, action: :correct_counter, quota_account_id: account.id,
+                                         from: stored, to: expected)
+        end
+
+        def apply(account)
+          Account.transaction(requires_new: true) do
+            locked = Account.lock.find_by(id: account.id)
+            next if locked.nil?
+
+            expected = expected_used_bytes(locked.id)
+            next if locked.used_bytes == expected
+
+            Account.where(id: locked.id)
+                   .update_all([ "used_bytes = ?, updated_at = ?", expected, Time.current ])
+
+            result.increment(:counters_corrected)
+            result.record(:counter_corrected, quota_account_id: locked.id,
+                                              from: locked.used_bytes, to: expected)
+          end
+        end
+
+        def expected_used_bytes(account_id)
+          Charge.where(quota_account_id: account_id).sum(:byte_size)
+        end
+    end
+  end
+end

data/lib/active_storage_quota/reconciliation/ledger_reconciler.rb

@@ -0,0 +1,250 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  module Reconciliation
+    # Level 2: repairs charges from live Active Storage ownership.
+    #
+    # The unit of work is one (account, blob) pair in its own short transaction,
+    # which is also the unit of retry and of failure. A pair that raises does not
+    # undo pairs already repaired; partial progress is reported rather than
+    # hidden.
+    #
+    # Lock order is charge first, then account -- the same order normal
+    # accounting uses -- so the two can run concurrently without deadlocking.
+    class LedgerReconciler
+      BATCH_SIZE = 500
+
+      def initialize(result:, budget:, dry_run:)
+        @result = result
+        @budget = budget
+        @dry_run = dry_run
+        @audit_result = Audit::Result.new
+        @resolver = ActiveStorageQuota::OwnerResolver.new(result: @audit_result, record_findings: true)
+      end
+
+      def run(owner)
+        account = Account.find_or_create_for!(owner) unless @dry_run
+        account ||= owner.storage_quota_account
+
+        result.increment(:accounts_checked)
+
+        blob_ids = candidate_blob_ids(owner, account)
+        return if blob_ids.empty?
+
+        account ||= owner.storage_quota_account
+        if account.nil?
+          # Dry run against an owner with no account yet: report what applying
+          # would create rather than creating it.
+          result.record(:planned_change, action: :create_quota_account,
+                                         owner_type: owner.class.polymorphic_name,
+                                         owner_id: owner.id.to_s)
+          return
+        end
+
+        repair_pairs(account, blob_ids)
+        carry_report_only_findings
+      end
+
+      private
+        attr_reader :result, :budget
+
+        def repair_pairs(account, blob_ids)
+          # Preflight is read-only, so a dry run runs exactly the same check and
+          # predicts the same breaker trips.
+          removable = preflight(account, blob_ids)
+          deletions = 0
+          bytes_removed = 0
+          # Bug found in testing: measuring the removed fraction against a
+          # used_bytes that our own deletions keep shrinking makes the
+          # denominator collapse and trips the breaker spuriously. The budget is
+          # a fraction of what the account held when this run started.
+          baseline_used_bytes = account.reload.used_bytes
+
+          blob_ids.each_slice(BATCH_SIZE) do |slice|
+            expected = expected_counts(account, slice)
+            sizes = blob_sizes(slice)
+
+            slice.each do |blob_id|
+              result.increment(:pairs_checked)
+              count = expected.fetch(blob_id, 0)
+
+              within_budget = within_budget?(baseline_used_bytes, deletions, bytes_removed,
+                                             sizes[blob_id])
+
+              outcome = repair_pair(account, blob_id, sizes[blob_id], count, removable && within_budget)
+
+              case outcome
+              when :removed
+                deletions += 1
+                bytes_removed += sizes[blob_id].to_i
+              when :deletion_skipped
+                # The breaker trips only when a real removal was refused because
+                # of the budget -- never for a pair that needed no removal.
+                trip_apply_breaker!(account, deletions, bytes_removed) if removable && !within_budget
+              end
+            end
+          end
+        end
+
+        # Would removing this pair keep the run inside its budget? Asked for
+        # every pair, not only planned candidates, because the account is live:
+        # a concurrent detach can turn a healthy pair into a removal between the
+        # plan and the locked recount. Deliberately pure -- whether the breaker
+        # trips depends on what the pair turned out to need.
+        def within_budget?(baseline_used_bytes, deletions, bytes_removed, byte_size)
+          return true unless budget.destructive?
+
+          !budget.exceeded?(deletions: deletions + 1,
+                            bytes_removed: bytes_removed + byte_size.to_i,
+                            used_bytes: baseline_used_bytes)
+        end
+
+        def trip_apply_breaker!(account, deletions, bytes_removed)
+          result.trip_circuit_breaker!(stage: :apply, quota_account_id: account.id,
+                                       deletions: deletions, bytes_removed: bytes_removed)
+        end
+
+        # Before touching anything, price the whole account's destructive work.
+        # Over budget means we do not start: skipping every repair for that
+        # account, not only the deletions, because a tripped breaker means this
+        # account's picture is not trusted.
+        def preflight(account, blob_ids)
+          return false unless budget.destructive?
+
+          planned_deletions = 0
+          planned_bytes = 0
+
+          blob_ids.each_slice(BATCH_SIZE) do |slice|
+            expected = expected_counts(account, slice)
+            sizes = blob_sizes(slice)
+            charges = charge_sizes(account, slice)
+
+            slice.each do |blob_id|
+              next unless charges.key?(blob_id)
+              next unless expected.fetch(blob_id, 0).zero?
+
+              planned_deletions += 1
+              planned_bytes += charges[blob_id].to_i
+            end
+          end
+
+          return true unless budget.exceeded?(deletions: planned_deletions,
+                                              bytes_removed: planned_bytes,
+                                              used_bytes: account.used_bytes)
+
+          result.trip_circuit_breaker!(stage: :preflight, quota_account_id: account.id,
+                                       planned_deletions: planned_deletions,
+                                       planned_bytes_removed: planned_bytes)
+          false
+        end
+
+        def repair_pair(account, blob_id, byte_size, expected_count, allow_deletion)
+          if @dry_run
+            plan_pair(account, blob_id, byte_size, expected_count, allow_deletion)
+          else
+            apply_pair(account, blob_id, byte_size, allow_deletion)
+          end
+        end
+
+        def plan_pair(account, blob_id, byte_size, expected_count, allow_deletion)
+          outcome = Charge.plan_pair(quota_account_id: account.id, blob_id: blob_id,
+                                     allow_deletion: allow_deletion, expected_count: expected_count)
+          tally(outcome, account, blob_id, byte_size, planned: true)
+          outcome
+        end
+
+        def apply_pair(account, blob_id, byte_size, allow_deletion)
+          outcome = Charge.reconcile_pair!(
+            quota_account_id: account.id, blob_id: blob_id,
+            blob_byte_size: byte_size, allow_deletion: allow_deletion
+          ) { live_attachment_count(account, blob_id) }
+
+          tally(outcome, account, blob_id, byte_size, planned: false)
+          outcome
+        rescue StandardError => e
+          # Reconciliation is explicit maintenance, not a fail-open callback, so
+          # failures stay visible. One bad pair does not abandon the rest.
+          result.record(:errors, quota_account_id: account.id, blob_id: blob_id,
+                                 exception_class: e.class.name, exception_message: e.message)
+          :error
+        end
+
+        def tally(outcome, account, blob_id, byte_size, planned:)
+          case outcome
+          when :created
+            result.increment(:charges_created).increment(:bytes_added, by: byte_size.to_i)
+          when :removed
+            result.increment(:charges_removed).increment(:bytes_removed, by: byte_size.to_i)
+          when :count_corrected
+            result.increment(:attachment_counts_corrected)
+          when :deletion_skipped
+            result.record(:skipped, reason: :destructive_repair_not_permitted,
+                                    quota_account_id: account.id, blob_id: blob_id)
+            return
+          else
+            return
+          end
+
+          return unless planned
+
+          result.record(:planned_change, action: outcome, quota_account_id: account.id,
+                                         blob_id: blob_id, byte_size: byte_size)
+        end
+
+        # The authoritative count, recomputed inside the repair transaction.
+        def live_attachment_count(account, blob_id)
+          @resolver.pairs_for_blobs([ blob_id ]).fetch([ account.id, blob_id.to_s ], 0)
+        end
+
+        def expected_counts(account, blob_ids)
+          @resolver.pairs_for_blobs(blob_ids).each_with_object(Hash.new(0)) do |((acct, blob), count), map|
+            map[blob] = count if acct == account.id
+          end
+        end
+
+        def candidate_blob_ids(owner, account)
+          from_charges = account ? Charge.where(quota_account_id: account.id).pluck(:blob_id) : []
+
+          (from_charges.map(&:to_s) + owner_attachment_blob_ids(owner)).uniq.sort
+        end
+
+        # Blobs attached to records this owner pays for. Where the payer is a
+        # plain belongs_to this is an indexed lookup; otherwise each candidate
+        # record's declared source is resolved in Ruby.
+        def owner_attachment_blob_ids(owner)
+          @resolver.managed_record_types.flat_map do |record_type, klass|
+            record_ids = @resolver.record_ids_for_owner(klass, owner)
+            next [] if record_ids.empty?
+
+            ActiveStorage::Attachment
+              .where(record_type: record_type, record_id: record_ids)
+              .pluck(:blob_id).map(&:to_s)
+          end
+        end
+
+        def blob_sizes(blob_ids)
+          ActiveStorage::Blob.where(id: blob_ids).pluck(:id, :byte_size)
+                             .each_with_object({}) { |(id, size), map| map[id.to_s] = size }
+        end
+
+        def charge_sizes(account, blob_ids)
+          Charge.where(quota_account_id: account.id, blob_id: blob_ids)
+                .pluck(:blob_id, :byte_size)
+                .each_with_object({}) { |(blob_id, size), map| map[blob_id.to_s] = size }
+        end
+
+        # Findings the resolver noticed that reconciliation deliberately does
+        # not repair. They are surfaced, never acted on.
+        def carry_report_only_findings
+          %i[missing_quota_owner invalid_quota_owner unresolvable_record_type].each do |category|
+            total = @audit_result.count(category)
+            next if total.zero?
+
+            @audit_result.sample(category).each do |finding|
+              result.record(:report_only, category: category, **finding.details)
+            end
+          end
+        end
+    end
+  end
+end

data/lib/active_storage_quota/reconciliation/result.rb

@@ -0,0 +1,107 @@
+# frozen_string_literal: true
+
+require "set"
+
+module ActiveStorageQuota
+  module Reconciliation
+    # What a reconciliation run did, or would have done.
+    #
+    # Bounded like Audit::Result: exact totals, capped samples, scalars only.
+    # Partial progress is deliberately visible -- per-pair transactions mean a
+    # failure part-way through leaves earlier repairs in place, and pretending
+    # otherwise would be a lie about the guarantee.
+    class Result
+      DEFAULT_SAMPLE_LIMIT = 100
+
+      COUNTERS = %i[
+        accounts_checked pairs_checked
+        charges_created charges_removed attachment_counts_corrected counters_corrected
+        bytes_added bytes_removed
+      ].freeze
+
+      attr_reader :sample_limit, :counts, :samples
+
+      def initialize(dry_run:, sample_limit: DEFAULT_SAMPLE_LIMIT)
+        @dry_run = dry_run
+        @sample_limit = sample_limit
+        @totals = Hash.new(0)
+        @counts = Hash.new(0)
+        @samples = Hash.new { |hash, key| hash[key] = [] }
+        @circuit_breaker_triggered = false
+      end
+
+      def dry_run?
+        @dry_run
+      end
+
+      def circuit_breaker_triggered?
+        @circuit_breaker_triggered
+      end
+
+      COUNTERS.each do |name|
+        define_method(name) { @totals[name] }
+      end
+
+      def increment(counter, by: 1)
+        @totals[counter] += by
+        self
+      end
+
+      # A planned change (dry run) or a category needing attention: skipped,
+      # unsupported, errors, report-only findings.
+      def record(category, **details)
+        category = category.to_sym
+        @counts[category] += 1
+        @samples[category] << details if @samples[category].size < sample_limit
+
+        self
+      end
+
+      def trip_circuit_breaker!(**details)
+        @circuit_breaker_triggered = true
+        record(:circuit_breaker_triggered, **details)
+      end
+
+      def count(category)
+        @counts[category.to_sym]
+      end
+
+      def sample(category)
+        @samples[category.to_sym]
+      end
+
+      def planned_changes
+        @samples[:planned_change]
+      end
+
+      def changed?
+        (charges_created + charges_removed + attachment_counts_corrected + counters_corrected).positive?
+      end
+
+      def to_h
+        {
+          dry_run: dry_run?,
+          circuit_breaker_triggered: circuit_breaker_triggered?,
+          totals: COUNTERS.to_h { |name| [ name, @totals[name] ] },
+          categories: @counts.dup,
+          samples: @samples.transform_values(&:dup)
+        }
+      end
+
+      def to_s
+        headline = dry_run? ? "reconciliation plan" : "reconciliation"
+        lines = [ "#{headline}: #{summary_line}" ]
+        @counts.sort_by { |_, total| -total }.each { |category, total| lines << "  #{category}: #{total}" }
+        lines << "  CIRCUIT BREAKER TRIPPED" if circuit_breaker_triggered?
+        lines.join("\n")
+      end
+
+      private
+        def summary_line
+          "#{accounts_checked} account(s), #{pairs_checked} pair(s), " \
+            "+#{charges_created} charge(s), -#{charges_removed} charge(s), " \
+            "#{attachment_counts_corrected} count(s), #{counters_corrected} counter(s)"
+        end
+    end
+  end
+end

data/lib/active_storage_quota/reconciliation.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+require "active_storage_quota/reconciliation/budget"
+require "active_storage_quota/reconciliation/result"
+require "active_storage_quota/reconciliation/counter_reconciler"
+require "active_storage_quota/reconciliation/ledger_reconciler"
+
+module ActiveStorageQuota
+  # Repair operations built on what the audit reports.
+  #
+  # Two levels, kept separate because they have different costs and different
+  # risks:
+  #
+  #   counters  Account#used_bytes from charges. Cheap, safe, no Active Storage.
+  #   ledger    Charges from live Active Storage ownership. Expensive, and the
+  #             only level that can destroy anything.
+  #
+  # Audit findings are candidates, never commands: every repair re-reads the
+  # current state under lock before it mutates anything, because the world moves
+  # between auditing and repairing.
+  module Reconciliation
+    class << self
+      def counters!(owner:, dry_run: false, sample_limit: Result::DEFAULT_SAMPLE_LIMIT)
+        result = Result.new(dry_run: dry_run, sample_limit: sample_limit)
+        CounterReconciler.new(result: result, dry_run: dry_run).run(accounts_for(owner))
+        result
+      end
+
+      def ledger!(owner:, dry_run: false, sample_limit: Result::DEFAULT_SAMPLE_LIMIT, **limits)
+        result = Result.new(dry_run: dry_run, sample_limit: sample_limit)
+        run_ledger(owner, result, budget(limits), dry_run)
+        result
+      end
+
+      # Ledger first, then counters: counters are derived from charges, so
+      # repairing them before the ledger would just recompute a stale total.
+      def all!(owner:, dry_run: false, sample_limit: Result::DEFAULT_SAMPLE_LIMIT, **limits)
+        result = Result.new(dry_run: dry_run, sample_limit: sample_limit)
+
+        run_ledger(owner, result, budget(limits), dry_run)
+        CounterReconciler.new(result: result, dry_run: dry_run).run(accounts_for(owner))
+
+        result
+      end
+
+      private
+        def run_ledger(owner, result, budget, dry_run)
+          unless active_storage_available?
+            result.record(:unsupported, reason: "Active Storage is not loaded")
+            return
+          end
+
+          LedgerReconciler.new(result: result, budget: budget, dry_run: dry_run).run(owner)
+        end
+
+        def budget(limits)
+          Budget.new(
+            max_charge_deletions: limits[:max_charge_deletions],
+            max_removed_fraction: limits[:max_removed_fraction]
+          )
+        end
+
+        def accounts_for(owner)
+          Account.where(owner: owner).order(:id)
+        end
+
+        def active_storage_available?
+          defined?(ActiveStorage::Attachment) && defined?(ActiveStorage::Blob)
+        end
+    end
+  end
+end

data/lib/active_storage_quota/record.rb

@@ -0,0 +1,20 @@
+# frozen_string_literal: true
+
+require "active_record"
+
+module ActiveStorageQuota
+  # Abstract base for this gem's models.
+  #
+  # Deliberately inherits from ActiveRecord::Base rather than the host
+  # application's ApplicationRecord: a gem cannot assume that constant exists,
+  # and inheriting from it would drag application-wide concerns, default scopes
+  # and callbacks into quota accounting.
+  #
+  # Having a base of our own also gives applications a single place to point
+  # these tables at another database:
+  #
+  #   ActiveStorageQuota::Record.connects_to database: { writing: :quota }
+  class Record < ActiveRecord::Base
+    self.abstract_class = true
+  end
+end