Your Name
437 lines
16 KiB
Python
437 lines
16 KiB
Python
"""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 secrets
|
|
import shlex
|
|
import time
|
|
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass, field
|
|
from typing import AsyncIterator, Dict, Optional
|
|
|
|
import asyncssh
|
|
from server.domain.models import Server
|
|
from server.utils.files_elevation import FilesElevation
|
|
from server.utils.posix_paths import posix_dirname
|
|
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: # noqa: S110 — best-effort cleanup
|
|
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
|
|
if len(self._pool) >= self.MAX_CONNECTIONS:
|
|
evicted = await self._evict_one()
|
|
if not evicted:
|
|
# All connections are busy — cannot exceed MAX_CONNECTIONS
|
|
raise ConnectionError(
|
|
f"SSH pool at capacity ({self.MAX_CONNECTIONS}) and all connections busy. "
|
|
f"Try again later."
|
|
)
|
|
|
|
# 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: # noqa: S110 — best-effort cleanup
|
|
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: # noqa: S110 — best-effort cleanup
|
|
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"""
|
|
from server.config import settings
|
|
|
|
connect_kwargs = {
|
|
"host": server.domain,
|
|
"port": server.port or 22,
|
|
"username": server.username or "root",
|
|
"connect_timeout": self.CONNECT_TIMEOUT,
|
|
}
|
|
|
|
# Host key verification: respect SSH_STRICT_HOST_CHECKING setting
|
|
# Default is "true" — reject unknown hosts for security
|
|
if settings.SSH_STRICT_HOST_CHECKING.lower() in ("false", "0", "no"):
|
|
connect_kwargs["known_hosts"] = None # Skip verification (insecure)
|
|
# If strict mode, asyncssh will use default known_hosts behavior
|
|
# (rejects unknown hosts — caller should configure ~/.ssh/known_hosts)
|
|
|
|
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}") from e
|
|
|
|
async def _evict_one(self) -> bool:
|
|
"""Evict the oldest idle connection from the pool.
|
|
|
|
Returns True if an idle connection was evicted, False if all are busy.
|
|
"""
|
|
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: # noqa: S110 — best-effort cleanup
|
|
pass
|
|
logger.debug(f"Evicted idle SSH connection: server={oldest_idle}")
|
|
return True
|
|
return False
|
|
|
|
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: # noqa: S110 — best-effort cleanup
|
|
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()
|
|
|
|
|
|
@asynccontextmanager
|
|
async def sftp_session(conn: asyncssh.SSHClientConnection) -> AsyncIterator[asyncssh.SFTPClient]:
|
|
"""Open an SFTP client on a pooled SSH connection (asyncssh 2.x: start_sftp_client)."""
|
|
async with await conn.start_sftp_client() as sftp:
|
|
yield sftp
|
|
|
|
|
|
async def sftp_write_bytes(sftp: asyncssh.SFTPClient, remote_path: str, data: bytes) -> None:
|
|
"""Write in-memory bytes to a remote file (asyncssh put() only accepts local paths)."""
|
|
async with sftp.open(remote_path, "wb") as remote_file:
|
|
await remote_file.write(data)
|
|
|
|
|
|
def _sftp_permission_denied(exc: Exception) -> bool:
|
|
if isinstance(exc, asyncssh.PermissionDenied):
|
|
return True
|
|
if isinstance(exc, asyncssh.SFTPError):
|
|
# SSH_FX_PERMISSION_DENIED = 3
|
|
if getattr(exc, "code", None) == 3:
|
|
return True
|
|
return "permission denied" in str(exc).lower()
|
|
|
|
|
|
async def _run_tee_stdin(
|
|
conn: asyncssh.SSHClientConnection,
|
|
command: str,
|
|
data: bytes,
|
|
) -> asyncssh.SSHCompletedProcess:
|
|
"""Feed raw bytes to stdin; default SSH session encoding expects str, not bytes."""
|
|
return await conn.run(command, input=data, encoding=None, check=False)
|
|
|
|
|
|
async def _ensure_remote_parent_dir(
|
|
conn: asyncssh.SSHClientConnection,
|
|
remote_path: str,
|
|
*,
|
|
use_sudo: bool,
|
|
) -> None:
|
|
parent = posix_dirname(remote_path)
|
|
if not parent or parent == "/":
|
|
return
|
|
prefix = "sudo -n " if use_sudo else ""
|
|
await conn.run(f"{prefix}mkdir -p {shlex.quote(parent)}", check=False)
|
|
|
|
|
|
async def _ssh_write_bytes_via_tee(
|
|
conn: asyncssh.SSHClientConnection,
|
|
remote_path: str,
|
|
data: bytes,
|
|
*,
|
|
use_sudo: bool,
|
|
) -> None:
|
|
"""Shell write fallback when SFTP cannot open the file (e.g. file not writable but directory is)."""
|
|
prefix = "sudo -n " if use_sudo else ""
|
|
await _ensure_remote_parent_dir(conn, remote_path, use_sudo=use_sudo)
|
|
|
|
direct = await _run_tee_stdin(
|
|
conn,
|
|
f"{prefix}tee {shlex.quote(remote_path)}",
|
|
data,
|
|
)
|
|
if direct.exit_status == 0:
|
|
return
|
|
|
|
tmp_path = f"/tmp/nexus-write-{secrets.token_hex(12)}"
|
|
staged = await _run_tee_stdin(
|
|
conn,
|
|
f"{prefix}tee {shlex.quote(tmp_path)}",
|
|
data,
|
|
)
|
|
if staged.exit_status != 0:
|
|
stderr = (staged.stderr or staged.stdout or direct.stderr or "")[:500]
|
|
raise PermissionError(stderr or "tee 写入失败")
|
|
|
|
moved = await conn.run(
|
|
f"{prefix}mv -f {shlex.quote(tmp_path)} {shlex.quote(remote_path)}",
|
|
check=False,
|
|
)
|
|
if moved.exit_status != 0:
|
|
await conn.run(f"{prefix}rm -f {shlex.quote(tmp_path)}", check=False)
|
|
stderr = (moved.stderr or "")[:500]
|
|
raise PermissionError(stderr or "无法覆盖目标文件（目录可能不可写）")
|
|
|
|
|
|
def _sudo_attempt_order(elevation: FilesElevation | None) -> tuple[bool, ...]:
|
|
mode = elevation if elevation is not None else FilesElevation.AUTO
|
|
if mode == FilesElevation.OFF:
|
|
return (False,)
|
|
if mode == FilesElevation.ALWAYS:
|
|
return (True,)
|
|
return (False, True)
|
|
|
|
|
|
async def remote_write_bytes(
|
|
conn: asyncssh.SSHClientConnection,
|
|
remote_path: str,
|
|
data: bytes,
|
|
*,
|
|
elevation: FilesElevation | None = None,
|
|
) -> None:
|
|
"""Write bytes to a remote path: SFTP first, then SSH tee / tmp+mv (incl. optional sudo -n)."""
|
|
eff = elevation if elevation is not None else FilesElevation.AUTO
|
|
skip_sftp = eff == FilesElevation.ALWAYS
|
|
|
|
if not skip_sftp:
|
|
try:
|
|
async with await conn.start_sftp_client() as sftp:
|
|
await sftp_write_bytes(sftp, remote_path, data)
|
|
return
|
|
except (asyncssh.PermissionDenied, asyncssh.SFTPError) as exc:
|
|
if not _sftp_permission_denied(exc):
|
|
raise
|
|
|
|
last_error: Exception | None = None
|
|
for use_sudo in _sudo_attempt_order(eff):
|
|
try:
|
|
await _ssh_write_bytes_via_tee(conn, remote_path, data, use_sudo=use_sudo)
|
|
return
|
|
except PermissionError as exc:
|
|
last_error = exc
|
|
|
|
reason = str(last_error) if last_error else "Permission denied"
|
|
raise asyncssh.SFTPError(3, reason)
|
|
|
|
|
|
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.
|
|
"""
|
|
conn = None
|
|
try:
|
|
conn = await ssh_pool.acquire(server)
|
|
result = await asyncio.wait_for(
|
|
conn.run(command, timeout=timeout),
|
|
timeout=timeout + 5,
|
|
)
|
|
await ssh_pool.release(server.id)
|
|
return {
|
|
"status": "success" if result.exit_status == 0 else "failed",
|
|
"stdout": result.stdout[:10000],
|
|
"stderr": result.stderr[:10000],
|
|
"exit_code": result.exit_status,
|
|
}
|
|
except asyncio.TimeoutError:
|
|
# Timed-out connection may be in an unknown state — evict so it is NOT
|
|
# returned to the pool and potentially reused for a future command.
|
|
if conn is not None:
|
|
await ssh_pool.close_connection(server.id)
|
|
return {
|
|
"status": "timeout",
|
|
"stdout": "",
|
|
"stderr": f"Command timed out after {timeout}s",
|
|
"exit_code": -1,
|
|
}
|
|
except ConnectionError as e:
|
|
if conn is not None:
|
|
await ssh_pool.close_connection(server.id)
|
|
return {
|
|
"status": "connection_error",
|
|
"stdout": "",
|
|
"stderr": str(e),
|
|
"exit_code": -1,
|
|
}
|
|
except Exception as e:
|
|
if conn is not None:
|
|
await ssh_pool.release(server.id)
|
|
return {
|
|
"status": "error",
|
|
"stdout": "",
|
|
"stderr": str(e),
|
|
"exit_code": -1,
|
|
} |