Nexus/server/infrastructure/ssh/asyncssh_pool.py

"""Nexus — asyncssh Connection Pool (Reference Counting Mode)

ADR-002: BFF mode with asyncssh embedded (single process, full Python).
ADR-004: Reference counting connection pool — reuse connections, reduce handshakes.

W1: Extracted from ssh_direct.py, provides asyncssh-based connection pooling.
W2: Replaces paramiko pool.py with asyncssh throughout the codebase.
"""

import asyncio
import logging
import time
from dataclasses import dataclass, field
from typing import Optional, Dict

import asyncssh
from server.domain.models import Server
from server.infrastructure.database.crypto import decrypt_value

logger = logging.getLogger("nexus.asyncssh_pool")


@dataclass
class PooledConnection:
    """A tracked asyncssh connection with reference counting"""
    conn: asyncssh.SSHClientConnection
    ref_count: int = 0
    created_at: float = field(default_factory=time.monotonic)
    last_used: float = field(default_factory=time.monotonic)
    server_id: int = 0

    @property
    def is_idle(self) -> bool:
        return self.ref_count == 0

    @property
    def idle_seconds(self) -> float:
        return time.monotonic() - self.last_used if self.is_idle else 0


class AsyncSSHPool:
    """asyncssh connection pool with reference counting (ADR-004)

    Features:
    - Reference counting: multiple operations can share one SSH connection
    - Auto-cleanup: idle connections closed after IDLE_TIMEOUT
    - Per-server pooling: one connection per server at a time
    - Thread-safe: uses asyncio.Lock for concurrent access
    """

    # Pool configuration
    MAX_CONNECTIONS = 100       # Global max connections
    IDLE_TIMEOUT = 300          # 5 minutes idle → close
    CONNECT_TIMEOUT = 10        # 10 seconds to establish connection
    CLEANUP_INTERVAL = 60       # Check for idle connections every 60s

    def __init__(self):
        self._pool: Dict[int, PooledConnection] = {}  # server_id → connection
        self._lock = asyncio.Lock()
        self._cleanup_task: Optional[asyncio.Task] = None

    async def start(self):
        """Start the connection pool and cleanup task"""
        self._cleanup_task = asyncio.create_task(self._cleanup_loop(), name="asyncssh_cleanup")
        logger.info(f"AsyncSSH pool started (max={self.MAX_CONNECTIONS}, idle_timeout={self.IDLE_TIMEOUT}s)")

    async def stop(self):
        """Stop the pool and close all connections"""
        if self._cleanup_task:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass

        async with self._lock:
            for server_id, pooled in self._pool.items():
                try:
                    pooled.conn.close()
                except Exception:
                    pass
            self._pool.clear()
        logger.info("AsyncSSH pool stopped — all connections closed")

    async def acquire(self, server: Server) -> asyncssh.SSHClientConnection:
        """Acquire a connection to a server (reference counting).

        If an idle connection exists for this server, reuse it.
        Otherwise, create a new connection.
        SSH handshake is done outside the lock to avoid blocking other pool operations.
        """
        async with self._lock:
            # Check for existing idle connection
            pooled = self._pool.get(server.id)
            if pooled and not pooled.conn.is_closed():
                pooled.ref_count += 1
                pooled.last_used = time.monotonic()
                logger.debug(f"Reused SSH connection: server={server.id}, refs={pooled.ref_count}")
                return pooled.conn

            # Need to create — evict if at capacity, then release lock for handshake
            if len(self._pool) >= self.MAX_CONNECTIONS:
                await self._evict_one()

        # Create connection outside lock (handshake can take seconds)
        conn = await self._create_connection(server)

        async with self._lock:
            # Re-check: another coroutine may have created one while we were unlocked
            pooled = self._pool.get(server.id)
            if pooled and not pooled.conn.is_closed():
                # Close our new connection, reuse the existing one
                try:
                    conn.close()
                except Exception:
                    pass
                pooled.ref_count += 1
                pooled.last_used = time.monotonic()
                return pooled.conn

            pooled = PooledConnection(
                conn=conn,
                ref_count=1,
                server_id=server.id,
            )
            self._pool[server.id] = pooled
            logger.debug(f"New SSH connection: server={server.id}, pool_size={len(self._pool)}")
            return conn

    async def release(self, server_id: int):
        """Release a connection (decrement reference count).

        If ref_count reaches 0, the connection stays in pool for reuse.
        It will be cleaned up by the idle timeout task.
        """
        async with self._lock:
            pooled = self._pool.get(server_id)
            if pooled:
                pooled.ref_count = max(0, pooled.ref_count - 1)
                pooled.last_used = time.monotonic()
                logger.debug(f"Released SSH connection: server={server_id}, refs={pooled.ref_count}")

    async def close_connection(self, server_id: int):
        """Force-close a connection to a specific server"""
        async with self._lock:
            pooled = self._pool.pop(server_id, None)
            if pooled:
                try:
                    pooled.conn.close()
                except Exception:
                    pass
                logger.debug(f"Force-closed SSH connection: server={server_id}")

    async def _create_connection(self, server: Server) -> asyncssh.SSHClientConnection:
        """Create a new asyncssh connection to a server"""
        connect_kwargs = {
            "host": server.domain,
            "port": server.port or 22,
            "username": server.username or "root",
            "known_hosts": None,  # Skip host key verification (like paramiko AutoAddPolicy)
            "connect_timeout": self.CONNECT_TIMEOUT,
        }

        if server.auth_method == "password" and server.password:
            password = decrypt_value(server.password)
            connect_kwargs["password"] = password
        elif server.ssh_key_configured and server.ssh_key_private:
            key_content = decrypt_value(server.ssh_key_private)
            connect_kwargs["client_keys"] = [asyncssh.import_private_key(key_content)]
        elif server.ssh_key_path:
            connect_kwargs["client_keys"] = [server.ssh_key_path]
        else:
            raise ValueError(f"No valid SSH credentials for server {server.id} ({server.name})")

        try:
            conn = await asyncssh.connect(**connect_kwargs)
            return conn
        except asyncssh.Error as e:
            raise ConnectionError(f"SSH connection failed to {server.domain}:{server.port}: {e}")

    async def _evict_one(self):
        """Evict the oldest idle connection from the pool"""
        oldest_idle = None
        oldest_time = float("inf")

        for server_id, pooled in self._pool.items():
            if pooled.is_idle and pooled.last_used < oldest_time:
                oldest_idle = server_id
                oldest_time = pooled.last_used

        if oldest_idle is not None:
            pooled = self._pool.pop(oldest_idle)
            try:
                pooled.conn.close()
            except Exception:
                pass
            logger.debug(f"Evicted idle SSH connection: server={oldest_idle}")

    async def _cleanup_loop(self):
        """Periodically clean up idle connections"""
        while True:
            await asyncio.sleep(self.CLEANUP_INTERVAL)
            try:
                async with self._lock:
                    to_remove = []
                    for server_id, pooled in self._pool.items():
                        if pooled.is_idle and pooled.idle_seconds > self.IDLE_TIMEOUT:
                            to_remove.append(server_id)
                        elif not pooled.is_idle and pooled.conn.is_closed():
                            # Connection died while in use — remove from pool
                            to_remove.append(server_id)

                    for server_id in to_remove:
                        pooled = self._pool.pop(server_id)
                        try:
                            pooled.conn.close()
                        except Exception:
                            pass

                    if to_remove:
                        logger.info(f"SSH pool cleanup: removed {len(to_remove)} idle/dead connections")
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"SSH pool cleanup error: {e}")

    @property
    def stats(self) -> dict:
        """Get pool statistics"""
        total = len(self._pool)
        active = sum(1 for p in self._pool.values() if not p.is_idle)
        idle = total - active
        return {"total": total, "active": active, "idle": idle}


# Global pool instance
ssh_pool = AsyncSSHPool()


async def exec_ssh_command(server: Server, command: str, timeout: int = 30) -> dict:
    """Execute a command via SSH using the asyncssh connection pool.

    W2: Replaces paramiko-based exec_command from pool.py.
    Returns structured result dict with stdout/stderr/exit_code.
    """
    try:
        conn = await ssh_pool.acquire(server)
        try:
            result = await asyncio.wait_for(
                conn.run(command, timeout=timeout),
                timeout=timeout + 5,
            )
            return {
                "status": "success" if result.exit_status == 0 else "failed",
                "stdout": result.stdout[:10000],
                "stderr": result.stderr[:10000],
                "exit_code": result.exit_status,
            }
        finally:
            await ssh_pool.release(server.id)
    except asyncio.TimeoutError:
        return {
            "status": "timeout",
            "stdout": "",
            "stderr": f"Command timed out after {timeout}s",
            "exit_code": -1,
        }
    except ConnectionError as e:
        return {
            "status": "connection_error",
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1,
        }
    except Exception as e:
        return {
            "status": "error",
            "stdout": "",
            "stderr": str(e),
            "exit_code": -1,
        }