stackpilot/backend/services/stack_job_service.py

"""The long-running backup/restore behind a stack lock, and stopping it.

``stack_lock_service`` answers one question — is this stack busy — and
deliberately knows nothing about who is doing the work. That was enough until
a backup got stuck: the lock a backup takes lives for twelve hours (a large
volume legitimately takes that long), and for as long as it is held the stack
cannot be started, stopped, restarted or deployed. So a backup that hangs — an
unreachable NFS destination, a helper container reading a volume that never
ends — or one whose work died without unwinding takes the whole stack with it
until the TTL runs out.

This module is the other half of the lock: the operation registers itself
here, and a second request can ask it to stop.

Cancellation is cooperative, and on purpose:

* ``cancel`` sets a flag and force-removes the helper containers the job is
  streaming through, which is what actually ends the long wait — the stream
  breaks, the worker thread returns, and the backup unwinds through its own
  ``finally``. That is what restarts the containers the backup stopped and
  releases the lock. Killing the task instead would leave both undone.
* Nothing is cancelled from the outside. A restore is not interruptible at all
  (it is halfway through writing a stack's data back), so it registers itself
  to be *seen*, not stopped.

When even that does not help — a job whose process is long gone, leaving only
the lock row behind — the escape hatch is releasing the lock itself
(:func:`stack_lock_service.release`), which the unlock endpoint does with
``force``.
"""
from __future__ import annotations

import logging
import threading
from contextlib import contextmanager
from contextvars import ContextVar
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Optional

from docker_client import get_long_client, safe_call

logger = logging.getLogger("stackpilot.stack_job")

#: Actions that can be stopped halfway without leaving a mess behind. A backup
#: only reads; a restore is mid-way through writing a stack's data back and
#: must be allowed to finish.
CANCELLABLE = {"backup"}

#: Actions that always register here while they run. A lock held by one of
#: these with no job behind it is therefore a leftover — the run is over, only
#: its row survived — and releasing it stops nothing. Compose operations
#: (start, update, git-deploy) are not in this set: they run as a subprocess
#: that outlives the request, so a lock of theirs may well still mean work.
TRACKED = {"backup", "restore"}


class JobCancelled(Exception):
    """Raised inside a job that has been asked to stop."""

    def __init__(self, stack_id: str, action: str, by: str = ""):
        self.stack_id = stack_id
        self.action = action
        self.by = by
        who = f" by {by}" if by else ""
        super().__init__(f"{action} of '{stack_id}' was cancelled{who}")


def _now() -> datetime:
    return datetime.now(timezone.utc)


@dataclass
class Job:
    """One backup or restore, while it runs."""

    stack_id: str
    action: str
    owner: str = ""
    started_at: datetime = field(default_factory=_now)
    #: What it is doing right now ("volume media-data", "writing the archive"),
    #: so a stack that looks stuck can be told from one that is just slow.
    phase: str = ""
    cancelled_by: Optional[str] = None
    _stop: threading.Event = field(default_factory=threading.Event)
    #: Helper containers this job streams through. Removing them is what ends
    #: a read that would otherwise run for hours.
    _helpers: set[str] = field(default_factory=set)

    @property
    def cancellable(self) -> bool:
        return self.action in CANCELLABLE

    @property
    def cancelled(self) -> bool:
        return self._stop.is_set()

    def to_dict(self) -> dict:
        return {
            "stack_id": self.stack_id,
            "action": self.action,
            "owner": self.owner,
            "started_at": self.started_at.isoformat(),
            "phase": self.phase,
            "cancellable": self.cancellable,
            "cancelled": self.cancelled,
            "cancelled_by": self.cancelled_by,
        }


_jobs: dict[str, Job] = {}
_guard = threading.Lock()
#: The job of the current task. ``asyncio.to_thread`` copies the context, so
#: the worker threads a backup fans out into see their own job too.
_current: ContextVar[Optional[Job]] = ContextVar("stackpilot_job", default=None)


@contextmanager
def track(stack_id: str, action: str, owner: str = ""):
    """Register the operation for the duration of the block.

    Always deregisters: a job that is gone from here is one nobody can be
    waiting on, so a leftover entry would make a free stack look busy.
    """
    job = Job(stack_id=stack_id, action=action, owner=owner)
    with _guard:
        _jobs[stack_id] = job
    token = _current.set(job)
    try:
        yield job
    finally:
        _current.reset(token)
        with _guard:
            if _jobs.get(stack_id) is job:
                del _jobs[stack_id]


def current() -> Optional[Job]:
    return _current.get()


def get(stack_id: str) -> Optional[Job]:
    with _guard:
        return _jobs.get(stack_id)


def active() -> list[dict]:
    with _guard:
        return [job.to_dict() for job in _jobs.values()]


def phase(text: str) -> None:
    """Record what the current job is doing. No-op outside a job."""
    job = _current.get()
    if job is not None:
        job.phase = text


def cancelled() -> bool:
    job = _current.get()
    return job is not None and job.cancelled


def check() -> None:
    """Raise :class:`JobCancelled` if the current job was asked to stop.

    Called at every point a backup can safely give up — between the files and
    volumes it archives, not in the middle of one.
    """
    job = _current.get()
    if job is not None and job.cancelled:
        raise JobCancelled(job.stack_id, job.action, job.cancelled_by or "")


# --------------------------------------------------------------------------- #
# Helper containers
# --------------------------------------------------------------------------- #


def register_helper(container_id: str) -> None:
    """Note a helper container the current job reads or writes through."""
    job = _current.get()
    if job is not None:
        with _guard:
            job._helpers.add(container_id)


def drop_helper(container_id: str) -> None:
    job = _current.get()
    if job is not None:
        with _guard:
            job._helpers.discard(container_id)


def _kill_helpers(job: Job) -> int:
    with _guard:
        ids = list(job._helpers)
    if not ids:
        return 0
    killed = 0
    client = get_long_client()
    for container_id in ids:
        try:
            container = safe_call(client.containers.get, container_id)
            safe_call(container.remove, force=True)
            killed += 1
        except Exception as exc:  # noqa: BLE001 - already gone is a fine outcome
            logger.debug("Helper %s of '%s' could not be removed: %s", container_id[:12], job.stack_id, exc)
    return killed


def cancel(stack_id: str, by: str = "") -> Optional[Job]:
    """Ask the job on ``stack_id`` to stop. Returns it, or None if there is none.

    Returns as soon as the job has been told; the job itself still has to
    unwind — restart what it stopped, drop the lock — which is what the caller
    waits for afterwards.
    """
    job = get(stack_id)
    if job is None:
        return None
    if not job.cancellable:
        return job
    job.cancelled_by = by or "unknown"
    job._stop.set()
    killed = _kill_helpers(job)
    logger.info(
        "Cancelling the %s of '%s' (requested by %s, %d helper container(s) removed)",
        job.action, stack_id, job.cancelled_by, killed,
    )
    return job