script-build-and-framework-tmp/script_build_host/src/script_build_host/infrastructure/ownership.py

"""Single-host mission ownership and database fencing.

The file descriptor lock remains held for the lifetime of an owner.  Database
epochs make every later command independently reject a stale process.
"""

from __future__ import annotations

import fcntl
import os
from collections.abc import Awaitable, Callable
from contextlib import AbstractAsyncContextManager
from datetime import UTC, datetime
from pathlib import Path
from types import TracebackType
from typing import Any, cast
from uuid import uuid4

from sqlalchemy import select, update
from sqlalchemy.engine import CursorResult
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from script_build_host.domain.errors import (
    BuildNotFound,
    MissionAlreadyOwned,
    MissionFencingTokenStale,
    MissionStopRequested,
)
from script_build_host.domain.records import BuildStatus, MissionOwnerToken
from script_build_host.infrastructure.legacy_tables import (
    script_build_record,
    script_build_runtime_record,
)
from script_build_host.infrastructure.tables import mission_binding_table


class OwnerLease(AbstractAsyncContextManager[MissionOwnerToken]):
    def __init__(
        self,
        sessions: async_sessionmaker[AsyncSession],
        lock_root: Path,
        *,
        owner_instance_id: str | None = None,
    ) -> None:
        self._sessions = sessions
        self._lock_root = lock_root.resolve()
        self._owner_instance_id = owner_instance_id or str(uuid4())
        self._fd: int | None = None
        self._token: MissionOwnerToken | None = None

    async def acquire(self, script_build_id: int) -> MissionOwnerToken:
        if self._fd is not None:
            if self._token is None or self._token.script_build_id != script_build_id:
                raise MissionAlreadyOwned()
            return self._token
        self._lock_root.mkdir(parents=True, exist_ok=True)
        path = self._lock_root / f"script-build-{script_build_id}.owner.lock"
        fd = os.open(path, os.O_RDWR | os.O_CREAT, 0o600)
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError as exc:
            os.close(fd)
            raise MissionAlreadyOwned() from exc
        try:
            async with self._sessions() as session, session.begin():
                row = (
                    (
                        await session.execute(
                            select(mission_binding_table)
                            .where(mission_binding_table.c.script_build_id == script_build_id)
                            .with_for_update()
                        )
                    )
                    .mappings()
                    .one_or_none()
                )
                if row is None:
                    raise BuildNotFound()
                epoch = int(row["owner_epoch"] or 0) + 1
                now = datetime.now(UTC)
                await session.execute(
                    update(mission_binding_table)
                    .where(mission_binding_table.c.id == row["id"])
                    .values(
                        owner_epoch=epoch,
                        owner_instance_id=self._owner_instance_id,
                        owner_acquired_at=now,
                        updated_at=now,
                    )
                )
                token = MissionOwnerToken(
                    script_build_id=script_build_id,
                    root_trace_id=str(row["root_trace_id"]),
                    owner_epoch=epoch,
                    stop_epoch=int(row["stop_epoch"] or 0),
                    owner_instance_id=self._owner_instance_id,
                )
        except BaseException:
            fcntl.flock(fd, fcntl.LOCK_UN)
            os.close(fd)
            raise
        self._fd = fd
        self._token = token
        return token

    async def release(self) -> None:
        fd, token = self._fd, self._token
        if fd is None:
            return
        try:
            if token is not None:
                async with self._sessions() as session, session.begin():
                    await session.execute(
                        update(mission_binding_table)
                        .where(
                            mission_binding_table.c.script_build_id == token.script_build_id,
                            mission_binding_table.c.owner_epoch == token.owner_epoch,
                            mission_binding_table.c.owner_instance_id == token.owner_instance_id,
                        )
                        .values(
                            owner_instance_id=None,
                            owner_acquired_at=None,
                            updated_at=datetime.now(UTC),
                        )
                    )
        finally:
            fcntl.flock(fd, fcntl.LOCK_UN)
            os.close(fd)
            self._fd = None
            self._token = None

    async def __aenter__(self) -> MissionOwnerToken:
        if self._token is None:
            raise RuntimeError("OwnerLease.acquire() must be called before entering the lease")
        return self._token

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        traceback: TracebackType | None,
    ) -> None:
        await self.release()


class FencedCommandGate:
    """Validate one owner token while holding binding/build row locks."""

    def __init__(self, sessions: async_sessionmaker[AsyncSession]) -> None:
        self._sessions = sessions
        bind = sessions.kw.get("bind")
        self._runtime_record = (
            script_build_runtime_record
            if getattr(getattr(bind, "dialect", None), "name", None) == "mysql"
            else script_build_record
        )

    async def verify(self, token: MissionOwnerToken) -> None:
        async with self._sessions() as session, session.begin():
            await self.verify_in_session(session, token)

    async def execute(
        self,
        token: MissionOwnerToken,
        mutation: Callable[[AsyncSession], Awaitable[Any]],
    ) -> Any:
        async with self._sessions() as session, session.begin():
            await self.verify_in_session(session, token)
            return await mutation(session)

    async def compare_and_set_input_snapshot(
        self,
        token: MissionOwnerToken,
        *,
        expected_snapshot_id: int,
        new_snapshot_id: int,
    ) -> None:
        """Move the binding to one immutable snapshot under the owner fence."""

        async with self._sessions() as session, session.begin():
            await self.verify_in_session(session, token)
            result = cast(
                CursorResult[Any],
                await session.execute(
                    update(mission_binding_table)
                    .where(
                        mission_binding_table.c.script_build_id == token.script_build_id,
                        mission_binding_table.c.input_snapshot_id == expected_snapshot_id,
                    )
                    .values(input_snapshot_id=new_snapshot_id, updated_at=datetime.now(UTC))
                ),
            )
            if result.rowcount != 1:
                current = await session.scalar(
                    select(mission_binding_table.c.input_snapshot_id).where(
                        mission_binding_table.c.script_build_id == token.script_build_id
                    )
                )
                if current != new_snapshot_id:
                    raise MissionFencingTokenStale()

    async def begin_phase(self, token: MissionOwnerToken) -> None:
        """Project the legacy build to running without a stop-overwrite window."""

        await self._update_runtime_state(
            token,
            expected_statuses=(BuildStatus.PARTIAL.value, BuildStatus.RUNNING.value),
            status=BuildStatus.RUNNING.value,
            end_time=None,
            error_message=None,
            summary=None,
        )

    async def set_checkpoint(
        self,
        token: MissionOwnerToken,
        *,
        checkpoint_code: str,
        summary: str,
    ) -> None:
        """Persist the Phase 3 checkpoint in the same transaction as fencing."""

        await self._update_runtime_state(
            token,
            expected_statuses=(BuildStatus.RUNNING.value,),
            idempotent_checkpoint=checkpoint_code[:255],
            status=BuildStatus.PARTIAL.value,
            error_message=checkpoint_code[:255],
            summary=summary[:2000],
            reson_trace_id=token.root_trace_id[:200],
            end_time=datetime.now(UTC),
        )

    async def _update_runtime_state(
        self,
        token: MissionOwnerToken,
        *,
        expected_statuses: tuple[str, ...],
        idempotent_checkpoint: str | None = None,
        **values: Any,
    ) -> None:
        async with self._sessions() as session, session.begin():
            await self.verify_in_session(session, token)
            result = cast(
                CursorResult[Any],
                await session.execute(
                    update(self._runtime_record)
                    .where(
                        self._runtime_record.c.id == token.script_build_id,
                        self._runtime_record.c.is_deleted.is_(False),
                        self._runtime_record.c.status.in_(expected_statuses),
                    )
                    .values(**values)
                ),
            )
            if result.rowcount != 1:
                current = (
                    (
                        await session.execute(
                            select(
                                script_build_record.c.status,
                                script_build_record.c.error_message,
                            ).where(script_build_record.c.id == token.script_build_id)
                        )
                    )
                    .mappings()
                    .one_or_none()
                )
                if current is None:
                    raise BuildNotFound()
                if (
                    idempotent_checkpoint is not None
                    and current["status"] == BuildStatus.PARTIAL.value
                    and current["error_message"] == idempotent_checkpoint
                ):
                    return
                raise MissionFencingTokenStale()

    async def verify_in_session(
        self,
        session: AsyncSession,
        token: MissionOwnerToken,
        *,
        allow_success: bool = False,
    ) -> None:
        binding = (
            (
                await session.execute(
                    select(mission_binding_table)
                    .where(mission_binding_table.c.script_build_id == token.script_build_id)
                    .with_for_update()
                )
            )
            .mappings()
            .one_or_none()
        )
        if binding is None:
            raise BuildNotFound()
        if (
            str(binding["root_trace_id"]) != token.root_trace_id
            or int(binding["owner_epoch"] or 0) != token.owner_epoch
            or binding["owner_instance_id"] != token.owner_instance_id
        ):
            raise MissionFencingTokenStale()
        if int(binding["stop_epoch"] or 0) != token.stop_epoch:
            raise MissionStopRequested()
        status = await session.scalar(
            select(script_build_record.c.status)
            .where(script_build_record.c.id == token.script_build_id)
            .with_for_update()
        )
        if status is None:
            raise BuildNotFound()
        if status in {BuildStatus.STOPPING.value, BuildStatus.STOPPED.value}:
            raise MissionStopRequested()
        if status == BuildStatus.FAILED.value:
            raise MissionFencingTokenStale()
        if status == BuildStatus.SUCCESS.value and not allow_success:
            raise MissionFencingTokenStale()

    async def request_stop(self, script_build_id: int) -> int:
        """Persist stop intent without acquiring the process OwnerLease."""

        async with self._sessions() as session, session.begin():
            binding = (
                (
                    await session.execute(
                        select(mission_binding_table)
                        .where(mission_binding_table.c.script_build_id == script_build_id)
                        .with_for_update()
                    )
                )
                .mappings()
                .one_or_none()
            )
            if binding is None:
                raise BuildNotFound()
            status = await session.scalar(
                select(script_build_record.c.status)
                .where(script_build_record.c.id == script_build_id)
                .with_for_update()
            )
            if status is None:
                raise BuildNotFound()
            if status in {BuildStatus.SUCCESS.value, BuildStatus.STOPPED.value}:
                return int(binding["stop_epoch"] or 0)
            stop_epoch = int(binding["stop_epoch"] or 0) + 1
            await session.execute(
                update(mission_binding_table)
                .where(mission_binding_table.c.id == binding["id"])
                .values(stop_epoch=stop_epoch, updated_at=datetime.now(UTC))
            )
            await session.execute(
                update(self._runtime_record)
                .where(self._runtime_record.c.id == script_build_id)
                .values(status=BuildStatus.STOPPING.value, end_time=None)
            )
            return stop_epoch


__all__ = ["FencedCommandGate", "OwnerLease"]