locking.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""Connection and state machine for GPU Memory Service.

This module handles:
- Connection: Represents an active client connection
- GMSLocalFSM: Explicit state transitions with validated permissions

State Diagram:

    EMPTY ──RW_CONNECT──► RW ──RW_COMMIT──► COMMITTED
      ▲                    │                   │
      │                    │                   │
      └───RW_ABORT─────────┘                   │
                                               ▼
    COMMITTED ◄──RO_DISCONNECT (last)── RO ◄──RO_CONNECT
                      │                  ▲
                      │                  │
                      └──RO_CONNECT──────┘
                      └──RO_DISCONNECT───┘ (not last)
"""

from __future__ import annotations

import asyncio
import logging
import socket
from dataclasses import dataclass, field
from typing import Callable, Optional, Set

from gpu_memory_service.common.types import (
    RO_ALLOWED,
    RW_ALLOWED,
    RW_REQUIRED,
    GrantedLockType,
    ServerState,
    StateEvent,
)

logger = logging.getLogger(__name__)


# =============================================================================
# Connection
# =============================================================================


@dataclass(eq=False)
class Connection:
    """Represents an active connection.

    The existence of Connection objects IS the state - we don't track
    sessions separately. When a Connection is removed, the lock is released.

    Note: eq=False disables auto-generated __eq__ so we can use default
    object identity for equality and add __hash__ for use in sets.
    """

    reader: asyncio.StreamReader
    writer: asyncio.StreamWriter
    mode: GrantedLockType
    session_id: str
    recv_buffer: bytearray = field(default_factory=bytearray)

    def __hash__(self) -> int:
        """Hash based on session_id (immutable identifier)."""
        return hash(self.session_id)

    @property
    def raw_socket(self) -> socket.socket:
        """Get underlying socket for FD passing."""
        return self.writer.get_extra_info("socket")

    async def close(self) -> None:
        """Close the connection."""
        self.writer.close()
        try:
            await self.writer.wait_closed()
        except Exception:
            pass


# =============================================================================
# State Machine
# =============================================================================


class InvalidTransition(Exception):
    """Raised when an invalid state transition is attempted."""

    pass


class OperationNotAllowed(Exception):
    """Raised when an operation is not allowed in the current state/mode."""

    pass


@dataclass(frozen=True)
class Transition:
    """A valid state transition.

    Attributes:
        from_states: Set of states this transition can originate from
        event: The event that triggers this transition
        to_state: The resulting state (or None if conditional)
        condition: Optional condition function for conditional transitions
    """

    from_states: frozenset[ServerState]
    event: StateEvent
    to_state: Optional[ServerState]
    condition: Optional[str] = None  # Name of condition method


# Transition table - the single source of truth for valid state transitions
TRANSITIONS: list[Transition] = [
    # From EMPTY or COMMITTED: RW can connect
    # Writer acquires exclusive lock
    Transition(
        from_states=frozenset({ServerState.EMPTY, ServerState.COMMITTED}),
        event=StateEvent.RW_CONNECT,
        to_state=ServerState.RW,
    ),
    # From RW: commit publishes and transitions to COMMITTED
    # Writer publishes and releases lock
    Transition(
        from_states=frozenset({ServerState.RW}),
        event=StateEvent.RW_COMMIT,
        to_state=ServerState.COMMITTED,
    ),
    # From RW: abort (disconnect without commit) transitions to EMPTY
    # Writer aborts, state invalidated
    Transition(
        from_states=frozenset({ServerState.RW}),
        event=StateEvent.RW_ABORT,
        to_state=ServerState.EMPTY,
    ),
    # From COMMITTED or RO: RO can connect
    # Reader acquires shared lock
    Transition(
        from_states=frozenset({ServerState.COMMITTED, ServerState.RO}),
        event=StateEvent.RO_CONNECT,
        to_state=ServerState.RO,
    ),
    # From RO: reader disconnect (not last) stays in RO
    # Reader leaves, others remain
    Transition(
        from_states=frozenset({ServerState.RO}),
        event=StateEvent.RO_DISCONNECT,
        to_state=ServerState.RO,
        condition="has_remaining_readers",
    ),
    # From RO: last reader disconnect transitions to COMMITTED
    # Last reader leaves
    Transition(
        from_states=frozenset({ServerState.RO}),
        event=StateEvent.RO_DISCONNECT,
        to_state=ServerState.COMMITTED,
        condition="is_last_reader",
    ),
]


@dataclass
class TransitionRecord:
    """Record of a state transition for debugging/auditing."""

    from_state: ServerState
    event: StateEvent
    to_state: ServerState
    session_id: Optional[str] = None


class GMSLocalFSM:
    """Explicit state machine for GPU Memory Service.

    State is DERIVED from actual connection objects:
    - _rw_conn: The active RW connection (or None)
    - _ro_conns: Set of active RO connections
    - _committed: Whether allocations have been committed

    All state mutations happen through explicit transitions.
    """

    def __init__(self, on_rw_abort: Optional[Callable[[], None]] = None):
        """Initialize the state machine.

        Args:
            on_rw_abort: Callback invoked when RW aborts (for cleanup)
        """
        # Connection state - THIS IS THE SOURCE OF TRUTH
        self._rw_conn: Optional[Connection] = None
        self._ro_conns: Set[Connection] = set()
        self._committed: bool = False

        # Callback for RW abort cleanup
        self._on_rw_abort = on_rw_abort

        # Transition history for debugging
        self._transition_log: list[TransitionRecord] = []

    # ==================== State Properties ====================

    @property
    def state(self) -> ServerState:
        """Derive current state from connection objects."""
        if self._rw_conn is not None:
            return ServerState.RW
        if len(self._ro_conns) > 0:
            return ServerState.RO
        if self._committed:
            return ServerState.COMMITTED
        return ServerState.EMPTY

    @property
    def rw_conn(self) -> Optional[Connection]:
        """The active RW connection, if any."""
        return self._rw_conn

    @property
    def ro_conns(self) -> Set[Connection]:
        """Set of active RO connections."""
        return self._ro_conns

    @property
    def ro_count(self) -> int:
        """Number of active RO connections."""
        return len(self._ro_conns)

    @property
    def committed(self) -> bool:
        """Whether allocations have been committed."""
        return self._committed

    @property
    def transition_log(self) -> list[TransitionRecord]:
        """History of state transitions."""
        return self._transition_log

    # ==================== Transition Conditions ====================

    def _has_remaining_readers(self, conn: Connection) -> bool:
        """Check if there are readers remaining after removing conn."""
        return len(self._ro_conns) > 1 or conn not in self._ro_conns

    def _is_last_reader(self, conn: Connection) -> bool:
        """Check if conn is the last reader."""
        return len(self._ro_conns) == 1 and conn in self._ro_conns

    def _check_condition(self, condition: Optional[str], conn: Connection) -> bool:
        """Evaluate a named condition."""
        if condition is None:
            return True
        if condition == "has_remaining_readers":
            return self._has_remaining_readers(conn)
        if condition == "is_last_reader":
            return self._is_last_reader(conn)
        raise ValueError(f"Unknown condition: {condition}")

    # ==================== State Transitions ====================

    def _find_transition(
        self, from_state: ServerState, event: StateEvent, conn: Connection
    ) -> Optional[Transition]:
        """Find the applicable transition for the given event."""
        for t in TRANSITIONS:
            if from_state not in t.from_states:
                continue
            if t.event != event:
                continue
            if not self._check_condition(t.condition, conn):
                continue
            return t
        return None

    def _apply_event(self, event: StateEvent, conn: Connection) -> None:
        """Mutate internal state based on event."""
        match event:
            case StateEvent.RW_CONNECT:
                self._rw_conn = conn
                self._committed = False  # Invalidate on RW connect
            case StateEvent.RW_COMMIT:
                self._committed = True
                self._rw_conn = None
            case StateEvent.RW_ABORT:
                self._rw_conn = None
                if self._on_rw_abort:
                    self._on_rw_abort()
            case StateEvent.RO_CONNECT:
                self._ro_conns.add(conn)
            case StateEvent.RO_DISCONNECT:
                self._ro_conns.discard(conn)

    def transition(self, event: StateEvent, conn: Connection) -> ServerState:
        """Execute a state transition.

        Args:
            event: The triggering event
            conn: The connection involved in the transition

        Returns:
            The new state after the transition

        Raises:
            InvalidTransition: If the transition is not valid from current state
        """
        from_state = self.state
        session_id = conn.session_id if conn else None

        # Find valid transition
        trans = self._find_transition(from_state, event, conn)
        if trans is None:
            raise InvalidTransition(
                f"No transition for {event.name} from state {from_state.name} "
                f"(session={session_id})"
            )

        # Apply the transition
        self._apply_event(event, conn)
        to_state = self.state

        # Validate we ended up in expected state
        if trans.to_state is not None and to_state != trans.to_state:
            raise InvalidTransition(
                f"Transition mismatch: expected {trans.to_state.name}, "
                f"got {to_state.name}"
            )

        # Record transition
        record = TransitionRecord(
            from_state,
            event,
            to_state,
            session_id=session_id,
        )
        self._transition_log.append(record)

        logger.info(
            f"State transition: {from_state.name} --{event.name}--> {to_state.name} "
            f"(session={session_id})"
        )

        return to_state

    # ==================== Operation Permissions ====================

    def check_operation(self, msg_type: type, conn: Connection) -> None:
        """Check if a request type is allowed for the given connection.

        Args:
            msg_type: The request message type (e.g., AllocateRequest)
            conn: The connection attempting the operation

        Raises:
            OperationNotAllowed: If the operation is not permitted
        """
        current_state = self.state

        # Determine allowed operations based on state
        if current_state == ServerState.RW:
            allowed = RW_ALLOWED
        elif current_state == ServerState.RO:
            allowed = RO_ALLOWED
        else:
            allowed = frozenset()  # EMPTY and COMMITTED have no connections

        if msg_type not in allowed:
            raise OperationNotAllowed(
                f"{msg_type.__name__} not allowed in state {current_state.name}"
            )

        # Check connection mode
        if msg_type in RW_REQUIRED and conn.mode != GrantedLockType.RW:
            raise OperationNotAllowed(
                f"{msg_type.__name__} requires RW connection, "
                f"but connection is {conn.mode.value}"
            )

    # ==================== Lock Acquisition Predicates ====================

    def can_acquire_rw(self) -> bool:
        """Check if RW lock can be acquired now.

        RW can only be acquired if:
        - No current RW holder
        - No RO holders

        Note: This allows RW from COMMITTED state (for explicit reload).
        For rw_or_ro mode, callers should also check `committed` to prefer RO.
        """
        return self._rw_conn is None and len(self._ro_conns) == 0

    def can_acquire_ro(self, waiting_writers: int) -> bool:
        """Check if RO lock can be acquired now.

        Args:
            waiting_writers: Number of writers waiting for the lock
        """
        return self._rw_conn is None and waiting_writers == 0 and self._committed