socks-manager/engine/instances.py

"""多实例管理器：每个实例是一个独立的 asyncio 服务器。"""
import asyncio
import logging
import threading
from dataclasses import dataclass, field

from engine.server import ConnectionHandler

log = logging.getLogger("socks.instances")


@dataclass
class InstanceConfig:
    """单个实例的配置快照（运行时使用）。"""
    instance_id: int
    name: str
    listen_host: str
    listen_port: int
    timeout: int
    auth_method: str
    bandwidth_down: float
    bandwidth_up: float
    max_concurrent: int

    def match(self, new_name, new_host, new_port):
        return (self.name == new_name and
                self.listen_host == new_host and
                self.listen_port == new_port)


class Socks5Server:
    """单个 SOCKS5 服务器的 asyncio 事件循环。"""

    def __init__(self, config: InstanceConfig, user_service):
        self.config = config
        self.user_service = user_service
        self.server = None
        self.loop = None
        self.thread = None
        self.running = False
        self._active_connections = 0
        self._lock = threading.Lock()

    def start(self):
        """启动服务器（在新线程中运行独立事件循环）。"""
        if self.running:
            return
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

        self._active_connections = 0
        self.running = True

        def run_loop():
            asyncio.set_event_loop(self.loop)
            self.server = self.loop.run_until_complete(
                self._create_server()
            )
            self.loop.run_forever()

        self.thread = threading.Thread(target=run_loop, daemon=True,
                                        name=f"socks5-{self.config.name}")
        self.thread.start()
        # 等 server 创建完成
        while self.server is None and self.running:
            time.sleep(0.05)
        if self.server:
            log.info("[%s] SOCKS5 服务器已启动: %s:%d",
                     self.config.name, self.config.listen_host, self.config.listen_port)

    def stop(self):
        """停止服务器。"""
        if not self.running:
            return
        self.running = False
        if self.server:
            self.server.close()
        if self.loop:
            try:
                self.loop.call_soon_threadsafe(self.loop.stop)
            except Exception:
                pass
        if self.thread:
            self.thread.join(timeout=5)
        log.info("[%s] SOCKS5 服务器已停止", self.config.name)

    async def _create_server(self):
        return await asyncio.start_server(
            self._accept_connection,
            self.config.listen_host, self.config.listen_port
        )

    async def _accept_connection(self, reader, writer):
        self._active_connections += 1
        handler = ConnectionHandler(
            reader, writer, self.config, self.user_service,
            on_close_cb=self._on_connection_close
        )
        await handler.handle()

    def _on_connection_close(self, handler):
        self._active_connections = max(0, self._active_connections - 1)

    @property
    def active_connections(self):
        with self._lock:
            return self._active_connections


import time  # noqa: E402


class InstanceManager:
    """管理所有 SOCKS5 实例。"""

    def __init__(self, user_service):
        self.user_service = user_service
        self._instances = {}       # name -> Socks5Server
        self._config_cache = {}    # name -> InstanceConfig
        self._lock = threading.Lock()

    def reload_config(self):
        """从数据库重新加载实例配置。"""
        from models import Instance
        from database import db
        new_configs = {}
        with db.app.app_context():
            for inst in Instance.query.filter_by(enabled=True).all():
                new_configs[inst.name] = InstanceConfig(
                    instance_id=inst.id,
                    name=inst.name,
                    listen_host=inst.listen_host,
                    listen_port=inst.listen_port,
                    timeout=inst.timeout,
                    auth_method=inst.auth_method,
                    bandwidth_down=inst.bandwidth_down,
                    bandwidth_up=inst.bandwidth_up,
                    max_concurrent=inst.max_concurrent,
                )
            self._config_cache = new_configs

    def sync_instances(self):
        """根据数据库配置同步实例运行状态。"""
        from models import Instance
        from database import db
        self.reload_config()

        current_names = set(self._config_cache.keys())
        running_names = set(self._instances.keys())

        # 启动缺失的实例
        for name in current_names - running_names:
            cfg = self._config_cache[name]
            srv = Socks5Server(cfg, self.user_service)
            srv.start()
            with self._lock:
                self._instances[name] = srv

        # 停止多余的实例
        for name in running_names - current_names:
            srv = self._instances.pop(name)
            srv.stop()

        # 更新实例运行状态到数据库
        with db.app.app_context():
            for inst in Instance.query.all():
                srv = self._instances.get(inst.name)
                inst.running = srv is not None and srv.running
                if srv:
                    inst.active_connections = srv.active_connections
                else:
                    inst.active_connections = 0
            db.session.commit()

    def start_instance(self, instance_id):
        """启动单个实例。"""
        from models import Instance
        from database import db
        with db.app.app_context():
            inst = Instance.query.get(instance_id)
            if not inst:
                return False
            cfg = InstanceConfig(
                instance_id=inst.id,
                name=inst.name,
                listen_host=inst.listen_host,
                listen_port=inst.listen_port,
                timeout=inst.timeout,
                auth_method=inst.auth_method,
                bandwidth_down=inst.bandwidth_down,
                bandwidth_up=inst.bandwidth_up,
                max_concurrent=inst.max_concurrent,
            )
            self._config_cache[inst.name] = cfg
            srv = Socks5Server(cfg, self.user_service)
            srv.start()
            with self._lock:
                self._instances[inst.name] = srv
            inst.running = True
            db.session.commit()
            return True

    def stop_instance(self, instance_id):
        """停止单个实例。"""
        from models import Instance
        from database import db
        with db.app.app_context():
            inst = Instance.query.get(instance_id)
            if not inst:
                return False
            srv = self._instances.pop(inst.name, None)
            if srv:
                srv.stop()
                del self._config_cache[inst.name]
            inst.running = False
            inst.active_connections = 0
            db.session.commit()
            return True

    def restart_instance(self, instance_id):
        """重启单个实例。"""
        self.stop_instance(instance_id)
        return self.start_instance(instance_id)

    def get_status(self):
        """获取所有实例状态。"""
        from models import Instance
        from database import db
        with db.app.app_context():
            result = []
            for inst in Instance.query.all():
                srv = self._instances.get(inst.name)
                result.append({
                    "id": inst.id,
                    "name": inst.name,
                    "host": inst.listen_host,
                    "port": inst.listen_port,
                    "enabled": inst.enabled,
                    "running": bool(srv and srv.running),
                    "active_connections": srv.active_connections if srv else 0,
                    "auth_method": inst.auth_method,
                    "timeout": inst.timeout,
                    "bandwidth_down": inst.bandwidth_down,
                    "bandwidth_up": inst.bandwidth_up,
                    "max_concurrent": inst.max_concurrent,
                    "notes": inst.notes,
                })
            return result