active_storage_quota 0.1.0

data/lib/active_storage_quota/blob_accounting.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+require "active_support/concern"
+
+module ActiveStorageQuota
+  # Mixed into ActiveStorage::Blob through the :active_storage_blob load hook.
+  #
+  # This is how purging is observed. ActiveStorage::Attachment#purge deletes the
+  # attachment row with `delete`, firing no callbacks, and then destroys the
+  # blob -- and that destruction does run callbacks. So the blob is where a
+  # purge becomes visible to accounting.
+  #
+  # Idempotent against AttachmentAccounting: whichever path deletes the charge
+  # row is the one that lowers used_bytes, and the other finds nothing left.
+  #
+  # A blob shared by several attachments is not destroyed at all (Blob's own
+  # before_destroy raises InvalidForeignKey, which purge swallows), so purging
+  # one attachment of a shared blob is invisible here too. That remains a
+  # reconciliation case.
+  module BlobAccounting
+    extend ActiveSupport::Concern
+
+    included do
+      after_destroy_commit :discard_storage_quota_charges
+    end
+
+    private
+      # Two independent cleanups. They are isolated from each other on purpose:
+      # a failure discarding charges must not stop reservations being released,
+      # and neither may ever roll back the deletion of a user's file.
+      def discard_storage_quota_charges
+        cleanup("blob_destroy") { ActiveStorageQuota::Charge.discard_blob!(blob_id: id) }
+        cleanup("blob_destroy_reservations") do
+          ActiveStorageQuota::Reservation.release_active_for_blob!(blob_id: id)
+        end
+      end
+
+      def cleanup(operation)
+        yield
+      rescue StandardError => e
+        ActiveStorageQuota.accounting_failed(operation: operation, exception: e, blob_id: id.to_s)
+      end
+  end
+end

data/lib/active_storage_quota/charge.rb

@@ -0,0 +1,379 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # The accounting ledger: one row per (account, blob) meaning "this account is
+  # charged this blob's size, once".
+  #
+  # Three layers hold storage state, and they are not interchangeable:
+  #
+  #   ActiveStorage::Attachment  ground truth for ownership. Only a full rebuild
+  #                              reading live attachments can restore it.
+  #   Charge                     canonical ledger during normal operation.
+  #   Account#used_bytes         denormalized aggregate of charges.
+  #
+  # The write-path invariant is:
+  #
+  #   account.used_bytes == account.charges.sum(:byte_size)
+  #
+  # It is structural, not hoped for: +record!+ is the only operation that adds
+  # to used_bytes and +discard!+ the only one that subtracts, and each does so
+  # in the same transaction that inserts or deletes the charge row. The charge
+  # row is the claim -- double counting would need a duplicate row, which the
+  # unique index forbids.
+  #
+  # Charges *can* drift from attachments, because Active Storage's detach and
+  # shared-blob purge paths fire no callbacks and deletion accounting fails
+  # open. That is what reconciliation is for, and why it has two levels:
+  # charges repair used_bytes, attachments rebuild charges.
+  class Charge < Record
+    self.table_name = "active_storage_quota_charges"
+
+    # A charge can be created and destroyed underneath us by concurrent attach
+    # and detach activity. Each retry is correct; this bounds how long we keep
+    # losing before giving up.
+    MAX_ATTEMPTS = 3
+
+    belongs_to :quota_account,
+               class_name: "ActiveStorageQuota::Account",
+               inverse_of: :charges
+
+    validates :blob_id, presence: true
+    validates :byte_size, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+    validates :attachment_count, numericality: { only_integer: true, greater_than_or_equal_to: 0 }
+
+    class << self
+      # Records one attachment of +blob_id+ against +account+.
+      #
+      # The first attachment inserts a charge and raises used_bytes by the
+      # blob's size, consuming a matching reservation if one is held. Later
+      # attachments of the same blob to the same account only raise the
+      # reference count.
+      #
+      # Raises QuotaExceeded when the bytes do not fit, and
+      # ConcurrentAccountingConflict when contention outlasts MAX_ATTEMPTS.
+      def record!(account:, blob_id:, byte_size:, limit_bytes:)
+        blob_id = blob_id.to_s
+
+        MAX_ATTEMPTS.times do
+          outcome = attempt_record(account, blob_id, byte_size, limit_bytes)
+
+          return outcome unless outcome == :contended
+        end
+
+        raise ConcurrentAccountingConflict,
+              "could not record a charge for account #{account.id} and blob " \
+              "#{blob_id} after #{MAX_ATTEMPTS} attempts; the charge is being " \
+              "created and destroyed concurrently"
+      end
+
+      # Records the removal of one attachment.
+      #
+      # Returns true only when this call removed the last attachment and so
+      # lowered used_bytes.
+      def discard!(quota_account_id:, blob_id:)
+        blob_id = blob_id.to_s
+
+        transaction(requires_new: true) do
+          charge = find_by(quota_account_id: quota_account_id, blob_id: blob_id)
+
+          if charge.nil?
+            false
+          else
+            discard_one(quota_account_id, blob_id, charge.byte_size)
+          end
+        end
+      end
+
+      # Re-points one of +account+'s attachments from one blob to another, as
+      # ActiveStorage::Attachment#update!(blob: other) does. Internal: called
+      # from the attachment's after_update, inside its transaction.
+      #
+      # One transaction, so a refusal leaves both blobs exactly as they were.
+      # The old reference is dropped before the new one is claimed, so the bytes
+      # an in-place replacement frees can make room for it.
+      #
+      # Every other accounting path locks the charge row and then the account,
+      # so this takes both charge rows first -- in blob_id order, so two moves in
+      # opposite directions queue rather than deadlock -- and only then touches
+      # the account. A charge that does not exist yet cannot be locked, so the
+      # new blob's row is inserted provisionally, carrying no usage, to hold it.
+      def move!(account:, from_blob_id:, to_blob_id:, byte_size:, limit_bytes:) # :nodoc:
+        from = from_blob_id.to_s
+        to = to_blob_id.to_s
+
+        transaction(requires_new: true) do
+          old_charge = nil
+          new_charge = nil
+          provisional = false
+
+          [ from, to ].sort.each do |blob_id|
+            if blob_id == from
+              old_charge = lock_pair(account.id, from)
+            else
+              new_charge, provisional = claim_pair(account.id, to, byte_size)
+            end
+          end
+
+          discard_one(account.id, from, old_charge.byte_size) if old_charge
+
+          if provisional
+            where(id: new_charge.id).update_all([ "attachment_count = 1, updated_at = ?", Time.current ])
+            claim_usage!(account, new_charge, limit_bytes)
+          else
+            where(id: new_charge.id)
+              .update_all([ "attachment_count = attachment_count + 1, updated_at = ?", Time.current ])
+            # Already charged, so a hold taken for this blob bought nothing.
+            Reservation.release_matching!(quota_account_id: account.id, blob_id: to)
+          end
+        end
+      end
+
+      # Removes every charge for a blob that no longer exists, across all
+      # accounts. Each account is settled in its own transaction so one failure
+      # cannot strand the others.
+      #
+      # Idempotent against +discard!+: whichever path deletes the row is the one
+      # that lowers used_bytes.
+      def discard_blob!(blob_id:)
+        rows = where(blob_id: blob_id.to_s).pluck(:id, :quota_account_id, :byte_size)
+
+        rows.count do |id, account_id, bytes|
+          transaction(requires_new: true) do
+            if where(id: id).delete_all == 1
+              Account.release_used_bytes!(account_id, bytes)
+              true
+            else
+              false
+            end
+          end
+        end
+      end
+
+      # Repairs one (account, blob) pair to match live Active Storage ownership.
+      #
+      # Deliberately NOT built on record!. Normal accounting is delta-based --
+      # one attach means attachment_count += 1 -- while repair is absolute:
+      # attachment_count becomes the live count, whatever it was before. Calling
+      # record! here would add one on every run and make reconciliation
+      # non-idempotent while amplifying the drift it is meant to remove.
+      #
+      # The block recounts live attachments and is called only *after* this
+      # transaction owns the charge row, which is what closes the race against a
+      # concurrent detach (see the provisional insert below).
+      #
+      # Returns :created, :removed, :count_corrected, :deletion_skipped or :noop.
+      def reconcile_pair!(quota_account_id:, blob_id:, blob_byte_size:, allow_deletion:, &recount)
+        blob_id = blob_id.to_s
+
+        transaction(requires_new: true) do
+          charge = lock_pair(quota_account_id, blob_id)
+
+          if charge
+            repair_existing(charge, allow_deletion, &recount)
+          elsif blob_byte_size.nil?
+            # No charge and no blob to size one from: nothing to create.
+            :noop
+          else
+            create_if_still_needed(quota_account_id, blob_id, blob_byte_size, allow_deletion, &recount)
+          end
+        end
+      end
+
+      # What reconcile_pair! would do, without locks or writes. Shares the
+      # decision table with the real thing so a plan cannot drift from the
+      # repair it predicts.
+      def plan_pair(quota_account_id:, blob_id:, allow_deletion:, expected_count:)
+        charge = find_by(quota_account_id: quota_account_id, blob_id: blob_id.to_s)
+
+        decide(
+          charge_present: !charge.nil?,
+          charge_count: charge&.attachment_count,
+          expected_count: expected_count,
+          allow_deletion: allow_deletion
+        )
+      end
+
+      # The single decision table, used by both apply and dry run.
+      def decide(charge_present:, charge_count:, expected_count:, allow_deletion:)
+        if !charge_present
+          expected_count.positive? ? :created : :noop
+        elsif expected_count.zero?
+          allow_deletion ? :removed : :deletion_skipped
+        elsif charge_count != expected_count
+          :count_corrected
+        else
+          :noop
+        end
+      end
+
+      private
+        def lock_pair(quota_account_id, blob_id)
+          lock.find_by(quota_account_id: quota_account_id, blob_id: blob_id)
+        end
+
+        def repair_existing(charge, allow_deletion, &recount)
+          expected = recount.call
+
+          case decide(charge_present: true, charge_count: charge.attachment_count,
+                      expected_count: expected, allow_deletion: allow_deletion)
+          when :removed
+            # We hold the row lock, so nothing else can have taken it.
+            unless where(id: charge.id).delete_all == 1
+              raise CounterInvariantViolation,
+                    "charge #{charge.id} vanished while this transaction held its lock"
+            end
+
+            Account.release_used_bytes!(charge.quota_account_id, charge.byte_size)
+            :removed
+          when :count_corrected
+            where(id: charge.id)
+              .update_all([ "attachment_count = ?, updated_at = ?", expected, Time.current ])
+            :count_corrected
+          when :deletion_skipped
+            :deletion_skipped
+          else
+            :noop
+          end
+        end
+
+        # A charge that does not exist cannot be locked, so we create a
+        # provisional row carrying no usage, and only then read the live
+        # attachment state. Because the row is ours until commit, a concurrent
+        # detach either lands before the recount (and we drop the provisional)
+        # or blocks behind us (and adjusts the finished charge afterwards).
+        # Reading attachments *before* owning a row would let a deletion slip
+        # between the read and the insert, leaving a charge nothing backs.
+        def create_if_still_needed(quota_account_id, blob_id, blob_byte_size, allow_deletion, &recount)
+          charge = insert_provisional(quota_account_id, blob_id, blob_byte_size)
+
+          if charge.nil?
+            # Live accounting won the insert. Take the winner under lock and
+            # treat it as an existing pair; never increment for our own sake.
+            winner = lock_pair(quota_account_id, blob_id)
+            return :noop if winner.nil?
+
+            return repair_existing(winner, allow_deletion, &recount)
+          end
+
+          expected = recount.call
+
+          if expected.zero?
+            # Nothing to charge after all. The provisional never carried usage,
+            # so removing it is not a destructive repair.
+            where(id: charge.id).delete_all
+            :noop
+          else
+            where(id: charge.id)
+              .update_all([ "attachment_count = ?, updated_at = ?", expected, Time.current ])
+            claim_usage_for_repair!(charge)
+            :created
+          end
+        end
+
+        # The charge row for a move's new blob, held by this transaction: the
+        # existing one under lock, or a provisional one we inserted. Returns
+        # [charge, provisional?]. Bounded like record!, for the same reason.
+        def claim_pair(quota_account_id, blob_id, byte_size)
+          MAX_ATTEMPTS.times do
+            inserted = insert_provisional(quota_account_id, blob_id, byte_size)
+            return [ inserted, true ] if inserted
+
+            existing = lock_pair(quota_account_id, blob_id)
+            return [ existing, false ] if existing
+          end
+
+          raise ConcurrentAccountingConflict,
+                "could not move an attachment onto blob #{blob_id} for account " \
+                "#{quota_account_id} after #{MAX_ATTEMPTS} attempts; its charge is " \
+                "being created and destroyed concurrently"
+        end
+
+        def insert_provisional(quota_account_id, blob_id, blob_byte_size)
+          transaction(requires_new: true) do
+            create!(quota_account_id: quota_account_id, blob_id: blob_id,
+                    byte_size: blob_byte_size, attachment_count: 0)
+          end
+        rescue ActiveRecord::RecordNotUnique
+          nil
+        end
+
+        # Repair is truth restoration, not admission control, so capacity is
+        # claimed with a nil limit: an owner whose files moved to them is
+        # charged even when that leaves them over quota. Future uploads are
+        # blocked normally by the usual guarded claim.
+        def claim_usage_for_repair!(charge)
+          return if Reservation.consume_matching!(charge: charge)
+
+          Account.find(charge.quota_account_id)
+                 .claim_used_capacity!(charge.byte_size, nil)
+        end
+
+        def attempt_record(account, blob_id, byte_size, limit_bytes)
+          transaction(requires_new: true) do
+            charge = insert_charge(account, blob_id, byte_size)
+
+            if charge
+              claim_usage!(account, charge, limit_bytes)
+              charge
+            else
+              add_attachment(account, blob_id)
+            end
+          end
+        end
+
+        # The INSERT is the claim. Only the transaction that wins it may raise
+        # used_bytes. A savepoint because on PostgreSQL a constraint violation
+        # poisons the enclosing transaction, and this one is nested.
+        def insert_charge(account, blob_id, byte_size)
+          transaction(requires_new: true) do
+            create!(
+              quota_account_id: account.id,
+              blob_id: blob_id,
+              byte_size: byte_size,
+              attachment_count: 1
+            )
+          end
+        rescue ActiveRecord::RecordNotUnique
+          nil
+        end
+
+        def add_attachment(account, blob_id)
+          affected = where(quota_account_id: account.id, blob_id: blob_id)
+                     .update_all([ "attachment_count = attachment_count + 1, updated_at = ?", Time.current ])
+
+          # A concurrent final detach deleted the row between our INSERT losing
+          # and this UPDATE running. Start over.
+          return :contended if affected.zero?
+
+          # These bytes were already charged, so any hold we took bought nothing.
+          Reservation.release_matching!(quota_account_id: account.id, blob_id: blob_id)
+
+          find_by(quota_account_id: account.id, blob_id: blob_id)
+        end
+
+        # used_bytes must rise by exactly the blob's size, once. A held
+        # reservation supplies that capacity; otherwise it is claimed fresh.
+        def claim_usage!(account, charge, limit_bytes)
+          return if Reservation.consume_matching!(charge: charge)
+
+          account.claim_used_capacity!(charge.byte_size, limit_bytes)
+        end
+
+        # Both statements share a transaction deliberately: the UPDATE holds the
+        # row's write lock through the DELETE, so a concurrent attach cannot
+        # change the count in between. Splitting them loses attachments.
+        def discard_one(quota_account_id, blob_id, byte_size)
+          now = Time.current
+          scope = where(quota_account_id: quota_account_id, blob_id: blob_id)
+
+          scope.where("attachment_count > 0")
+               .update_all([ "attachment_count = attachment_count - 1, updated_at = ?", now ])
+
+          return false unless scope.where(attachment_count: 0).delete_all == 1
+
+          Account.release_used_bytes!(quota_account_id, byte_size, now)
+          true
+        end
+    end
+  end
+end

data/lib/active_storage_quota/configuration.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # Gem-wide settings. Access through ActiveStorageQuota.config or set them in
+  # an initializer:
+  #
+  #   ActiveStorageQuota.configure do |config|
+  #     config.reservation_ttl = 30.minutes
+  #   end
+  class Configuration
+    # One hour, in seconds. Long enough for a slow client to finish a large
+    # direct upload, short enough that a crashed process does not hold quota
+    # for the rest of the day.
+    DEFAULT_RESERVATION_TTL = 3_600
+
+    # How long a reservation stays committable, as an Integer number of seconds.
+    #
+    # Assigning an ActiveSupport::Duration is supported and normalised, so
+    # +config.reservation_ttl = 1.hour+ reads back as +3600+. Every byte and
+    # duration this gem exposes is an Integer.
+    attr_reader :reservation_ttl
+
+    def initialize
+      @reservation_ttl = DEFAULT_RESERVATION_TTL
+    end
+
+    def reservation_ttl=(value)
+      unless value.is_a?(Numeric) && value.to_i.positive?
+        raise ConfigurationError,
+              "reservation_ttl must be a positive number of seconds " \
+              "(got #{value.inspect})"
+      end
+
+      @reservation_ttl = value.to_i
+    end
+  end
+end

data/lib/active_storage_quota/definition.rb

@@ -0,0 +1,100 @@
+# frozen_string_literal: true
+
+module ActiveStorageQuota
+  # Resolves the storage limit declared by +has_storage_quota+.
+  #
+  # A declaration is one of:
+  #
+  #   limit: 5.gigabytes            # Integer number of bytes
+  #   limit: :storage_limit_for_plan # Symbol, a method on the owner
+  #   limit: ->(owner) { ... }      # callable taking the owner
+  #   limit: -> { ... }             # callable evaluated against the owner
+  #   limit: nil                    # unlimited
+  #
+  # Whatever the declaration, the resolved value must be an Integer of zero or
+  # more bytes, or nil for unlimited. Nothing is coerced: a limit of "5000" or
+  # 5.0 is a bug in the host application, not a number to round off, so it
+  # raises ConfigurationError rather than silently calling +to_i+.
+  #
+  # Static declarations are validated when the class body is evaluated. Symbol
+  # and callable declarations can only be validated once resolved, so they are
+  # checked on every call.
+  class Definition
+    # Arities a callable limit may declare: zero (evaluated against the owner)
+    # or one (passed the owner).
+    SUPPORTED_ARITIES = [ 0, 1 ].freeze
+
+    attr_reader :limit
+
+    def initialize(limit:)
+      @limit = limit
+      validate_declaration!
+    end
+
+    # The limit for +owner+, in bytes. Returns nil when unlimited.
+    #
+    # Raises ConfigurationError if a Symbol or callable declaration resolves to
+    # something that is not an Integer >= 0 or nil.
+    def resolve(owner)
+      validate_resolved(call_limit(owner), owner)
+    end
+
+    private
+      def call_limit(owner)
+        case limit
+        when nil, Integer then limit
+        when Symbol       then owner.public_send(limit)
+        else
+          limit.arity.zero? ? owner.instance_exec(&limit) : limit.call(owner)
+        end
+      end
+
+      def validate_declaration!
+        case limit
+        when nil then nil
+        when Integer then validate_non_negative!(limit)
+        when Symbol then nil
+        else
+          unless limit.respond_to?(:call)
+            raise ConfigurationError,
+                  "has_storage_quota limit: must be an Integer number of bytes, " \
+                  "nil (unlimited), a Symbol naming a method on the owner, or a " \
+                  "callable; got #{limit.inspect} (#{limit.class})"
+          end
+
+          unless SUPPORTED_ARITIES.include?(limit.arity)
+            raise ConfigurationError,
+                  "has_storage_quota limit: a callable must take no arguments " \
+                  "(evaluated against the owner) or one argument (the owner); " \
+                  "got arity #{limit.arity}"
+          end
+        end
+      end
+
+      def validate_non_negative!(bytes)
+        return bytes unless bytes.negative?
+
+        raise ConfigurationError,
+              "has_storage_quota limit: must be zero or more bytes; got #{bytes}"
+      end
+
+      def validate_resolved(bytes, owner)
+        return nil if bytes.nil?
+
+        unless bytes.is_a?(Integer)
+          raise ConfigurationError,
+                "#{owner.class.name} resolved its storage limit to " \
+                "#{bytes.inspect} (#{bytes.class}); expected an Integer number " \
+                "of bytes, or nil for unlimited"
+        end
+
+        if bytes.negative?
+          raise ConfigurationError,
+                "#{owner.class.name} resolved its storage limit to #{bytes}; " \
+                "expected zero or more bytes"
+        end
+
+        bytes
+      end
+  end
+end

data/lib/active_storage_quota/engine.rb

@@ -0,0 +1,28 @@
+# frozen_string_literal: true
+
+require "rails/engine"
+
+module ActiveStorageQuota
+  # Makes the gem's migrations available to the host application via
+  #
+  #   bin/rails active_storage_quota:install:migrations
+  #
+  # and its app/controllers autoloadable, so applications can subclass
+  # ActiveStorageQuota::DirectUploadsController without requiring anything.
+  #
+  # There is deliberately no rake_tasks block: Rails::Engine#run_tasks_blocks
+  # already loads every lib/tasks/*.rake in the engine. Loading them again here
+  # would define each task twice, and Rake *appends* actions rather than
+  # replacing them -- so every task body would run twice.
+  #
+  # The engine is otherwise thin: no routes, no controllers to mount, nothing to
+  # isolate. Applications opt into the direct-upload endpoint by subclassing and
+  # routing it themselves.
+  class Engine < ::Rails::Engine
+    # Without this, Rails derives the namespace from the class name and the
+    # migration task becomes active_storage_quota_engine:install:migrations,
+    # which is both undocumented and inconsistent with every other task this
+    # gem ships under active_storage_quota:.
+    railtie_name "active_storage_quota"
+  end
+end