webmail-client/backend/app/indexer/gc.py

"""The nightly clean-up of stored files nothing uses any more.

Messages leave in many ways — Delete forever, Empty Trash, the Trash/Spam
retention purge, an expunge on the server, a folder re-synced — and none of
them removes the message's file: files are stored by content, so one file can
serve several rows (the same email in two folders, a forwarded file in a
draft), and removing it is only safe once nothing refers to it. This pass does
that for every user, whatever the way the rows went:

* uploads (compose attachments) no draft has referred to for a week — left by
  a sent message (its draft is gone) or a composer closed without saving;
* message/attachment files no row refers to;
* attachment previews (thumbnails) of attachments that are gone;
* streaming copies past their keep time (app/indexer/attachments.py).

A file younger than an hour is left alone: it may belong to work in progress
(a message being stored, a sent message waiting for its Sent copy).
"""

from __future__ import annotations

import datetime as dt
import logging
import re
import time
import uuid
from pathlib import Path
from typing import Any

from sqlalchemy import delete, or_, select

from app.config import get_settings
from app.indexer import attachments as att_store

logger = logging.getLogger(__name__)

MIN_AGE_S = 3600
UPLOAD_KEEP = dt.timedelta(days=7)
_SHA = re.compile(r"^[0-9a-f]{64}$")


# ---- files a job still needs ---------------------------------------------------------
#
# A sent message whose Sent copy could not be filed yet is referenced only by
# its retrying job (app/outbox/jobs.py), for up to a day: held here meanwhile.
# Entries older than HOLD_MAX are ignored, in case a job died without letting go.

HOLD_MAX_S = 3 * 24 * 3600


def _hold_key(user_id: uuid.UUID | str) -> str:
    return f"blob:held:{user_id}"


async def hold(user_id: uuid.UUID | str, sha: str) -> None:
    from app.redis_client import get_redis

    await get_redis().zadd(_hold_key(user_id), {sha: time.time()})


async def release(user_id: uuid.UUID | str, sha: str) -> None:
    from app.redis_client import get_redis

    await get_redis().zrem(_hold_key(user_id), sha)


async def _held(user_id: uuid.UUID) -> set[str]:
    from app.redis_client import get_redis

    redis = get_redis()
    oldest = time.time() - HOLD_MAX_S
    await redis.zremrangebyscore(_hold_key(user_id), "-inf", oldest)
    return {
        m.decode() if isinstance(m, bytes) else str(m)
        for m in await redis.zrangebyscore(_hold_key(user_id), oldest, "+inf")
    }


def _draft_upload_ids(payloads: list[dict[str, Any]]) -> set[uuid.UUID]:
    out: set[uuid.UUID] = set()
    for payload in payloads:
        for value in payload.get("attachment_ids") or []:
            try:
                out.add(uuid.UUID(str(value)))
            except ValueError:
                continue
    return out


def _old(path: Path, now: float) -> bool:
    return now - path.stat().st_mtime >= MIN_AGE_S


async def sweep_user(user_id: uuid.UUID) -> dict[str, int]:
    """Clean one user's stored files; returns what was removed."""
    from app.db import session_scope
    from app.models import Attachment, ComposeAttachment, Draft, Message

    stats = {"uploads": 0, "files": 0, "bytes": 0, "thumbs": 0}
    async with session_scope() as db:
        # Uploads no draft refers to, once a week old.
        used = _draft_upload_ids(
            list(await db.scalars(select(Draft.payload).where(Draft.user_id == user_id)))
        )
        stale = select(ComposeAttachment.id).where(
            ComposeAttachment.user_id == user_id,
            ComposeAttachment.created_at < dt.datetime.now(dt.UTC) - UPLOAD_KEEP,
        )
        if used:
            stale = stale.where(ComposeAttachment.id.not_in(used))
        gone = list(await db.scalars(stale))
        if gone:
            await db.execute(delete(ComposeAttachment).where(ComposeAttachment.id.in_(gone)))
        stats["uploads"] = len(gone)

    async with session_scope() as db:
        # Every file still referred to. Attachments count only where their file
        # is their only copy (not found in the message, or not yet looked for).
        referenced = set(
            await db.scalars(select(Message.blob_sha256).where(Message.user_id == user_id))
        )
        referenced |= set(
            await db.scalars(
                select(ComposeAttachment.blob_sha256).where(ComposeAttachment.user_id == user_id)
            )
        )
        referenced |= set(
            await db.scalars(
                select(Attachment.blob_sha256).where(
                    Attachment.user_id == user_id,
                    or_(
                        Attachment.part_index.is_(None),
                        Attachment.part_index == att_store.NOT_IN_MESSAGE,
                    ),
                )
            )
        )
        referenced |= await _held(user_id)
        attachment_shas = set(
            await db.scalars(select(Attachment.blob_sha256).where(Attachment.user_id == user_id))
        )

    root = Path(get_settings().blob_root) / str(user_id)
    if not root.is_dir():
        return stats
    now = time.time()
    for path in root.glob("??/*"):
        if not path.is_file() or not _SHA.match(path.name) or path.name in referenced:
            continue
        if not _old(path, now):
            continue
        size = path.stat().st_size
        path.unlink(missing_ok=True)
        stats["files"] += 1
        stats["bytes"] += size
    thumbs = root / "thumbs"
    if thumbs.is_dir():
        for path in thumbs.iterdir():
            sha = path.name.split("-", 1)[0]
            if (
                path.is_file()
                and _SHA.match(sha)
                and sha not in attachment_shas
                and _old(path, now)
            ):
                path.unlink(missing_ok=True)
                stats["thumbs"] += 1
    cache = root / "cache"
    if cache.is_dir():
        att_store.prune_cache(cache)
    return stats


async def sweep(user_id: uuid.UUID | None = None) -> dict[str, int]:
    """:func:`sweep_user` for one user, or every user."""
    from app.db import session_scope
    from app.models import User

    async with session_scope() as db:
        users = [user_id] if user_id is not None else list(await db.scalars(select(User.id)))
    total = {"users": 0, "uploads": 0, "files": 0, "bytes": 0, "thumbs": 0}
    for uid in users:
        try:
            stats = await sweep_user(uid)
        except Exception:  # noqa: BLE001 - one user's trouble must not stop the rest
            logger.warning("stored-file clean-up failed for user %s", uid, exc_info=True)
            continue
        total["users"] += 1
        for key, value in stats.items():
            total[key] += value
    return total


async def sweep_job(ctx: dict[str, Any]) -> dict[str, int]:
    """arq cron job: the nightly :func:`sweep`."""
    stats = await sweep()
    logger.info("stored-file clean-up: %s", stats)
    return stats