stackpilot/backend/services/stack_lock_service.py

"""One compose operation per stack at a time.

``docker compose`` does no locking. Two ``update`` calls against the same
project — two open browser tabs, or the auto-update pass landing on a stack
somebody just clicked — both run ``pull`` and then ``up -d``, and race each
other recreating the same containers.

There *was* a busy flag in ``compose_service``, but it only ever fed the status
column: no lifecycle handler consulted it before acting. This module is the
actual guard, and it lives in the database so it holds across workers and
across a restart. ``compose_service.compute_status`` therefore reports only
what the containers say; callers overlay the lock to show "updating".
"""
from __future__ import annotations

import logging
from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
from typing import Optional

from sqlalchemy.exc import IntegrityError
from sqlmodel import Session, delete, select

from models.runtime_state import StackLock

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

#: Long enough to outlast the slowest legitimate operation (compose commands
#: time out at 600s, a full pull of a large stack can chain several), short
#: enough that a lock orphaned by a killed worker clears itself within an hour.
DEFAULT_TTL = timedelta(minutes=30)


class StackBusy(Exception):
    """The stack is already running an 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 _now() -> datetime:
    return datetime.now(timezone.utc)


def _aware(value: Optional[datetime]) -> Optional[datetime]:
    """SQLite hands datetimes back naive; compare them as UTC."""
    if value is not None and value.tzinfo is None:
        return value.replace(tzinfo=timezone.utc)
    return value


def acquire(
    session: Session,
    stack_id: str,
    action: str,
    owner: str = "",
    ttl: timedelta = DEFAULT_TTL,
) -> None:
    """Take the lock for ``stack_id`` or raise :class:`StackBusy`.

    An expired lock is taken over — that is the recovery path for a worker that
    died mid-deploy, which would otherwise leave the stack unusable.
    """
    now = _now()
    existing = session.get(StackLock, stack_id)
    if existing is not None:
        if (_aware(existing.expires_at) or now) > now:
            raise StackBusy(stack_id, existing.action)
        logger.warning(
            "Taking over an expired %s lock on '%s' (held by %r since %s)",
            existing.action, stack_id, existing.owner, existing.acquired_at,
        )
        session.delete(existing)
        session.commit()

    session.add(
        StackLock(
            stack_id=stack_id,
            action=action,
            owner=owner,
            acquired_at=now,
            expires_at=now + ttl,
        )
    )
    try:
        session.commit()
    except IntegrityError as exc:
        # Another worker inserted between our check and our commit. The primary
        # key is what actually makes this safe; the read above is only there to
        # give a useful error and to clear stale rows.
        session.rollback()
        raise StackBusy(stack_id, action) from exc


def release(session: Session, stack_id: str) -> None:
    """Drop the lock. Safe to call when it is not held."""
    existing = session.get(StackLock, stack_id)
    if existing is not None:
        session.delete(existing)
        session.commit()


@contextmanager
def hold(session: Session, stack_id: str, action: str, owner: str = ""):
    """Hold the lock for the duration of the block.

    Raises :class:`StackBusy` if somebody else has it. Always releases, so a
    failed deploy does not leave the stack locked.
    """
    acquire(session, stack_id, action, owner)
    try:
        yield
    finally:
        try:
            release(session, stack_id)
        except Exception:  # noqa: BLE001 - never mask the original error
            logger.exception("Failed to release the lock on '%s'", stack_id)


def active(session: Session) -> dict[str, str]:
    """``{stack_id: action}`` for every lock still in force.

    One query for the whole stacks list, rather than a lookup per row.
    """
    now = _now()
    return {
        lock.stack_id: lock.action
        for lock in session.exec(select(StackLock)).all()
        if (_aware(lock.expires_at) or now) > now
    }


def info(session: Session, stack_id: str) -> Optional[dict]:
    """The lock in force on ``stack_id``, as plain data — or None.

    What the UI needs to explain a stack that will not budge: which operation
    holds it, since when, and when it would give up on its own.
    """
    lock = session.get(StackLock, stack_id)
    if lock is None:
        return None
    now = _now()
    expires = _aware(lock.expires_at) or now
    if expires <= now:
        return None
    acquired = _aware(lock.acquired_at) or now
    return {
        "action": lock.action,
        "owner": lock.owner,
        "acquired_at": acquired.isoformat(),
        "expires_at": expires.isoformat(),
        "held_for": int((now - acquired).total_seconds()),
    }


def clear_all(session: Session) -> list[str]:
    """Drop every lock. Returns the stack ids that were locked.

    Called at startup: a lock only means something while the operation holding
    it runs, and nothing of ours is in flight yet (StackPilot runs a single
    uvicorn worker — see the Dockerfile's CMD). Before this, a backup whose
    process died left its stack unusable for the twelve hours of its TTL.
    """
    locks = session.exec(select(StackLock)).all()
    ids = [lock.stack_id for lock in locks]
    for lock in locks:
        session.delete(lock)
    session.commit()
    return ids


def is_busy(session: Session, stack_id: str) -> bool:
    lock = session.get(StackLock, stack_id)
    return lock is not None and (_aware(lock.expires_at) or _now()) > _now()


def prune_expired(session: Session) -> int:
    """Drop locks that have timed out. Called at startup and by the scheduler."""
    result = session.exec(delete(StackLock).where(StackLock.expires_at < _now()))
    session.commit()
    return result.rowcount or 0