"""Disk-usage accounting, per-user quota, and automatic eviction. An artifact's on-disk footprint is the primary stored file plus, for kinds that preserve their original source (see `SIDECAR_SUFFIXES`), the sidecar file. We record it on the row (`size_bytes`) at ingest so usage is a cheap SUM rather than a stat of every file. """ import shutil import threading from datetime import timedelta from django.conf import settings from django.db import connection from django.db.models import F, Sum from django.utils import timezone from .models import Artifact, Asset, UserProfile class QuotaExceeded(Exception): """Raised when a post cannot fit even after evicting every unkept artifact.""" # Kinds that preserve their original source alongside the rendered .html file. SIDECAR_SUFFIXES = { Artifact.MARKDOWN: ".md", Artifact.JSON: ".json", Artifact.CSV: ".csv", } def artifact_disk_size(relative_path, kind): """Bytes an artifact occupies on disk, including any sidecar file.""" root = settings.ARTIFACTS_ROOT total = 0 primary = root / relative_path if primary.exists(): total += primary.stat().st_size sidecar_suffix = SIDECAR_SUFFIXES.get(kind) if sidecar_suffix: sidecar = primary.with_suffix(sidecar_suffix) if sidecar.exists(): total += sidecar.stat().st_size return total def delete_artifact_files(artifact): """Remove an artifact's files from disk (primary, sidecar, assets). Assets are cascade-deleted from the DB by the FK, but their files must be removed here explicitly. """ root = settings.ARTIFACTS_ROOT primary = root / artifact.path primary.unlink(missing_ok=True) sidecar_suffix = SIDECAR_SUFFIXES.get(artifact.kind) if sidecar_suffix: primary.with_suffix(sidecar_suffix).unlink(missing_ok=True) for asset in artifact.assets.all(): delete_asset_file(asset) def delete_asset_file(asset): (settings.ARTIFACTS_ROOT / asset.path).unlink(missing_ok=True) def artifact_total_bytes(artifact): """An artifact's own bytes plus those of the assets embedded in it.""" asset_bytes = artifact.assets.aggregate(total=Sum("size_bytes")).get("total") or 0 return artifact.size_bytes + asset_bytes def user_usage_bytes(user): # Trashed artifacts are off-quota (deleting relieves quota immediately); # their files linger on disk until Empty Trash / purge, but don't count. artifact_bytes = ( Artifact.objects.filter(owner=user, deleted_at__isnull=True) .aggregate(total=Sum("size_bytes")) .get("total") or 0 ) asset_bytes = ( Asset.objects.filter(owner=user).aggregate(total=Sum("size_bytes")).get("total") or 0 ) return artifact_bytes + asset_bytes def user_quota_bytes(user): profile, _ = UserProfile.objects.get_or_create(user=user) return profile.quota_bytes def make_room_for(user, incoming_bytes): """Evict oldest unkept artifacts until `incoming_bytes` fits under quota. Returns the ids of evicted artifacts (already deleted - Django clears an instance's pk on .delete(), so ids are captured beforehand). Raises QuotaExceeded, without deleting anything, if even evicting every unkept artifact wouldn't free enough room. Must be called inside `transaction.atomic()` by the caller, wrapping through the subsequent `Artifact.objects.create()` too. The UPDATE below is a real write (not a plain SELECT) specifically so SQLite acquires its write lock *before* `used` is read - without that, two concurrent requests can each read a stale `used` during their unlocked read phase and both conclude "this fits" (issue #010). SQLite has no row-level locking, so this necessarily serializes quota-affecting writes across *all* users, not just the current one - an inherent SQLite limitation, acceptable at this project's scale. """ UserProfile.objects.get_or_create(user=user) UserProfile.objects.filter(user=user).update(quota_bytes=F("quota_bytes")) quota = user_quota_bytes(user) if incoming_bytes > quota: raise QuotaExceeded("Artifact is larger than the entire quota.") used = user_usage_bytes(user) if used + incoming_bytes <= quota: return [] # Trashed artifacts are eviction-protected (recovery window is guaranteed) # and already off-quota, so they're never eviction candidates. unkept = list( Artifact.objects.filter(owner=user, is_kept=False, deleted_at__isnull=True).order_by("created_at") ) # Check whether eviction can possibly succeed before deleting anything - # otherwise a doomed request would still destroy every unkept artifact on # its way to failing anyway. evictable_total = sum(artifact_total_bytes(artifact) for artifact in unkept) if used - evictable_total + incoming_bytes > quota: raise QuotaExceeded("Quota is full of kept artifacts; free space to post.") evicted = [] for artifact in unkept: if used + incoming_bytes <= quota: break # Freeing an artifact also frees the assets embedded in it. used -= artifact_total_bytes(artifact) evicted_id = artifact.id delete_artifact_files(artifact) artifact.delete() evicted.append(evicted_id) return evicted def purge_trashed(user): """Permanently remove all of a user's trashed artifacts (files + rows). Backs the "Empty Trash" action.""" trashed = list(Artifact.objects.filter(owner=user, deleted_at__isnull=False)) for artifact in trashed: delete_artifact_files(artifact) artifact.delete() return len(trashed) TRASH_RETENTION = timedelta(days=3) def purge_expired_trash(now=None): """Permanently remove trashed artifacts older than the retention window (all users). Tolerant of already-gone rows so concurrent purges are safe. Returns the number purged.""" now = now or timezone.now() cutoff = now - TRASH_RETENTION expired = list(Artifact.objects.filter(deleted_at__isnull=False, deleted_at__lt=cutoff)) purged = 0 for artifact in expired: try: delete_artifact_files(artifact) artifact.delete() purged += 1 except Exception: # Another concurrent purge got there first, or a file's already # gone — a purge error must never surface to anyone. pass return purged def spawn_trash_purge(): """Run purge_expired_trash off the request thread (cron-less: triggered by deletes). Falls back to inline when async is disabled (tests).""" if not getattr(settings, "AGENTFEED_ASYNC_PURGE", True): purge_expired_trash() return def _run(): try: purge_expired_trash() finally: connection.close() # threads get their own connection; don't leak it threading.Thread(target=_run, daemon=True).start() ASSET_ORPHAN_GRACE = timedelta(hours=24) def sweep_orphan_assets(user=None): """Delete assets uploaded but never referenced by an artifact. Runs opportunistically on upload. The grace period gives an agent time to upload images and then post the artifact that references them. """ cutoff = timezone.now() - ASSET_ORPHAN_GRACE orphans = Asset.objects.filter(artifact__isnull=True, created_at__lt=cutoff) if user is not None: orphans = orphans.filter(owner=user) count = 0 for asset in orphans: delete_asset_file(asset) asset.delete() count += 1 return count def disk_free_bytes(): """Free space on the volume that holds the artifact store.""" return shutil.disk_usage(settings.ARTIFACTS_ROOT).free def total_allocated_bytes(): return UserProfile.objects.aggregate(total=Sum("quota_bytes")).get("total") or 0