NODEDC_MISSION_CORE/src/k1link/sessions/active.py

from __future__ import annotations

import importlib
import os
import stat
from dataclasses import dataclass
from pathlib import Path
from typing import Any, cast

ACTIVE_SESSION_MARKER = ".current_session"


class ActiveSessionLeaseError(RuntimeError):
    """The evidence root already has an active writer or cannot be leased."""


@dataclass(slots=True)
class ActiveSessionLease:
    """Cross-process lease that keeps an in-progress session out of replay.

    The marker is intentionally written before the session directory is
    created. Discovery therefore observes either no candidate yet or a
    candidate protected by an already locked marker. A process crash releases
    the OS lock; startup recovery can then remove the stale marker while
    preserving the interrupted evidence directory.
    """

    sessions_root: Path
    session_root: Path
    _descriptor: int
    _marker_identity: tuple[int, int]
    _released: bool = False

    @classmethod
    def acquire(cls, sessions_root: Path, session_root: Path) -> ActiveSessionLease:
        root = sessions_root.expanduser().resolve()
        root.mkdir(mode=0o700, parents=True, exist_ok=True)
        target = session_root.expanduser().absolute()
        if target.parent.resolve() != root or not target.name:
            raise ActiveSessionLeaseError("active session must be a direct child of evidence root")

        marker = root / ACTIVE_SESSION_MARKER
        flags = os.O_RDWR | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
        try:
            descriptor = os.open(marker, flags, 0o600)
        except FileExistsError as exc:
            raise ActiveSessionLeaseError(
                "another observation session owns the evidence root"
            ) from exc
        try:
            payload = f"{target.name}\n".encode()
            os.write(descriptor, payload)
            os.fsync(descriptor)
            _lock_descriptor(descriptor, blocking=False)
            _fsync_directory(root)
            metadata = os.fstat(descriptor)
            return cls(
                sessions_root=root,
                session_root=target,
                _descriptor=descriptor,
                _marker_identity=(metadata.st_dev, metadata.st_ino),
            )
        except BaseException:
            try:
                marker.unlink(missing_ok=True)
                _fsync_directory(root)
            finally:
                os.close(descriptor)
            raise

    def release(self) -> None:
        if self._released:
            return
        marker = self.sessions_root / ACTIVE_SESSION_MARKER
        try:
            try:
                metadata = marker.lstat()
            except FileNotFoundError:
                metadata = None
            if metadata is not None and (metadata.st_dev, metadata.st_ino) == self._marker_identity:
                marker.unlink()
                _fsync_directory(self.sessions_root)
        finally:
            _unlock_descriptor(self._descriptor)
            os.close(self._descriptor)
            self._released = True

    def __enter__(self) -> ActiveSessionLease:
        return self

    def __exit__(self, *_: object) -> None:
        self.release()


def recover_stale_active_session_marker(sessions_root: Path) -> bool:
    """Remove only an unlocked marker left by a terminated writer."""

    root = sessions_root.expanduser().resolve()
    marker = root / ACTIVE_SESSION_MARKER
    try:
        marker_stat = marker.lstat()
    except FileNotFoundError:
        return False
    except OSError:
        return False
    if not stat.S_ISREG(marker_stat.st_mode) or marker_stat.st_size > 4096:
        return False
    try:
        descriptor = os.open(marker, os.O_RDWR | getattr(os, "O_NOFOLLOW", 0))
    except OSError:
        return False
    locked = False
    try:
        try:
            _lock_descriptor(descriptor, blocking=False)
            locked = True
        except OSError:
            return False
        opened = os.fstat(descriptor)
        try:
            current = marker.lstat()
        except OSError:
            return False
        if (opened.st_dev, opened.st_ino) != (current.st_dev, current.st_ino):
            return False
        marker.unlink()
        _fsync_directory(root)
        return True
    finally:
        if locked:
            _unlock_descriptor(descriptor)
        os.close(descriptor)


def _lock_descriptor(descriptor: int, *, blocking: bool) -> None:
    if os.name == "nt":
        msvcrt = cast(Any, importlib.import_module("msvcrt"))

        os.lseek(descriptor, 0, os.SEEK_SET)
        mode = msvcrt.LK_LOCK if blocking else msvcrt.LK_NBLCK
        msvcrt.locking(descriptor, mode, 1)
        return
    import fcntl

    operation = fcntl.LOCK_EX | (0 if blocking else fcntl.LOCK_NB)
    fcntl.flock(descriptor, operation)


def _unlock_descriptor(descriptor: int) -> None:
    try:
        if os.name == "nt":
            msvcrt = cast(Any, importlib.import_module("msvcrt"))

            os.lseek(descriptor, 0, os.SEEK_SET)
            msvcrt.locking(descriptor, msvcrt.LK_UNLCK, 1)
            return
        import fcntl

        fcntl.flock(descriptor, fcntl.LOCK_UN)
    except OSError:
        return


def _fsync_directory(path: Path) -> None:
    try:
        descriptor = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
    except OSError:
        return
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)