stackpilot/backend/services/prune_job.py

"""Prunes with visible progress, shared by images and volumes.

``docker image prune`` / ``docker volume prune`` are single blocking calls that
report nothing until they are done, which on a host with a lot to delete means
minutes of a spinner and, past the proxy's read timeout, an error while the
daemon carries on. Instead, a service lists what the prune would remove and a
PruneJob deletes it one item at a time, publishing how far it is so the UI can
poll it. One prune of each kind runs at a time.
"""
from __future__ import annotations

import logging
import threading
from dataclasses import dataclass
from typing import Any, Callable, Optional

from docker_client import DockerError

logger = logging.getLogger(__name__)

_MAX_ERRORS = 10


class PruneBusy(Exception):
    """A prune of this kind is already running."""


@dataclass
class PrunePlan:
    """What one prune will do, built by the owning service.

    ``remove`` deletes one item and returns the entries it contributes to the
    result plus whether the item is actually gone (removing one of several
    image tags is not). It raises DockerError to mark just that item failed.
    ``prepare`` runs before the first removal (e.g. measuring disk usage) and
    ``finish`` after the last; it gets every collected entry and returns the
    API result together with the space reclaimed, or None if unknown.
    """

    all: bool
    items: list
    label: Callable[[Any], str]
    remove: Callable[[Any], tuple[list, bool]]
    finish: Callable[[list], tuple[dict, Optional[int]]]
    prepare: Optional[Callable[[], None]] = None


class PruneJob:
    def __init__(self, what: str):
        self._what = what
        self._lock = threading.Lock()
        self._status: dict = {
            "active": False,
            "run": 0,
            "all": False,
            # preparing -> removing -> finishing, so the UI can say what the
            # parts without a count are doing.
            "phase": None,
            "total": 0,
            "done": 0,
            "current": None,
            "deleted": 0,
            "failed": 0,
            "errors": [],
            "reclaimed": None,
            "error": None,
        }

    def status(self) -> dict:
        return {**self._status, "errors": list(self._status["errors"])}

    def run(self, build: Callable[[], PrunePlan]) -> dict:
        """Prune and wait for the result."""
        plan = self._begin(build)
        error = None
        try:
            return self._execute(plan)
        except DockerError as exc:
            error = exc.detail or exc.error
            raise
        finally:
            self._finish(error)

    def start(
        self,
        build: Callable[[], PrunePlan],
        on_done: Optional[Callable[[dict], None]] = None,
    ) -> dict:
        """Prune in a background thread and return the initial status.

        The plan is still built here, so a daemon that cannot even list the
        candidates fails the request instead of a status nobody polls yet.
        ``on_done`` gets the result from the worker thread.
        """
        plan = self._begin(build)

        def work() -> None:
            error = None
            result = None
            try:
                result = self._execute(plan)
            except Exception as exc:  # noqa: BLE001 - reported through the status
                logger.warning("%s prune failed: %s", self._what.capitalize(), exc)
                error = exc.detail if isinstance(exc, DockerError) else str(exc)
            finally:
                self._finish(error)
            if result is not None and on_done:
                try:
                    on_done(result)
                except Exception as exc:  # noqa: BLE001 - the prune itself succeeded
                    logger.warning("%s prune follow-up failed: %s", self._what, exc)

        threading.Thread(target=work, name=f"{self._what}-prune", daemon=True).start()
        return self.status()

    def _begin(self, build: Callable[[], PrunePlan]) -> PrunePlan:
        if not self._lock.acquire(blocking=False):
            raise PruneBusy(f"{self._what.capitalize()} prune is already running")
        try:
            plan = build()
        except BaseException:
            self._lock.release()
            raise
        self._status.update(
            {
                "active": True,
                "run": self._status["run"] + 1,
                "all": plan.all,
                "phase": "preparing",
                "total": len(plan.items),
                "done": 0,
                "current": None,
                "deleted": 0,
                "failed": 0,
                "errors": [],
                "reclaimed": None,
                "error": None,
            }
        )
        return plan

    def _execute(self, plan: PrunePlan) -> dict:
        status = self._status
        if plan.prepare:
            plan.prepare()
        status["phase"] = "removing"
        entries: list = []
        for item in plan.items:
            status["current"] = plan.label(item)
            try:
                produced, gone = plan.remove(item)
                entries.extend(produced)
                if gone:
                    status["deleted"] += 1
            except DockerError as exc:
                # One item must not stop the rest, e.g. something started
                # using it since the candidates were listed.
                logger.warning("Prune: could not remove %s: %s", status["current"], exc)
                status["failed"] += 1
                if len(status["errors"]) < _MAX_ERRORS:
                    status["errors"].append(f"{status['current']}: {exc.detail or exc.error}")
            status["done"] += 1
        status.update({"phase": "finishing", "current": None})
        result, reclaimed = plan.finish(entries)
        status["reclaimed"] = reclaimed
        return result

    def _finish(self, error: Optional[str]) -> None:
        self._status.update({"active": False, "phase": None, "current": None, "error": error})
        self._lock.release()