fbuild 1.2.8__py3-none-any.whl

fbuild/daemon/async_client.py

@@ -0,0 +1,531 @@
+"""
+Async Client Connection Manager for fbuild daemon.
+
+This module provides classes for managing asynchronous client connections,
+including connection tracking, heartbeat mechanism, and connection lifecycle.
+
+Features:
+- Track connected clients by unique client_id (UUID string)
+- Heartbeat mechanism to detect dead clients
+- Track client metadata: pid, connect_time, last_heartbeat, attached resources
+- Cleanup on client disconnect (release locks, detach from serial sessions)
+- Support for registering cleanup callbacks
+"""
+
+import logging
+import threading
+import time
+import uuid
+from dataclasses import asdict, dataclass, field
+from typing import Any, Callable
+
+# Default heartbeat timeout: clients not sending heartbeat in this time are considered dead
+DEFAULT_HEARTBEAT_TIMEOUT = 30.0
+
+# Cleanup interval for checking dead clients
+CLEANUP_INTERVAL = 10.0
+
+
+@dataclass
+class ClientInfo:
+    """Information about a connected client.
+
+    Attributes:
+        client_id: Unique identifier for the client (UUID string)
+        pid: Process ID of the client
+        connect_time: Unix timestamp when client connected
+        last_heartbeat: Unix timestamp of last heartbeat received
+        metadata: Additional client metadata (e.g., version, hostname)
+        attached_resources: Set of resource keys this client is attached to
+    """
+
+    client_id: str
+    pid: int
+    connect_time: float = field(default_factory=time.time)
+    last_heartbeat: float = field(default_factory=time.time)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    attached_resources: set[str] = field(default_factory=set)
+
+    def is_alive(self, timeout_seconds: float = DEFAULT_HEARTBEAT_TIMEOUT) -> bool:
+        """Check if client is still alive based on heartbeat timeout.
+
+        Args:
+            timeout_seconds: Maximum time since last heartbeat before considered dead.
+
+        Returns:
+            True if client is alive (heartbeat within timeout), False otherwise.
+        """
+        return (time.time() - self.last_heartbeat) <= timeout_seconds
+
+    def time_since_heartbeat(self) -> float:
+        """Get time in seconds since last heartbeat."""
+        return time.time() - self.last_heartbeat
+
+    def connection_duration(self) -> float:
+        """Get total connection duration in seconds."""
+        return time.time() - self.connect_time
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return {
+            "client_id": self.client_id,
+            "pid": self.pid,
+            "connect_time": self.connect_time,
+            "last_heartbeat": self.last_heartbeat,
+            "metadata": self.metadata,
+            "attached_resources": list(self.attached_resources),
+            "is_alive": self.is_alive(),
+            "time_since_heartbeat": self.time_since_heartbeat(),
+            "connection_duration": self.connection_duration(),
+        }
+
+
+@dataclass
+class ClientConnectMessage:
+    """Client -> Daemon: Connection request message.
+
+    Sent when a client connects to the daemon to register itself.
+
+    Attributes:
+        client_id: Unique identifier for the client (generated by client or daemon)
+        pid: Process ID of the client
+        metadata: Additional client metadata (e.g., version, hostname)
+        timestamp: Unix timestamp when message was created
+    """
+
+    client_id: str
+    pid: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+    timestamp: float = field(default_factory=time.time)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ClientConnectMessage":
+        """Create ClientConnectMessage from dictionary."""
+        return cls(
+            client_id=data["client_id"],
+            pid=data["pid"],
+            metadata=data.get("metadata", {}),
+            timestamp=data.get("timestamp", time.time()),
+        )
+
+
+@dataclass
+class ClientHeartbeatMessage:
+    """Client -> Daemon: Periodic heartbeat message.
+
+    Sent periodically by clients to indicate they are still alive.
+
+    Attributes:
+        client_id: Unique identifier for the client
+        timestamp: Unix timestamp when heartbeat was sent
+    """
+
+    client_id: str
+    timestamp: float = field(default_factory=time.time)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ClientHeartbeatMessage":
+        """Create ClientHeartbeatMessage from dictionary."""
+        return cls(
+            client_id=data["client_id"],
+            timestamp=data.get("timestamp", time.time()),
+        )
+
+
+@dataclass
+class ClientDisconnectMessage:
+    """Client -> Daemon: Graceful disconnect message.
+
+    Sent when a client is gracefully disconnecting.
+
+    Attributes:
+        client_id: Unique identifier for the client
+        reason: Optional reason for disconnection
+        timestamp: Unix timestamp when disconnect was initiated
+    """
+
+    client_id: str
+    reason: str | None = None
+    timestamp: float = field(default_factory=time.time)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for JSON serialization."""
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ClientDisconnectMessage":
+        """Create ClientDisconnectMessage from dictionary."""
+        return cls(
+            client_id=data["client_id"],
+            reason=data.get("reason"),
+            timestamp=data.get("timestamp", time.time()),
+        )
+
+
+class ClientConnectionManager:
+    """Manages client connections with heartbeat monitoring and cleanup.
+
+    This class provides a centralized manager for tracking client connections
+    to the daemon. It handles:
+    - Client registration and unregistration
+    - Heartbeat tracking for connection health
+    - Resource attachment tracking per client
+    - Automatic cleanup of dead clients
+    - Cleanup callbacks for resource release
+
+    Example:
+        >>> manager = ClientConnectionManager()
+        >>>
+        >>> # Register a new client
+        >>> client_id = manager.generate_client_id()
+        >>> client_info = manager.register_client(client_id, pid=12345)
+        >>>
+        >>> # Send periodic heartbeats
+        >>> manager.heartbeat(client_id)
+        >>>
+        >>> # Attach resources
+        >>> manager.attach_resource(client_id, "port:/dev/ttyUSB0")
+        >>>
+        >>> # Cleanup dead clients
+        >>> dead_clients = manager.cleanup_dead_clients()
+    """
+
+    def __init__(self) -> None:
+        """Initialize the ClientConnectionManager."""
+        self._lock = threading.Lock()
+        self._clients: dict[str, ClientInfo] = {}
+        self._cleanup_callbacks: list[Callable[[str], None]] = []
+
+    def generate_client_id(self) -> str:
+        """Generate a unique client ID using UUID4.
+
+        Returns:
+            A unique client ID string.
+        """
+        return str(uuid.uuid4())
+
+    def register_client(
+        self,
+        client_id: str,
+        pid: int,
+        metadata: dict[str, Any] | None = None,
+    ) -> ClientInfo:
+        """Register a new client connection.
+
+        If a client with the same ID already exists, it will be replaced
+        (after calling cleanup callbacks for the old client).
+
+        Args:
+            client_id: Unique identifier for the client.
+            pid: Process ID of the client.
+            metadata: Optional additional client metadata.
+
+        Returns:
+            ClientInfo object for the registered client.
+        """
+        with self._lock:
+            # If client already exists, unregister it first
+            if client_id in self._clients:
+                logging.warning(f"Client {client_id} already registered, replacing existing connection")
+                # Call cleanup callbacks outside the lock
+                self._call_cleanup_callbacks_unlocked(client_id)
+
+            current_time = time.time()
+            client_info = ClientInfo(
+                client_id=client_id,
+                pid=pid,
+                connect_time=current_time,
+                last_heartbeat=current_time,
+                metadata=metadata or {},
+                attached_resources=set(),
+            )
+            self._clients[client_id] = client_info
+            logging.info(f"Client registered: {client_id} (pid={pid})")
+            return client_info
+
+    def unregister_client(self, client_id: str) -> bool:
+        """Unregister a client connection.
+
+        This calls any registered cleanup callbacks before removing the client.
+
+        Args:
+            client_id: Unique identifier for the client.
+
+        Returns:
+            True if client was unregistered, False if client not found.
+        """
+        with self._lock:
+            if client_id not in self._clients:
+                logging.warning(f"Cannot unregister unknown client: {client_id}")
+                return False
+
+            client_info = self._clients[client_id]
+            resource_count = len(client_info.attached_resources)
+
+            # Call cleanup callbacks (this may release resources)
+            self._call_cleanup_callbacks_unlocked(client_id)
+
+            # Remove client from registry
+            del self._clients[client_id]
+            logging.info(f"Client unregistered: {client_id} " f"(had {resource_count} attached resources)")
+            return True
+
+    def heartbeat(self, client_id: str) -> bool:
+        """Update last heartbeat time for a client.
+
+        Args:
+            client_id: Unique identifier for the client.
+
+        Returns:
+            True if heartbeat was recorded, False if client not found.
+        """
+        with self._lock:
+            if client_id not in self._clients:
+                logging.debug(f"Heartbeat from unknown client: {client_id}")
+                return False
+
+            self._clients[client_id].last_heartbeat = time.time()
+            logging.debug(f"Heartbeat received from client: {client_id}")
+            return True
+
+    def get_client(self, client_id: str) -> ClientInfo | None:
+        """Get client information by ID.
+
+        Args:
+            client_id: Unique identifier for the client.
+
+        Returns:
+            ClientInfo if client exists, None otherwise.
+        """
+        with self._lock:
+            return self._clients.get(client_id)
+
+    def get_all_clients(self) -> dict[str, ClientInfo]:
+        """Get all registered clients.
+
+        Returns:
+            Dictionary mapping client_id to ClientInfo for all registered clients.
+        """
+        with self._lock:
+            # Return a copy to prevent external modification
+            return dict(self._clients)
+
+    def is_client_alive(
+        self,
+        client_id: str,
+        timeout_seconds: float = DEFAULT_HEARTBEAT_TIMEOUT,
+    ) -> bool:
+        """Check if a client is alive based on heartbeat timeout.
+
+        Args:
+            client_id: Unique identifier for the client.
+            timeout_seconds: Maximum time since last heartbeat before considered dead.
+
+        Returns:
+            True if client exists and is alive, False otherwise.
+        """
+        with self._lock:
+            client_info = self._clients.get(client_id)
+            if client_info is None:
+                return False
+            return client_info.is_alive(timeout_seconds)
+
+    def cleanup_dead_clients(
+        self,
+        timeout_seconds: float = DEFAULT_HEARTBEAT_TIMEOUT,
+    ) -> list[str]:
+        """Clean up clients that have not sent a heartbeat within the timeout.
+
+        This calls cleanup callbacks for each dead client before removing them.
+
+        Args:
+            timeout_seconds: Maximum time since last heartbeat before considered dead.
+
+        Returns:
+            List of client IDs that were cleaned up.
+        """
+        dead_clients: list[str] = []
+
+        with self._lock:
+            # Find dead clients
+            for client_id, client_info in list(self._clients.items()):
+                if not client_info.is_alive(timeout_seconds):
+                    dead_clients.append(client_id)
+                    logging.warning(f"Client {client_id} is dead " f"(no heartbeat for {client_info.time_since_heartbeat():.1f}s)")
+
+            # Clean up dead clients
+            for client_id in dead_clients:
+                # Call cleanup callbacks
+                self._call_cleanup_callbacks_unlocked(client_id)
+
+                # Remove from registry
+                del self._clients[client_id]
+                logging.info(f"Dead client cleaned up: {client_id}")
+
+        return dead_clients
+
+    def attach_resource(self, client_id: str, resource_key: str) -> bool:
+        """Attach a resource to a client for tracking.
+
+        This allows tracking which resources a client is using, so they
+        can be released when the client disconnects.
+
+        Args:
+            client_id: Unique identifier for the client.
+            resource_key: Key identifying the resource (e.g., "port:/dev/ttyUSB0").
+
+        Returns:
+            True if resource was attached, False if client not found.
+        """
+        with self._lock:
+            if client_id not in self._clients:
+                logging.warning(f"Cannot attach resource to unknown client: {client_id}")
+                return False
+
+            self._clients[client_id].attached_resources.add(resource_key)
+            logging.debug(f"Resource attached to client {client_id}: {resource_key}")
+            return True
+
+    def detach_resource(self, client_id: str, resource_key: str) -> bool:
+        """Detach a resource from a client.
+
+        Args:
+            client_id: Unique identifier for the client.
+            resource_key: Key identifying the resource.
+
+        Returns:
+            True if resource was detached, False if client or resource not found.
+        """
+        with self._lock:
+            if client_id not in self._clients:
+                logging.warning(f"Cannot detach resource from unknown client: {client_id}")
+                return False
+
+            resources = self._clients[client_id].attached_resources
+            if resource_key not in resources:
+                logging.warning(f"Resource not attached to client {client_id}: {resource_key}")
+                return False
+
+            resources.discard(resource_key)
+            logging.debug(f"Resource detached from client {client_id}: {resource_key}")
+            return True
+
+    def get_client_resources(self, client_id: str) -> set[str]:
+        """Get all resources attached to a client.
+
+        Args:
+            client_id: Unique identifier for the client.
+
+        Returns:
+            Set of resource keys attached to the client (empty set if not found).
+        """
+        with self._lock:
+            client_info = self._clients.get(client_id)
+            if client_info is None:
+                return set()
+            # Return a copy to prevent external modification
+            return set(client_info.attached_resources)
+
+    def register_cleanup_callback(
+        self,
+        callback: Callable[[str], None],
+    ) -> None:
+        """Register a callback to be called when a client disconnects.
+
+        The callback receives the client_id of the disconnecting client.
+        Callbacks are called in the order they were registered.
+
+        Args:
+            callback: Function to call with client_id when client disconnects.
+        """
+        with self._lock:
+            self._cleanup_callbacks.append(callback)
+            logging.debug(f"Cleanup callback registered " f"(total callbacks: {len(self._cleanup_callbacks)})")
+
+    def _call_cleanup_callbacks_unlocked(self, client_id: str) -> None:
+        """Call all cleanup callbacks for a client (must hold lock).
+
+        This is an internal method that calls cleanup callbacks while
+        the lock is held. Callbacks should be fast and non-blocking.
+
+        Args:
+            client_id: Unique identifier for the disconnecting client.
+        """
+        for callback in self._cleanup_callbacks:
+            try:
+                callback(client_id)
+            except KeyboardInterrupt:  # noqa: KBI002
+                raise
+            except Exception as e:
+                logging.error(f"Error in cleanup callback for client {client_id}: {e}")
+
+    def get_client_count(self) -> int:
+        """Get the number of registered clients.
+
+        Returns:
+            Number of currently registered clients.
+        """
+        with self._lock:
+            return len(self._clients)
+
+    def get_alive_client_count(
+        self,
+        timeout_seconds: float = DEFAULT_HEARTBEAT_TIMEOUT,
+    ) -> int:
+        """Get the number of alive clients.
+
+        Args:
+            timeout_seconds: Maximum time since last heartbeat before considered dead.
+
+        Returns:
+            Number of alive clients.
+        """
+        with self._lock:
+            return sum(1 for client in self._clients.values() if client.is_alive(timeout_seconds))
+
+    def get_status(self) -> dict[str, Any]:
+        """Get status information about the client manager.
+
+        Returns:
+            Dictionary with client manager status information.
+        """
+        with self._lock:
+            clients_info = {client_id: client.to_dict() for client_id, client in self._clients.items()}
+
+            return {
+                "total_clients": len(self._clients),
+                "alive_clients": sum(1 for client in self._clients.values() if client.is_alive()),
+                "dead_clients": sum(1 for client in self._clients.values() if not client.is_alive()),
+                "callback_count": len(self._cleanup_callbacks),
+                "clients": clients_info,
+            }
+
+    def clear_all_clients(self) -> int:
+        """Clear all clients (use with caution - for daemon restart).
+
+        This calls cleanup callbacks for each client before clearing.
+
+        Returns:
+            Number of clients cleared.
+        """
+        with self._lock:
+            count = len(self._clients)
+
+            # Call cleanup callbacks for each client
+            for client_id in list(self._clients.keys()):
+                self._call_cleanup_callbacks_unlocked(client_id)
+
+            self._clients.clear()
+
+            if count > 0:
+                logging.info(f"Cleared all {count} clients")
+
+            return count