Attachments stored once; nightly clean-up of unused files

Attachments are no longer stored apart from their message (they were kept twice: base64 inside the raw message and decoded on their own). The indexer records the MIME part (attachments.part_index, migration 0023); downloads, zip, inline images, thumbnails and forwarding take it out of the message; videos stream from a short-lived copy (kept a day, 1 GiB per user). A background pass on worker start removes the copies stored before, each only after checking it can be taken out again (spec App. D §D.6 item 1, criterion corrected to <= 60%). A nightly job (03:47, after the Trash/Spam purge) removes files nothing refers to any more — deleted or purged messages, previews of attachments gone, uploads no draft has used for a week, expired streaming copies. Files younger than an hour are left alone; a sent message waiting for its Sent copy is held until its job is done. Account removal keeps files a draft still uses. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-08 20:20:21 +02:00
parent 0514053252
commit 3957dd540a
17 changed files with 1063 additions and 27 deletions
@@ -0,0 +1,237 @@
+"""Attachments, kept once: inside their message (SPEC App. D §D.6 item 1).
+
+An attachment is not stored apart from its message any more — it was, which
+kept most of a mailbox's bytes twice (once base64 in the raw message, once
+decoded on its own). It is taken out of the raw message when wanted:
+
+* :func:`load` — its bytes (downloads, zip, inline images, forwarding, thumbnails);
+* :func:`file` — a file holding it, for streaming with HTTP Range (videos):
+  taken out into ``{user}/cache/`` and kept a day after its last use;
+* :func:`dedupe_existing` — the one-time pass removing the separate copies
+  stored before, each only after checking it can be taken out of its message.
+
+Blocking (parsing, disk): call these in a thread.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import hashlib
+import logging
+import time
+import uuid
+from collections import defaultdict
+from pathlib import Path
+
+from app.indexer import blobstore
+from app.indexer.mime import load_message, part_payload
+
+logger = logging.getLogger(__name__)
+
+CACHE_KEEP_S = 24 * 3600  # a taken-out file is kept this long after its last use
+CACHE_MAX_BYTES = 1024**3  # per user; the least recently used go first beyond it
+# part_index of a row whose content is not in its message (kept as its own file).
+NOT_IN_MESSAGE = -1
+
+
+def _legacy(user_id: uuid.UUID, sha: str) -> Path | None:
+    """The separate copy stored before attachments were kept once, if still there."""
+    path = blobstore.blob_path(user_id, sha)
+    return path if path.exists() else None
+
+
+def load(user_id: uuid.UUID, message_sha: str, sha: str, part_index: int | None) -> bytes:
+    """An attachment's bytes. Raises FileNotFoundError when they cannot be had."""
+    legacy = _legacy(user_id, sha)
+    if legacy is not None:
+        return legacy.read_bytes()
+    raw = blobstore.blob_path(user_id, message_sha)
+    if part_index == NOT_IN_MESSAGE or not raw.exists():
+        raise FileNotFoundError(sha)
+    payload = part_payload(load_message(path=raw), part_index, sha)
+    if payload is None:
+        raise FileNotFoundError(sha)
+    return payload
+
+
+def materialize(user_id: uuid.UUID, message_sha: str, sha: str, part_index: int | None) -> bool:
+    """Make sure the attachment exists as a blob of its own — for a compose
+    attachment (forward, a draft reopened), which is read by its hash. False
+    when its bytes cannot be had."""
+    if blobstore.exists(user_id, sha):
+        return True
+    try:
+        blobstore.store(user_id, load(user_id, message_sha, sha, part_index))
+    except FileNotFoundError:
+        return False
+    return True
+
+
+def _cache_dir(user_id: uuid.UUID) -> Path:
+    return blobstore.blob_path(user_id, "00").parent.parent / "cache"
+
+
+def prune_cache(cache: Path) -> None:
+    """Drop streaming copies past their keep time, and the oldest beyond the cap."""
+    files = [(p, p.stat()) for p in cache.iterdir() if p.is_file()]
+    now = time.time()
+    total = 0
+    for path, st in sorted(files, key=lambda f: f[1].st_mtime, reverse=True):
+        total += st.st_size
+        if now - st.st_mtime > CACHE_KEEP_S or total > CACHE_MAX_BYTES:
+            path.unlink(missing_ok=True)
+
+
+def file(user_id: uuid.UUID, message_sha: str, sha: str, part_index: int | None) -> Path:
+    """A file holding the attachment, to stream from. Raises FileNotFoundError."""
+    legacy = _legacy(user_id, sha)
+    if legacy is not None:
+        return legacy
+    cache = _cache_dir(user_id)
+    path = cache / sha
+    if path.exists():
+        path.touch()  # used again: kept longer
+        return path
+    data = load(user_id, message_sha, sha, part_index)
+    cache.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(f".{uuid.uuid4().hex}.tmp")
+    tmp.write_bytes(data)
+    tmp.replace(path)
+    with contextlib.suppress(OSError):  # a file removed meanwhile by another request
+        prune_cache(cache)
+    return path
+
+
+def drop_cached(user_id: uuid.UUID) -> None:
+    """Forget every taken-out file of a user (account removed)."""
+    cache = _cache_dir(user_id)
+    if cache.exists():
+        for path in cache.iterdir():
+            path.unlink(missing_ok=True)
+
+
+# ---- removing the copies stored before ------------------------------------------------
+
+# A file this recent may belong to work in progress (a sent message waiting
+# for its Sent copy is referenced only by its job): left for the next pass.
+_MIN_AGE_S = 3600
+
+
+def _part_indexes(raw_path: Path) -> dict[str, int]:
+    """Content hash → MIME part position, for every leaf part of a message."""
+    out: dict[str, int] = {}
+    for index, part in enumerate(load_message(path=raw_path).walk()):
+        if part.is_multipart():
+            continue
+        payload = part.get_payload(decode=True)
+        if isinstance(payload, bytes):
+            out.setdefault(hashlib.sha256(payload).hexdigest(), index)
+    return out
+
+
+async def dedupe_existing(batch: int = 200, *, user_id: uuid.UUID | None = None) -> dict[str, int]:
+    """Record where each older attachment sits in its message, and delete its
+    separate copy once it is known to be there. Idempotent; resumable."""
+    import asyncio
+
+    from sqlalchemy import select, update
+
+    from app.db import session_scope
+    from app.models import Attachment, ComposeAttachment, Message
+
+    stats = {"rows": 0, "found": 0, "not_found": 0, "files_removed": 0, "bytes_freed": 0}
+    while True:
+        async with session_scope() as db:
+            rows = (
+                await db.execute(
+                    select(
+                        Attachment.id,
+                        Attachment.user_id,
+                        Attachment.blob_sha256,
+                        Message.blob_sha256,
+                    )
+                    .join(Message, Message.id == Attachment.message_id)
+                    .where(
+                        Attachment.part_index.is_(None),
+                        *([Attachment.user_id == user_id] if user_id is not None else []),
+                    )
+                    .order_by(Attachment.message_id)
+                    .limit(batch)
+                )
+            ).all()
+        if not rows:
+            return stats
+        by_message: dict[tuple[uuid.UUID, str], list[tuple[uuid.UUID, str]]] = defaultdict(list)
+        for att_id, owner, sha, msg_sha in rows:
+            by_message[(owner, msg_sha)].append((att_id, sha))
+
+        located: dict[uuid.UUID, int] = {}
+        for (owner, msg_sha), atts in by_message.items():
+            raw = blobstore.blob_path(owner, msg_sha)
+            try:
+                indexes = await asyncio.to_thread(_part_indexes, raw) if raw.exists() else {}
+            except Exception:  # noqa: BLE001 - unreadable message: keep the copies
+                logger.warning("attachments: cannot parse %s", raw, exc_info=True)
+                indexes = {}
+            for att_id, sha in atts:
+                located[att_id] = indexes.get(sha, NOT_IN_MESSAGE)
+
+        async with session_scope() as db:
+            for att_id, index in located.items():
+                await db.execute(
+                    update(Attachment).where(Attachment.id == att_id).values(part_index=index)
+                )
+        stats["rows"] += len(located)
+        stats["found"] += sum(1 for i in located.values() if i != NOT_IN_MESSAGE)
+        stats["not_found"] += sum(1 for i in located.values() if i == NOT_IN_MESSAGE)
+
+        # A separate copy goes only when every attachment row with that content
+        # is found in its message and nothing else uses the file.
+        candidates = {
+            (owner, sha)
+            for att_id, owner, sha, _ in rows
+            if located.get(att_id, NOT_IN_MESSAGE) != NOT_IN_MESSAGE
+        }
+        async with session_scope() as db:
+            for owner, sha in candidates:
+                path = blobstore.blob_path(owner, sha)
+                if not path.exists() or time.time() - path.stat().st_mtime < _MIN_AGE_S:
+                    continue
+                in_use = (
+                    await db.scalar(
+                        select(Message.id)
+                        .where(Message.user_id == owner, Message.blob_sha256 == sha)
+                        .limit(1)
+                    )
+                    or await db.scalar(
+                        select(ComposeAttachment.id)
+                        .where(
+                            ComposeAttachment.user_id == owner,
+                            ComposeAttachment.blob_sha256 == sha,
+                        )
+                        .limit(1)
+                    )
+                    or await db.scalar(
+                        select(Attachment.id)
+                        .where(
+                            Attachment.user_id == owner,
+                            Attachment.blob_sha256 == sha,
+                            (Attachment.part_index.is_(None))
+                            | (Attachment.part_index == NOT_IN_MESSAGE),
+                        )
+                        .limit(1)
+                    )
+                )
+                if in_use:
+                    continue
+                size = path.stat().st_size
+                path.unlink(missing_ok=True)
+                stats["files_removed"] += 1
+                stats["bytes_freed"] += size
+
+
+async def dedupe_existing_job(ctx: dict[str, object]) -> dict[str, int]:
+    """arq job (bulk queue): :func:`dedupe_existing`, run once after an update."""
+    stats = await dedupe_existing()
+    logger.info("attachments kept once: %s", stats)
+    return stats