210 lines
8.5 KiB
Python
210 lines
8.5 KiB
Python
"""Space quotas, including descendants, trash, and reserved resumable uploads.
|
||
|
||
A child can be constrained by both its own allocation and its ancestors'.
|
||
Call assert_space_capacity() inside transaction.atomic() before reserving/writing.
|
||
"""
|
||
from contextlib import contextmanager
|
||
from django.db import transaction
|
||
from django.db.models import Q, Sum
|
||
from .models import DriveItem, DriveSpace, UploadSession
|
||
|
||
MAX_QUOTA_BYTES = 9_007_199_254_740_991 # exact integer in JS too
|
||
|
||
|
||
class QuotaExceeded(Exception):
|
||
pass
|
||
|
||
|
||
def parse_quota(value):
|
||
"""None = unlimited. Accept only an integer number of bytes (0 is valid)."""
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, bool) or isinstance(value, float):
|
||
raise ValueError('Le quota doit être un entier en octets, ou null pour illimité.')
|
||
if isinstance(value, str) and (not value.strip().isdigit()):
|
||
raise ValueError('Le quota doit être un entier en octets, ou null pour illimité.')
|
||
try:
|
||
quota = int(value)
|
||
except (TypeError, ValueError, OverflowError):
|
||
raise ValueError('Quota invalide.') from None
|
||
if quota < 0 or quota > MAX_QUOTA_BYTES:
|
||
raise ValueError('Quota hors limites (0 à 9 Po environ).')
|
||
return quota
|
||
|
||
|
||
def ancestry(space):
|
||
result = []
|
||
seen = set()
|
||
current = space
|
||
while current is not None:
|
||
if current.pk in seen:
|
||
raise ValueError('Hiérarchie des espaces invalide.')
|
||
seen.add(current.pk)
|
||
result.append(current)
|
||
current = current.parent
|
||
return result
|
||
|
||
|
||
def subtree_space_ids(space):
|
||
"""Use the stable DB path rather than a filesystem traversal."""
|
||
# The logical root has an empty relative_path and contains every space.
|
||
if space.parent_id is None and space.relative_path == '':
|
||
return DriveSpace.objects.all().values_list('pk', flat=True)
|
||
return DriveSpace.objects.filter(
|
||
Q(relative_path=space.relative_path) | Q(relative_path__startswith=space.relative_path + '/')
|
||
).values_list('pk', flat=True)
|
||
|
||
|
||
def occupancy(space):
|
||
"""Count ALL files, even in the trash. Pending sessions reserve their full size."""
|
||
ids = subtree_space_ids(space)
|
||
files = DriveItem.objects.filter(space_id__in=ids, is_folder=False).aggregate(n=Sum('size'))['n'] or 0
|
||
pending = UploadSession.objects.filter(space_id__in=ids).aggregate(n=Sum('total_size'))['n'] or 0
|
||
return int(files), int(pending)
|
||
|
||
|
||
def trash_occupancy(space):
|
||
"""Bytes still on disk in the trash, including descendants of trashed folders.
|
||
|
||
Ordinary (not trashed) folders require no traversal. Only traverse the
|
||
descendants of explicitly trashed folders; avoid counting trashed files twice.
|
||
"""
|
||
ids = subtree_space_ids(space)
|
||
bytes_in_trash = DriveItem.objects.filter(
|
||
space_id__in=ids, is_folder=False, trashed_at__isnull=False
|
||
).aggregate(n=Sum('size'))['n'] or 0
|
||
folder_ids = list(DriveItem.objects.filter(
|
||
space_id__in=ids, is_folder=True, trashed_at__isnull=False
|
||
).values_list('pk', flat=True))
|
||
visited = set(folder_ids)
|
||
while folder_ids:
|
||
next_folders = []
|
||
for start in range(0, len(folder_ids), 300):
|
||
for pk, is_folder, size, trashed in DriveItem.objects.filter(
|
||
space_id__in=ids, parent_id__in=folder_ids[start:start + 300]
|
||
).values_list('pk', 'is_folder', 'size', 'trashed_at'):
|
||
if is_folder:
|
||
if pk not in visited:
|
||
visited.add(pk)
|
||
next_folders.append(pk)
|
||
elif trashed is None:
|
||
bytes_in_trash += size
|
||
folder_ids = next_folders
|
||
return int(bytes_in_trash)
|
||
|
||
|
||
def _budget_tree(space):
|
||
"""One DB snapshot of all descendant spaces and their direct physical usage.
|
||
|
||
A child's fixed quota reserves that much capacity for it; a child without
|
||
a quota reserves only its actual consumption (including its descendants).
|
||
Thus quotas of siblings cannot silently overcommit their common parent.
|
||
"""
|
||
ids=list(subtree_space_ids(space))
|
||
nodes={node.pk:node for node in DriveSpace.objects.filter(pk__in=ids)}
|
||
children={pk:[] for pk in nodes}
|
||
for node in nodes.values():
|
||
if node.parent_id in children:children[node.parent_id].append(node.pk)
|
||
direct={pk:0 for pk in nodes}
|
||
for pk,amount in DriveItem.objects.filter(space_id__in=ids,is_folder=False).values('space_id').annotate(total=Sum('size')).values_list('space_id','total'):
|
||
direct[pk]+=int(amount or 0)
|
||
for pk,amount in UploadSession.objects.filter(space_id__in=ids).values('space_id').annotate(total=Sum('total_size')).values_list('space_id','total'):
|
||
direct[pk]+=int(amount or 0)
|
||
memo={}
|
||
def committed(pk):
|
||
if pk not in memo:
|
||
memo[pk]=direct[pk]+sum(max(nodes[c].quota_bytes or 0,committed(c)) for c in children[pk])
|
||
return memo[pk]
|
||
committed(space.pk)
|
||
return memo
|
||
|
||
|
||
def _committed(space):
|
||
"""Allocated bytes inside an entire subtree, including nested quota reservations."""
|
||
return _budget_tree(space)[space.pk]
|
||
|
||
|
||
def available_for_growth(space):
|
||
"""Bytes that can be added directly to space without breaking any ancestor allocation.
|
||
|
||
Capacity already reserved for this child does not count against its parent's
|
||
available budget a second time. A zero-byte quota correctly permits no growth.
|
||
"""
|
||
maximum=MAX_QUOTA_BYTES
|
||
absorbed=0
|
||
chain=ancestry(space)
|
||
committed_tree=_budget_tree(chain[-1])
|
||
for ancestor in chain:
|
||
committed=committed_tree[ancestor.pk]
|
||
if ancestor.quota_bytes is not None:
|
||
maximum=min(maximum,ancestor.quota_bytes-committed+absorbed)
|
||
absorbed+=max(0,ancestor.quota_bytes-committed)
|
||
return None if maximum==MAX_QUOTA_BYTES else max(0,int(maximum))
|
||
|
||
|
||
def assert_new_child_quota(parent, quota):
|
||
"""A new nested space cannot reserve more than its parent's unallocated capacity."""
|
||
if quota is None:return
|
||
remaining=available_for_growth(parent)
|
||
if remaining is not None and quota>remaining:
|
||
raise ValueError(f'Le quota dépasse les {remaining} octets disponibles dans l’espace parent.')
|
||
|
||
|
||
def assert_updated_quota(space, quota):
|
||
"""Validate nested allocations, including sibling limits and actual disk use.
|
||
|
||
The old reservation is replaced atomically with the requested value. Parent
|
||
quotas must be large enough for their direct content and all child quotas.
|
||
"""
|
||
chain=ancestry(space)
|
||
committed_tree=_budget_tree(chain[-1])
|
||
committed=committed_tree[space.pk]
|
||
if quota is not None and quota<committed:
|
||
raise ValueError('Le quota est inférieur aux fichiers et sous-espaces déjà attribués.')
|
||
current=space.quota_bytes
|
||
if current==quota:return
|
||
# Compute the changed reservation at every ancestor: each parent's
|
||
# committed usage has this subtree's reserved/actual allocation substituted.
|
||
old_child=max(current or 0,committed)
|
||
new_child=max(quota or 0,committed)
|
||
for ancestor in chain[1:]:
|
||
amount=committed_tree[ancestor.pk]-old_child+new_child
|
||
if ancestor.quota_bytes is not None and amount>ancestor.quota_bytes:
|
||
raise ValueError(f'Quota impossible : capacité de « {ancestor.display_name} » dépassée.')
|
||
old_child=max(ancestor.quota_bytes or 0,committed_tree[ancestor.pk])
|
||
new_child=max(ancestor.quota_bytes or 0,amount)
|
||
|
||
|
||
def quota_state(space):
|
||
used,reserved=occupancy(space)
|
||
trashed=trash_occupancy(space)
|
||
available=available_for_growth(space)
|
||
return {
|
||
'quota_bytes':space.quota_bytes,
|
||
'used_bytes':used,
|
||
'active_bytes':max(0,used-trashed),
|
||
'trashed_bytes':trashed,
|
||
'reserved_bytes':reserved,
|
||
'available_bytes':available,
|
||
'effective_limit_bytes':None if available is None else used+reserved+available,
|
||
}
|
||
|
||
|
||
@contextmanager
|
||
def locked_quota_tree(space):
|
||
"""Lock ancestors in deterministic order so uploads and limit changes serialize."""
|
||
if not transaction.get_connection().in_atomic_block:
|
||
raise RuntimeError('locked_quota_tree must run inside transaction.atomic()')
|
||
ids = sorted(a.pk for a in ancestry(space))
|
||
list(DriveSpace.objects.select_for_update().filter(pk__in=ids).order_by('pk'))
|
||
yield
|
||
|
||
|
||
def assert_space_capacity(space, additional_bytes):
|
||
"""Called with ancestor locks; include sibling allocations when checking uploads."""
|
||
available=available_for_growth(space)
|
||
if available is not None and additional_bytes>available:
|
||
raise QuotaExceeded(
|
||
f'Espace insuffisant : {available} octets encore disponibles après les quotas des sous-espaces.'
|
||
)
|