active_storage_quota 0.1.0

data/lib/active_storage_quota/reservation.rb

@@ -0,0 +1,297 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # A claim on an owner's storage, held while an upload is in flight.
+  #
+  # A reservation is created +active+ and reaches exactly one terminal state:
+  #
+  #   active --> committed   the bytes became stored data
+  #   active --> released    the upload failed, expired, or was abandoned
+  #
+  # Terminal rows are kept. That makes +commit!+ and +release!+ idempotent
+  # under job retries, and leaves an account's whole history readable in one
+  # query when its counters need explaining.
+  #
+  # Every transition is a conditional UPDATE guarded on the current status.
+  # Winning that UPDATE -- getting one affected row -- is what grants the right
+  # to move the account's counters, and both happen in one transaction. Two
+  # callers racing on the same reservation therefore cannot both move its
+  # bytes: the loser's UPDATE matches nothing.
+  class Reservation < Record
+    self.table_name = "active_storage_quota_reservations"
+
+    STATUS_ACTIVE = "active"
+    STATUS_COMMITTED = "committed"
+    STATUS_RELEASED = "released"
+    STATUSES = [ STATUS_ACTIVE, STATUS_COMMITTED, STATUS_RELEASED ].freeze
+
+    belongs_to :quota_account,
+               class_name: "ActiveStorageQuota::Account",
+               inverse_of: :reservations
+
+    validates :byte_size, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+    validates :status, inclusion: { in: STATUSES }
+    validates :expires_at, presence: true
+
+    class << self
+      # Releases reservations whose expiry has passed, returning how many this
+      # invocation actually transitioned.
+      #
+      # Safe to run from several processes at once with no coordination. Two
+      # sweepers may select the same rows; only one can win each guarded
+      # transition, and only won transitions are counted.
+      #
+      #   ActiveStorageQuota.release_expired_reservations!(limit: 1_000)
+      def release_expired!(limit: 1_000)
+        unless limit.is_a?(Integer) && limit.positive?
+          raise ArgumentError, "limit must be a positive Integer, got #{limit.inspect}"
+        end
+
+        now = Time.current
+
+        expired(now).limit(limit).order(:id).to_a.count do |reservation|
+          reservation.release_if_expired!(now)
+        end
+      end
+
+      # Consumes the hold covering +charge+, if this account holds one.
+      #
+      # Matching is account-scoped on purpose. A reservation taken by one owner
+      # for a blob must never be consumed or released by another owner attaching
+      # that same blob: it stays active and is reclaimed by its own TTL.
+      def consume_matching!(charge:, now: Time.current)
+        candidate = where(
+          quota_account_id: charge.quota_account_id,
+          blob_id: charge.blob_id,
+          byte_size: charge.byte_size,
+          status: STATUS_ACTIVE
+        ).where("expires_at > ?", now).order(:id).first
+
+        return false if candidate.nil?
+
+        candidate.consume!(charge, now)
+      end
+
+      # Hands back a hold that turned out to be unnecessary, because the blob
+      # was already charged to this account. Tolerates losing the race.
+      def release_matching!(quota_account_id:, blob_id:, now: Time.current)
+        candidate = where(
+          quota_account_id: quota_account_id,
+          blob_id: blob_id.to_s,
+          status: STATUS_ACTIVE
+        ).where("expires_at > ?", now).order(:id).first
+
+        return false if candidate.nil?
+
+        candidate.release!
+      rescue InvalidReservation
+        false
+      end
+
+      # Releases every active hold for a blob, returning how many this call
+      # actually released.
+      #
+      # A destroyed blob can never be attached, so no hold naming it can ever be
+      # consumed; keeping the bytes reserved until the TTL expires would be pure
+      # waste. Direct uploads make this common, since an abandoned upload is
+      # usually purged as an unattached blob long before its hold lapses.
+      #
+      # Every hold for the blob is released regardless of account: they are all
+      # equally unusable. Each goes through the ordinary guarded release, so this
+      # is idempotent and races with the expiry sweeper the same way every other
+      # transition does -- whichever wins, the bytes come back exactly once.
+      def release_active_for_blob!(blob_id:, now: Time.current)
+        where(blob_id: blob_id.to_s, status: STATUS_ACTIVE).order(:id).to_a.count do |reservation|
+          reservation.release!
+        rescue InvalidReservation
+          # Committed or released between the read and the write. Not our bytes.
+          false
+        end
+      end
+
+      # Active reservations whose expiry has passed. The exact complement of
+      # the guard commit! uses, so at any instant an active reservation is
+      # claimable by exactly one of the two.
+      def expired(now = Time.current)
+        where(status: STATUS_ACTIVE).where("expires_at <= ?", now)
+      end
+    end
+
+    def active?
+      status == STATUS_ACTIVE
+    end
+
+    def committed?
+      status == STATUS_COMMITTED
+    end
+
+    def released?
+      status == STATUS_RELEASED
+    end
+
+    # Whether this reservation's hard expiry has passed. Compared against the
+    # application clock rather than the database's, so expiry can be tested
+    # with travel_to and behaves identically on every adapter.
+    def expired?(now = Time.current)
+      expires_at <= now
+    end
+
+    # Consumes this hold on behalf of +charge+, in the charge's transaction.
+    #
+    # Returns false when the guarded transition matched nothing -- the sweeper
+    # released it, it expired, or another charge consumed it first. That is a
+    # benign lost race, not an accounting failure: the caller simply claims
+    # fresh capacity instead. An exception here means something worse, that the
+    # transition WAS won but the counters could not be moved.
+    #
+    # Requires a persisted, matching charge, so a hold can never become usage
+    # without a ledger row to account for it. That is what makes
+    # used_bytes == SUM(charges.byte_size) structural.
+    def consume!(charge, now = Time.current)
+      validate_charge!(charge)
+
+      # Its own savepoint, so the status transition and the counter move are
+      # atomic even though this normally runs inside Charge's transaction. If
+      # the counters cannot be moved, the transition is undone rather than
+      # leaving a committed reservation whose bytes were never transferred.
+      claimed = self.class.transaction(requires_new: true) do
+        won = self.class
+          .where(id: id, status: STATUS_ACTIVE, quota_account_id: charge.quota_account_id,
+                 blob_id: charge.blob_id, byte_size: charge.byte_size)
+          .where("expires_at > ?", now)
+          .update_all(status: STATUS_COMMITTED, updated_at: now) == 1
+
+        move_reserved_to_used!(now) if won
+        won
+      end
+
+      return false unless claimed
+
+      reload
+      true
+    end
+
+    # Hands this reservation's bytes back to the account.
+    #
+    # Returns true when the reservation is released, including when it was
+    # already released. Raises InvalidReservation if it was committed, since
+    # those bytes are stored data now.
+    #
+    # Deliberately has no expiry guard: freeing capacity is always safe, so an
+    # expired reservation can still be released explicitly.
+    def release!
+      now = Time.current
+
+      claimed = self.class.transaction(requires_new: true) do
+        if claim_transition(to: STATUS_RELEASED, now: now)
+          return_reserved!(now)
+          true
+        else
+          false
+        end
+      end
+
+      return explain_failed_release! unless claimed
+
+      reload
+      true
+    end
+
+    # The sweeper's transition. Same as release!, but guarded on the expiry
+    # having passed so a reservation that is still valid is never swept.
+    #
+    # Returns true only if this call is the one that released it.
+    def release_if_expired!(now = Time.current)
+      claimed = self.class.transaction(requires_new: true) do
+        if claim_transition(to: STATUS_RELEASED, now: now) { |rows| rows.where("expires_at <= ?", now) }
+          return_reserved!(now)
+          true
+        else
+          false
+        end
+      end
+
+      reload if claimed
+      claimed
+    end
+
+    private
+      def validate_charge!(charge)
+        return if charge.is_a?(Charge) && charge.persisted? &&
+                  charge.quota_account_id == quota_account_id &&
+                  charge.byte_size == byte_size &&
+                  blob_id.present? && charge.blob_id == blob_id
+
+        raise ArgumentError,
+              "reservation #{id} can only be consumed by a persisted charge on " \
+              "the same account, for the same blob, for the same number of " \
+              "bytes. A reservation taken without a blob can only be released."
+      end
+
+      # The guarded state transition. One UPDATE, matching only while this row
+      # is still active, so exactly one concurrent caller can win it.
+      def claim_transition(to:, now:)
+        rows = self.class.where(id: id, status: STATUS_ACTIVE)
+        rows = yield(rows) if block_given?
+
+        rows.update_all(status: to, updated_at: now) == 1
+      end
+
+      def move_reserved_to_used!(now)
+        apply_counters!(
+          now,
+          [ "reserved_bytes = reserved_bytes - ?, used_bytes = used_bytes + ?, updated_at = ?",
+            byte_size, byte_size, now ]
+        )
+      end
+
+      def return_reserved!(now)
+        apply_counters!(
+          now,
+          [ "reserved_bytes = reserved_bytes - ?, updated_at = ?", byte_size, now ]
+        )
+      end
+
+      # Moves the account's counters, refusing to act unless the account still
+      # holds the bytes this reservation claims.
+      #
+      # The reserved_bytes >= byte_size guard is what keeps a drifted counter
+      # from becoming a negative one. Failing here rolls back the status
+      # transition made moments ago in this same transaction, so the
+      # reservation stays active and the account is untouched.
+      def apply_counters!(_now, updates)
+        affected = Account
+          .where(id: quota_account_id)
+          .where("reserved_bytes >= ?", byte_size)
+          .update_all(updates)
+
+        return if affected == 1
+
+        raise CounterInvariantViolation,
+              "reservation #{id} claims #{byte_size} reserved bytes, but account " \
+              "#{quota_account_id} does not hold that many; refusing to move " \
+              "counters. The account's counters have drifted from its reservations."
+      end
+
+      # Only reached when the guarded UPDATE matched nothing, so it costs one
+      # SELECT on the failure path and none on the success path.
+      def explain_failed_release!
+        current = self.class.find_by(id: id)
+
+        raise InvalidReservation, "reservation #{id} no longer exists" if current.nil?
+
+        case current.status
+        when STATUS_RELEASED
+          reload
+          true
+        when STATUS_COMMITTED
+          raise InvalidReservation,
+                "reservation #{id} was committed and cannot be released; " \
+                "those bytes are stored data"
+        else
+          raise InvalidReservation,
+                "reservation #{id} changed state concurrently and could not be released"
+        end
+      end
+  end
+end

data/lib/active_storage_quota/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  VERSION = "0.1.0"
+end

data/lib/active_storage_quota.rb

@@ -0,0 +1,207 @@
+# frozen_string_literal: true
+
+require "active_support"
+
+require "active_storage_quota/version"
+require "active_storage_quota/errors"
+require "active_storage_quota/configuration"
+require "active_storage_quota/definition"
+require "active_storage_quota/owner_source"
+require "active_storage_quota/owner_resolver"
+
+# Storage quota management for Rails Active Storage.
+#
+# See README.md for usage. The short version:
+#
+#   class Company < ApplicationRecord
+#     has_storage_quota limit: 5.gigabytes
+#   end
+#
+#   company.storage_remaining # => 5368709120
+module ActiveStorageQuota
+  # Quota accounts are keyed by owner *and* scope. Owners declared with
+  # +has_storage_quota+ use this scope; named scopes arrive in a later version.
+  DEFAULT_SCOPE = "default"
+
+  # The Active-Record-dependent public API, declared with autoload rather than
+  # required from the :active_record load hook.
+  #
+  # That hook only fires when ActiveRecord::Base is first referenced. A rake
+  # task has not touched a model by the time it runs, so requiring these there
+  # left every one of them undefined -- which broke every task this gem ships.
+  # Autoload resolves them on first reference instead, whenever that happens.
+  autoload :Record, "active_storage_quota/record"
+  autoload :Account, "active_storage_quota/account"
+  autoload :Reservation, "active_storage_quota/reservation"
+  autoload :Charge, "active_storage_quota/charge"
+  autoload :Audit, "active_storage_quota/audit"
+  autoload :Reconciliation, "active_storage_quota/reconciliation"
+  autoload :Backfill, "active_storage_quota/backfill"
+
+  class << self
+    # The current Configuration.
+    def config
+      @config ||= Configuration.new
+    end
+
+    # Yields the Configuration for an initializer block.
+    def configure
+      yield config
+      config
+    end
+
+    # Restores every setting to its default. Intended for test suites.
+    def reset_configuration!
+      @config = Configuration.new
+    end
+
+    # Where accounting failures are reported. Defaults to Rails' logger when
+    # Rails is present, otherwise Active Record's.
+    attr_writer :logger
+
+    def logger
+      return @logger if defined?(@logger) && @logger
+
+      if defined?(Rails) && Rails.respond_to?(:logger) && Rails.logger
+        Rails.logger
+      elsif defined?(ActiveRecord::Base)
+        ActiveRecord::Base.logger
+      end
+    end
+
+    # Reports an accounting failure that was deliberately not raised.
+    #
+    # Deletion accounting fails open: it can never undo a deletion, because
+    # blocking a delete on drifted counters would trap an over-quota account
+    # with no way back under its limit. Failing open without saying so would be
+    # worse, so every swallowed failure is logged and instrumented.
+    #
+    # Subscribe with:
+    #
+    #   ActiveSupport::Notifications.subscribe("active_storage_quota.accounting_failed") do |*, payload|
+    #     Sentry.capture_message(payload[:exception_message])
+    #   end
+    #
+    # The payload carries scalars only, never Active Record objects.
+    def accounting_failed(operation:, exception:, quota_account_id: nil, blob_id: nil,
+                          record_type: nil, record_id: nil)
+      payload = {
+        operation: operation,
+        quota_account_id: quota_account_id,
+        blob_id: blob_id,
+        record_type: record_type,
+        record_id: record_id&.to_s,
+        exception_class: exception.class.name,
+        exception_message: exception.message
+      }
+
+      logger&.error do
+        "[active_storage_quota] #{operation} accounting failed for account " \
+          "#{quota_account_id.inspect} blob #{blob_id.inspect}: " \
+          "#{exception.class}: #{exception.message}"
+      end
+
+      ActiveSupport::Notifications.instrument("active_storage_quota.accounting_failed", payload)
+
+      nil
+    end
+
+    # Inspects quota accounting and reports what it finds, writing nothing.
+    #
+    #   result = ActiveStorageQuota.audit
+    #   puts result unless result.clean?
+    #
+    # See ActiveStorageQuota::Audit. Safe to run against production: no repairs,
+    # no deletions, no account creation, and no dependence on whether quota
+    # enforcement is switched on.
+    def audit(sample_limit: Audit::Result::DEFAULT_SAMPLE_LIMIT)
+      Audit.run(sample_limit: sample_limit)
+    end
+
+    # Repairs Account#used_bytes from the charges it aggregates.
+    #
+    #   ActiveStorageQuota.reconcile_counters!(owner: company)
+    #
+    # Cheap and safe to run often. Locks each account row briefly so concurrent
+    # uploads cannot lose bytes to it.
+    def reconcile_counters!(owner:, dry_run: false, **options)
+      Reconciliation.counters!(owner: owner, dry_run: dry_run, **options)
+    end
+
+    # Repairs charges from live Active Storage ownership.
+    #
+    #   ActiveStorageQuota.reconcile_ledger!(owner: company, dry_run: true)
+    #
+    # Destructive repair -- removing charges nothing backs any more -- happens
+    # only when a budget is supplied; see Reconciliation::Budget.
+    def reconcile_ledger!(owner:, dry_run: false, **options)
+      Reconciliation.ledger!(owner: owner, dry_run: dry_run, **options)
+    end
+
+    # Ledger repair followed by counter repair.
+    #
+    #   ActiveStorageQuota.reconcile!(owner: company, dry_run: true)
+    #   ActiveStorageQuota.reconcile!(owner: company,
+    #                                 max_charge_deletions: 10,
+    #                                 max_removed_fraction: 0.25)
+    #
+    # +owner+ is required. A reconciliation across every tenant is reachable
+    # only through the rake task, which makes you ask for it explicitly -- it is
+    # the operation most worth thinking about before running.
+    def reconcile!(owner:, dry_run: false, **options)
+      Reconciliation.all!(owner: owner, dry_run: dry_run, **options)
+    end
+
+    # Releases reservations whose expiry has passed, returning how many this
+    # call actually released.
+    #
+    #   ActiveStorageQuota.release_expired_reservations!
+    #   ActiveStorageQuota.release_expired_reservations!(limit: 5_000)
+    #
+    # A reservation is abandoned whenever the process holding it dies before
+    # committing or releasing. Without this, those bytes stay reserved forever.
+    #
+    # There is no scheduler here on purpose: run it from cron, a recurring job,
+    # or the bundled rake task, as often as your reservation_ttl warrants. It is
+    # safe to run several at once -- no locks, no leader election -- because
+    # each reservation can only be released by whichever call wins its guarded
+    # transition.
+    #
+    # +limit+ bounds one invocation so a backlog cannot turn into an unbounded
+    # run. It must be a positive Integer.
+    def release_expired_reservations!(limit: 1_000)
+      Reservation.release_expired!(limit: limit)
+    end
+  end
+end
+
+# Active Record may be loaded before or after this gem. Deferring to the load
+# hook means the models are defined once Active Record is ready either way, and
+# the gem works with or without Rails booted.
+# The macros have to be installed on ActiveRecord::Base the moment it loads;
+# everything else resolves through the autoloads above, on first reference.
+ActiveSupport.on_load(:active_record) do
+  require "active_storage_quota/owner"
+  require "active_storage_quota/attachable"
+
+  extend ActiveStorageQuota::Owner::Macro
+  extend ActiveStorageQuota::Attachable::Macro
+end
+
+# Active Storage integration. These hooks are the documented extension point for
+# both models and exist in every supported Rails version, so nothing here
+# patches Active Storage. When Active Storage is never loaded, they simply never
+# run -- which is why it is a development dependency and not a runtime one.
+ActiveSupport.on_load(:active_storage_attachment) do
+  require "active_storage_quota/attachment_accounting"
+
+  include ActiveStorageQuota::AttachmentAccounting
+end
+
+ActiveSupport.on_load(:active_storage_blob) do
+  require "active_storage_quota/blob_accounting"
+
+  include ActiveStorageQuota::BlobAccounting
+end
+
+require "active_storage_quota/engine" if defined?(Rails::Engine)

data/lib/tasks/active_storage_quota.rake

@@ -0,0 +1,11 @@
+# frozen_string_literal: true
+
+namespace :active_storage_quota do
+  desc "Release quota reservations whose expiry has passed"
+  task release_expired_reservations: :environment do
+    limit = Integer(ENV.fetch("LIMIT", 1_000))
+    released = ActiveStorageQuota.release_expired_reservations!(limit: limit)
+
+    puts "Released #{released} expired reservation#{'s' unless released == 1}."
+  end
+end

data/lib/tasks/active_storage_quota_audit.rake

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+namespace :active_storage_quota do
+  desc "Report quota accounting inconsistencies (read only, changes nothing)"
+  task audit: :environment do
+    result = ActiveStorageQuota.audit(
+      sample_limit: Integer(ENV.fetch("SAMPLE_LIMIT", 100))
+    )
+
+    puts result
+    result.findings.each { |finding| puts "  - #{finding}" } if ENV["VERBOSE"]
+
+    exit(1) unless result.clean?
+  end
+end

data/lib/tasks/active_storage_quota_backfill.rake

@@ -0,0 +1,104 @@
+# frozen_string_literal: true
+
+require "shellwords"
+
+namespace :active_storage_quota do
+  desc "Show what historical backfill would create (dry run, never writes)"
+  task backfill: :environment do
+    ActiveStorageQuota::BackfillRakeSupport.run(dry_run: true)
+  end
+
+  namespace :backfill do
+    desc "Create historical quota charges (writes)"
+    task apply: :environment do
+      ActiveStorageQuota::BackfillRakeSupport.run(dry_run: false)
+    end
+  end
+end
+
+module ActiveStorageQuota
+  module BackfillRakeSupport # :nodoc:
+    class << self
+      def run(dry_run:)
+        record_type = ENV["RECORD_TYPE"].presence
+        require_explicit_scope!(record_type) unless dry_run
+
+        result = ActiveStorageQuota::Backfill.run(
+          record_type: record_type,
+          cursor: ENV["CURSOR"].presence,
+          batch_size: integer_env("BATCH_SIZE") || ActiveStorageQuota::Backfill::DEFAULT_BATCH_SIZE,
+          max_batches: integer_env("MAX_BATCHES"),
+          dry_run: dry_run,
+          allow_active_enforcement: ENV["ALLOW_ACTIVE_ENFORCEMENT"] == "1"
+        )
+
+        report(result, record_type, dry_run)
+      rescue ActiveStorageQuota::EnforcementActive => e
+        abort <<~MESSAGE
+          #{e.message}
+
+          Backfill writes charges but not used_bytes, so counters stay incomplete
+          until reconcile_counters! runs. Enforcing against them meanwhile would
+          admit uploads that should have been refused.
+
+          Either let limits resolve to nil during adoption, or re-run with
+          ALLOW_ACTIVE_ENFORCEMENT=1 if you accept that.
+        MESSAGE
+      end
+
+      private
+        def require_explicit_scope!(record_type)
+          return if record_type
+          return if ENV["ALL"] == "1"
+
+          abort "Refusing to back-fill every record type. Pass RECORD_TYPE=..., or ALL=1."
+        end
+
+        def report(result, record_type, dry_run)
+          puts result
+          puts "  scope: #{record_type || 'all record types'}"
+
+          if result.next_cursor
+            puts
+            puts "  Not finished. Resume with:"
+            puts "    #{resume_command(result.next_cursor, record_type, dry_run)}"
+          end
+
+          return unless result.counters_require_reconciliation?
+
+          puts
+          puts "  Charges were created but used_bytes was not touched."
+          puts "  Run ActiveStorageQuota.reconcile_counters!(owner: ...) before enabling limits."
+        end
+
+        # Resumes exactly this run. The task is the same one -- a dry run must
+        # never hint at apply -- and so is the scope, which apply requires and
+        # which the cursor is bound to. Batching and an explicit enforcement
+        # override are carried forward only when the operator gave them.
+        def resume_command(cursor, record_type, dry_run)
+          env = [ "CURSOR=#{Shellwords.escape(cursor)}" ]
+          env << (record_type ? "RECORD_TYPE=#{Shellwords.escape(record_type)}" : "ALL=1")
+          %w[BATCH_SIZE MAX_BATCHES].each do |name|
+            env << "#{name}=#{Shellwords.escape(ENV[name])}" if ENV[name].present?
+          end
+          env << "ALLOW_ACTIVE_ENFORCEMENT=1" if ENV["ALLOW_ACTIVE_ENFORCEMENT"] == "1"
+
+          task = dry_run ? "active_storage_quota:backfill" : "active_storage_quota:backfill:apply"
+          "#{env.join(" ")} bin/rails #{task}"
+        end
+
+        # A malformed value must never become zero or unlimited.
+        def integer_env(name)
+          raw = ENV[name].presence
+          return nil if raw.nil?
+
+          value = Integer(raw)
+          abort "#{name} must be a positive Integer, got #{raw.inspect}" unless value.positive?
+
+          value
+        rescue ArgumentError, TypeError
+          abort "#{name} must be a positive Integer, got #{raw.inspect}"
+        end
+    end
+  end
+end