stackpilot/backend/services/compose_run_service.py

"""Streamed compose runs (deploy console, update progress) that outlive their socket.

The deploy console and the stacks list's update bar watch ``compose up`` /
``compose pull && up`` over a WebSocket. That run used to *be* the socket
handler, which went wrong the moment somebody closed the modal or navigated
away halfway through a deploy:

* the stack lock was released in the handler's ``finally`` at once, while
  compose was still running — so the next deploy, update or backup could start
  on top of it;
* nobody read compose's output any more, so a chatty pull filled the pipe and
  compose blocked on it;
* the audit entry and the "stack started" notification went out with
  ``rc=None`` before anything had finished, as a success.

Here the run is its own task. It holds the lock and drains the output until
compose exits, then records the real exit code. A socket only *watches*: it can
leave at any time, and one that connects while the same operation is already in
flight on that stack is attached to it (backlog first, then live) instead of
being turned away or starting a second compose.
"""
from __future__ import annotations

import asyncio
import collections
import logging
from collections.abc import AsyncIterator, Awaitable, Callable
from dataclasses import dataclass, field
from datetime import timedelta
from typing import Optional

from sqlmodel import Session

from database import engine
from services import stack_lock_service

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

#: Lines kept for a socket that attaches to a run already in progress.
BACKLOG_LINES = 2000

#: A streamed run has no timeout of its own (a large pull legitimately takes a
#: while), so its lock may live longer than a REST call's. The unlock endpoint
#: is the way out of one that hangs.
LOCK_TTL = timedelta(hours=2)

#: (kind, payload): ("log", line) or ("done", returncode).
Item = tuple[str, object]


@dataclass
class ComposeRun:
    stack_id: str
    action: str
    owner: str
    returncode: Optional[int] = None
    finished: bool = False
    backlog: collections.deque = field(
        default_factory=lambda: collections.deque(maxlen=BACKLOG_LINES)
    )
    _subscribers: list[asyncio.Queue] = field(default_factory=list)

    def subscribe(self) -> asyncio.Queue:
        """A queue that gets the backlog, every later item, then ``None``."""
        queue: asyncio.Queue = asyncio.Queue()
        for line in self.backlog:
            queue.put_nowait(("log", line))
        if self.finished:
            queue.put_nowait(("done", self.returncode))
            queue.put_nowait(None)
        else:
            self._subscribers.append(queue)
        return queue

    def unsubscribe(self, queue: asyncio.Queue) -> None:
        if queue in self._subscribers:
            self._subscribers.remove(queue)

    def _publish(self, item: Optional[Item]) -> None:
        for queue in list(self._subscribers):
            queue.put_nowait(item)


_runs: dict[str, ComposeRun] = {}
#: asyncio keeps only weak references to tasks; these must not be collected
#: while compose is still running.
_tasks: set[asyncio.Task] = set()


class RunBusy(Exception):
    """The stack is held by a different operation."""

    def __init__(self, stack_id: str, action: str):
        self.stack_id = stack_id
        self.action = action
        super().__init__(f"Stack '{stack_id}' is busy: {action} in progress")


def get(stack_id: str) -> Optional[ComposeRun]:
    return _runs.get(stack_id)


def start(
    stack_id: str,
    action: str,
    owner: str,
    stream: Callable[[], AsyncIterator[Item]],
    on_finish: Optional[Callable[[Optional[int]], Awaitable[None]]] = None,
) -> tuple[ComposeRun, asyncio.Queue, bool]:
    """Start ``stream`` under the stack lock, or attach to the same run in flight.

    Returns ``(run, queue, attached)``. Raises :class:`RunBusy` when another
    operation holds the stack. Must be called from the event loop; the queue is
    subscribed before the task gets to run, so not a line is missed.
    """
    current = _runs.get(stack_id)
    if current is not None and not current.finished:
        if current.action != action:
            raise RunBusy(stack_id, current.action)
        return current, current.subscribe(), True

    with Session(engine) as session:
        try:
            stack_lock_service.acquire(session, stack_id, action, owner, ttl=LOCK_TTL)
        except stack_lock_service.StackBusy as exc:
            raise RunBusy(stack_id, exc.action) from exc

    run = ComposeRun(stack_id=stack_id, action=action, owner=owner)
    _runs[stack_id] = run
    queue = run.subscribe()
    task = asyncio.create_task(_drive(run, stream, on_finish), name=f"compose-{action}-{stack_id}")
    _tasks.add(task)
    task.add_done_callback(_tasks.discard)
    return run, queue, False


async def _drive(
    run: ComposeRun,
    stream: Callable[[], AsyncIterator[Item]],
    on_finish: Optional[Callable[[Optional[int]], Awaitable[None]]],
) -> None:
    rc: Optional[int] = None
    try:
        async for kind, payload in stream():
            if kind == "log":
                run.backlog.append(payload)
                run._publish(("log", payload))
            else:
                rc = payload  # type: ignore[assignment]
    except asyncio.CancelledError:
        raise
    except Exception as exc:  # noqa: BLE001 - reported to whoever is watching
        logger.warning("compose %s of '%s' failed: %s", run.action, run.stack_id, exc)
        line = f"error: {exc}"
        run.backlog.append(line)
        run._publish(("log", line))
        rc = rc if rc is not None else -1
    finally:
        run.returncode = rc
        run.finished = True
        try:
            with Session(engine) as session:
                stack_lock_service.release(session, run.stack_id)
        except Exception:  # noqa: BLE001 - never mask the outcome
            logger.exception("Failed to release the lock on '%s'", run.stack_id)
        run._publish(("done", rc))
        run._publish(None)
        if _runs.get(run.stack_id) is run:
            del _runs[run.stack_id]
    if on_finish is not None:
        try:
            await on_finish(rc)
        except Exception:  # noqa: BLE001 - audit/notify are best-effort
            logger.exception("Post-run hook for '%s' failed", run.stack_id)