stackpilot/backend/services/image_service.py

"""Image listing and pruning for the images router."""
from __future__ import annotations

import logging
import threading
from typing import Callable, Optional

from docker_client import DockerError, get_client, get_long_client, safe_call
from services import update_service

logger = logging.getLogger(__name__)

COMPOSE_PROJECT_LABEL = "com.docker.compose.project"


def list_images() -> list[dict]:
    """Return one row per image tag, annotated with using-stacks + update status."""
    try:
        client = get_client()
        images = safe_call(client.images.list)
        containers = safe_call(client.containers.list, all=True)
    except DockerError:
        return []

    # Map image ref -> stacks using it.
    usage: dict[str, set[str]] = {}
    for c in containers:
        ref = c.attrs.get("Config", {}).get("Image")
        stack = c.labels.get(COMPOSE_PROJECT_LABEL)
        if ref:
            usage.setdefault(ref, set())
            if stack:
                usage[ref].add(stack)

    cache = update_service.get_cache()
    result = []
    for img in images:
        tags = img.tags or []
        if not tags:
            continue
        for tag in tags:
            result.append(
                {
                    "id": img.short_id,
                    "tag": tag,
                    "size": img.attrs.get("Size", 0),
                    "created": img.attrs.get("Created"),
                    "stacks": sorted(usage.get(tag, set())),
                    "update": cache.get(tag),
                }
            )
    result.sort(key=lambda r: r["tag"])
    return result


# --------------------------------------------------------------------------- #
# Prune
# --------------------------------------------------------------------------- #

#: One prune at a time: a second pass over the same images would only collect
#: errors for everything the first one already removed.
_prune_lock = threading.Lock()
_pruning: dict = {
    "active": False,
    "run": 0,
    "all": False,
    "total": 0,
    "done": 0,
    "current": None,
    "deleted": 0,
    "failed": 0,
    "errors": [],
    "reclaimed": None,
    "error": None,
}
_MAX_ERRORS = 10


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


def prune_status() -> dict:
    return {**_pruning, "errors": list(_pruning["errors"])}


def _layers_size(client) -> Optional[int]:
    """Total size of all image layers on the host, or None if unavailable.

    Taken before and after the prune, the difference is the space actually
    freed. Summing the sizes of the deleted images would count base layers
    they share with the images that stay. ``type=image`` keeps the daemon from
    also walking every volume (API 1.42+; older daemons ignore it).
    """
    try:
        api = client.api
        response = api._get(api._url("/system/df"), params={"type": "image"})
        return api._result(response, True).get("LayersSize")
    except Exception as exc:  # noqa: BLE001 - the size is a nicety, not the job
        logger.warning("Could not read image disk usage: %s", exc)
        return None


def _candidates(client, all_unused: bool) -> list[dict]:
    """The images ``docker image prune [-a]`` would remove, largest first."""
    in_use = {c.get("ImageID") for c in safe_call(client.api.containers, all=True)}
    filters = None if all_unused else {"dangling": True}
    images = safe_call(client.api.images, filters=filters)
    unused = [img for img in images if img["Id"] not in in_use]
    return sorted(unused, key=lambda img: img.get("Size", 0), reverse=True)


def _label(image: dict) -> str:
    tags = [t for t in image.get("RepoTags") or [] if t != "<none>:<none>"]
    if tags:
        return tags[0] if len(tags) == 1 else f"{tags[0]} (+{len(tags) - 1})"
    return image["Id"].removeprefix("sha256:")[:12]


def _remove(client, image: dict) -> list[dict]:
    """Delete one image the way the prune would, returning Docker's entries.

    Tags are removed one by one rather than with force=True: the last one takes
    the image with it, and a container created since the candidates were listed
    still makes the daemon refuse instead of losing its image.
    """
    tags = [t for t in image.get("RepoTags") or [] if t != "<none>:<none>"]
    entries: list[dict] = []
    for ref in tags or [image["Id"]]:
        entries.extend(safe_call(client.api.remove_image, ref) or [])
    return entries


def _begin(all_unused: bool) -> tuple:
    client = get_long_client()
    candidates = _candidates(client, all_unused)
    _pruning.update(
        {
            "active": True,
            "run": _pruning["run"] + 1,
            "all": all_unused,
            "total": len(candidates),
            "done": 0,
            "current": None,
            "deleted": 0,
            "failed": 0,
            "errors": [],
            "reclaimed": None,
            "error": None,
        }
    )
    return client, candidates


def _execute(client, candidates: list[dict], all_unused: bool) -> dict:
    before = _layers_size(client)
    deleted: list[dict] = []
    for image in candidates:
        _pruning["current"] = _label(image)
        try:
            entries = _remove(client, image)
            deleted.extend(entries)
            if any("Deleted" in e for e in entries):
                _pruning["deleted"] += 1
        except DockerError as exc:
            # One image must not stop the rest, e.g. a container started from
            # it since the list was taken.
            logger.warning("Prune: could not remove %s: %s", _pruning["current"], exc)
            _pruning["failed"] += 1
            if len(_pruning["errors"]) < _MAX_ERRORS:
                _pruning["errors"].append(f"{_pruning['current']}: {exc.detail or exc.error}")
        _pruning["done"] += 1
    _pruning["current"] = None

    # Let the daemon sweep up whatever the per-image pass left, such as
    # untagged images referenced from several repositories.
    filters = {"dangling": False} if all_unused else {"dangling": True}
    sweep = safe_call(client.images.prune, filters=filters)
    deleted.extend(sweep.get("ImagesDeleted") or [])

    after = _layers_size(client)
    if before is not None and after is not None:
        reclaimed = max(before - after, 0)
    else:
        reclaimed = sweep.get("SpaceReclaimed", 0)
    _pruning["reclaimed"] = reclaimed
    return {"ImagesDeleted": deleted, "SpaceReclaimed": reclaimed}


def _finish(error: Optional[str] = None) -> None:
    _pruning.update({"active": False, "current": None, "error": error})
    _prune_lock.release()


def prune_images(all_unused: bool = False) -> dict:
    """Remove unused images and wait for the result.

    By default only dangling (untagged) images are removed; ``all_unused=True``
    removes every image not referenced by a container (``docker image prune
    -a``). Progress is visible through prune_status() while it runs.
    """
    if not _prune_lock.acquire(blocking=False):
        raise PruneBusy("An image prune is already running")
    error = None
    try:
        client, candidates = _begin(all_unused)
        return _execute(client, candidates, all_unused)
    except DockerError as exc:
        error = exc.detail or exc.error
        raise
    finally:
        _finish(error)


def start_prune(
    all_unused: bool = False, on_done: Optional[Callable[[dict], None]] = None
) -> dict:
    """Start a prune in the background and return its initial status.

    A prune removing gigabytes of layers can outlast the proxy's read timeout,
    so the UI starts it this way and polls prune_status(). ``on_done`` gets
    the same result prune_images() returns, from the worker thread.
    """
    if not _prune_lock.acquire(blocking=False):
        raise PruneBusy("An image prune is already running")
    try:
        client, candidates = _begin(all_unused)
    except BaseException:
        _prune_lock.release()
        raise

    def work() -> None:
        error = None
        result = None
        try:
            result = _execute(client, candidates, all_unused)
        except Exception as exc:  # noqa: BLE001 - reported through the status
            logger.warning("Image prune failed: %s", exc)
            error = exc.detail if isinstance(exc, DockerError) else str(exc)
        finally:
            _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("Image prune follow-up failed: %s", exc)

    threading.Thread(target=work, name="image-prune", daemon=True).start()
    return prune_status()