"""Disk-usage accounting, per-user quota, and automatic eviction. An artifact's on-disk footprint is the primary stored file plus, for markdown, the preserved `.md` sidecar. 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 from datetime import timedelta from django.conf import settings from django.db.models import 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.""" def artifact_disk_size(relative_path, kind): """Bytes an artifact occupies on disk, including the markdown sidecar.""" root = settings.ARTIFACTS_ROOT total = 0 primary = root / relative_path if primary.exists(): total += primary.stat().st_size if kind == Artifact.MARKDOWN: sidecar = primary.with_suffix(".md") if sidecar.exists(): total += sidecar.stat().st_size return total def delete_artifact_files(artifact): """Remove an artifact's files from disk (primary, markdown 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) if artifact.kind == Artifact.MARKDOWN: primary.with_suffix(".md").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): artifact_bytes = ( Artifact.objects.filter(owner=user).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. """ 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 [] unkept = list( Artifact.objects.filter(owner=user, is_kept=False).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 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