"""Per-user download limits + footprint accounting. A user's *footprint* is the total size of the ready assets their library references (cache hits count — it's "how much disk you're responsible for", not "how much you personally downloaded"). An upload counts from the moment it lands, ready or not: its original is on disk while its conversion waits, and stays there when the conversion fails (for Retry) — bytes a ready-only sum never saw. Limits come from a per-user DownloadQuota row if the admin set one, else the sysconfig defaults. `unlimited` (e.g. the admin's own account) bypasses the byte cap. Enforcement split: max_jobs + max_bytes are checked at enqueue (they bound holdings); per-user max_concurrent is enforced by the worker at claim time (it bounds simultaneous downloads). DOWNLOAD jobs additionally consult `remaining_bytes` in the worker, because enqueue can only compare the footprint as it stood when the job was created and earlier jobs may have filled the cap since. That is a per-job ceiling, not an aggregate ledger: concurrent jobs of one user each read the same remainder (in-flight bytes are not `ready`, so `footprint` cannot see them), so a wave of max_concurrent downloads can still land over the cap — the deliberate trade, since a download's size is unknown until it finishes. CLIP jobs consult it before their fetch or encode and again with their output's real size before it is placed. Edit jobs do not consult it at all (tracked as SEC-19). """ from dataclasses import dataclass from sqlalchemy import and_, func, or_, select from sqlalchemy.orm import Session from app import sysconfig from app.downloads.status import JobKind, JobStatus from app.downloads.models import DownloadJob, DownloadQuota, MediaAsset from app.userscope import owned_by # Jobs a user is considered to be "holding" (occupy a job slot / footprint). Terminal failures # (error, canceled) don't count. _HOLDING = (JobStatus.QUEUED, JobStatus.RUNNING, JobStatus.PAUSED, JobStatus.DONE) _RUNNING = (JobStatus.QUEUED, JobStatus.RUNNING) # An asset whose file is on disk: a ready one, or an upload in any state (its original lands at # ingest, see the module docstring). Every byte sum — a user's footprint, the GC's total-cap pass, # the admin's storage view — reads it, so they all count the same disk. ON_DISK = or_( MediaAsset.status == "ready", and_(MediaAsset.source_kind == "upload", MediaAsset.rel_path.is_not(None)), ) @dataclass class QuotaLimits: max_bytes: int max_concurrent: int max_jobs: int unlimited: bool class QuotaExceeded(Exception): """Raised by check_enqueue; `reason` is one of max_jobs|max_bytes (i18n key on the client).""" def __init__(self, reason: str, limit: int, current: int): self.reason = reason self.limit = limit self.current = current super().__init__(f"download quota exceeded: {reason} ({current}/{limit})") def resolve(db: Session, user_id: int) -> QuotaLimits: row = db.get(DownloadQuota, user_id) if row is not None: # Floor max_concurrent at 1 (admin_set_quota already does on write; the sysconfig default's # min is 1). A stored 0 is nonsensical for a concurrency limit AND would diverge the two # "at limit" checks — the worker's claim-window exclusion vs at_concurrency_limit — since a # 0-limit user is "capped" with zero running jobs, which the SQL exclusion can't see. Keep # them consistent by never letting the effective limit drop below 1. (To BLOCK a user's # downloads, suspend them — the worker leaves a suspended user's queued jobs unclaimed — or # set max_jobs=0, not max_concurrent=0.) return QuotaLimits(row.max_bytes, max(1, row.max_concurrent), row.max_jobs, row.unlimited) return QuotaLimits( max_bytes=sysconfig.get_int(db, "download_default_max_bytes"), max_concurrent=max(1, sysconfig.get_int(db, "download_default_max_concurrent")), max_jobs=sysconfig.get_int(db, "download_default_max_jobs"), unlimited=False, ) def footprint(db: Session, user_id: int) -> int: """Total bytes of the distinct assets this user's jobs hold on disk: the ready ones their active jobs reference, and every upload of theirs — including one whose conversion failed, which keeps its original (see the module docstring).""" held_assets = ( select(DownloadJob.asset_id) .where( owned_by(DownloadJob.user_id, user_id), or_( DownloadJob.status.in_(_HOLDING), and_(DownloadJob.job_kind == JobKind.UPLOAD, DownloadJob.status == JobStatus.ERROR), ), DownloadJob.asset_id.is_not(None), ) .distinct() ) return int( db.execute( select(func.coalesce(func.sum(MediaAsset.size_bytes), 0)).where( MediaAsset.id.in_(held_assets), ON_DISK ) ).scalar() or 0 ) def _count(db: Session, user_id: int, states) -> int: return ( db.execute( select(func.count()) .select_from(DownloadJob) .where(owned_by(DownloadJob.user_id, user_id), DownloadJob.status.in_(states)) ).scalar() or 0 ) def active_jobs(db: Session, user_id: int) -> int: return _count(db, user_id, _HOLDING) def running_jobs(db: Session, user_id: int) -> int: return _count(db, user_id, _RUNNING) def at_concurrency_limit(db: Session, user_id: int) -> bool: """Worker-side gate: is this user already downloading their max? Unlimited users never are.""" lim = resolve(db, user_id) if lim.unlimited: return False running = _count(db, user_id, (JobStatus.RUNNING,)) return running >= lim.max_concurrent def check_enqueue(db: Session, user_id: int) -> None: """Raise QuotaExceeded if the user can't take another job (holdings caps). Concurrency is enforced later by the worker, so a big queue is allowed; it just drains max_concurrent-wide.""" lim = resolve(db, user_id) if lim.unlimited: return n = active_jobs(db, user_id) if n >= lim.max_jobs: raise QuotaExceeded("max_jobs", lim.max_jobs, n) used = footprint(db, user_id) if used >= lim.max_bytes: raise QuotaExceeded("max_bytes", lim.max_bytes, used) def remaining_bytes(db: Session, user_id: int) -> int | None: """Storage this user may still spend on ONE download: ``None`` for an unlimited user, else the cap minus the current footprint, floored at 0. Zero is a real answer, not "no limit" — it means the cap is already met and the caller must REFUSE rather than proceed unbounded. `check_enqueue` cannot make that refusal on the caller's behalf: it ran when the job was queued, while the user was still under the cap, and the jobs that filled it have landed since. Keeping the arithmetic here (rather than inline in the worker) is what keeps every max_bytes decision in this module.""" lim = resolve(db, user_id) if lim.unlimited: return None return max(0, lim.max_bytes - footprint(db, user_id)) def check_upload(db: Session, user_id: int, add_bytes: int) -> None: """Like check_enqueue, but for an upload whose size is KNOWN at ingest: also reject when adding this file would push the footprint over max_bytes. A download can't pre-check its size (unknown until the worker finishes) so check_enqueue only compares the current footprint; an upload can, so it's charged against the cap before it lands. Raises QuotaExceeded(max_jobs|max_bytes).""" lim = resolve(db, user_id) if lim.unlimited: return n = active_jobs(db, user_id) if n >= lim.max_jobs: raise QuotaExceeded("max_jobs", lim.max_jobs, n) projected = footprint(db, user_id) + add_bytes if projected > lim.max_bytes: raise QuotaExceeded("max_bytes", lim.max_bytes, projected) def fits_after_conversion(db: Session, user_id: int, old_bytes: int, new_bytes: int) -> bool: """Whether an upload converted from `old_bytes` (what it is charged at now) to `new_bytes` keeps its owner within the cap. A re-encode can come out larger than the original, and nothing else re-checks it once the upload was let in.""" if new_bytes <= old_bytes: # a smaller file never breaks a cap the larger one was let in under return True lim = resolve(db, user_id) if lim.unlimited: return True return footprint(db, user_id) - old_bytes + new_bytes <= lim.max_bytes def usage(db: Session, user_id: int) -> dict: """Compact usage snapshot for the UI (footprint + counts vs limits).""" lim = resolve(db, user_id) return { "footprint_bytes": footprint(db, user_id), "active_jobs": active_jobs(db, user_id), "running_jobs": running_jobs(db, user_id), "max_bytes": lim.max_bytes, "max_concurrent": lim.max_concurrent, "max_jobs": lim.max_jobs, "unlimited": lim.unlimited, } def set_custom( db: Session, user_id: int, *, max_bytes: int | None = None, max_concurrent: int | None = None, max_jobs: int | None = None, unlimited: bool | None = None, ) -> None: """The admin's own limits for a user: starts from what applies now, then each value given (a size floored at 0, a count at 1). The handler commits.""" row = db.get(DownloadQuota, user_id) if row is None: cur = resolve(db, user_id) row = DownloadQuota( user_id=user_id, max_bytes=cur.max_bytes, max_concurrent=cur.max_concurrent, max_jobs=cur.max_jobs, unlimited=cur.unlimited, ) db.add(row) if max_bytes is not None: row.max_bytes = max(0, max_bytes) if max_concurrent is not None: row.max_concurrent = max(1, max_concurrent) if max_jobs is not None: row.max_jobs = max(1, max_jobs) if unlimited is not None: row.unlimited = unlimited db.flush() def reset_custom(db: Session, user_id: int) -> bool: """Back to the defaults; False when the user had no limits of their own.""" row = db.get(DownloadQuota, user_id) if row is None: return False db.delete(row) db.flush() return True