252 lines
8.4 KiB
Python
252 lines
8.4 KiB
Python
import asyncio
|
|
import logging
|
|
import socket
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Optional
|
|
|
|
try:
|
|
import nacos
|
|
except Exception:
|
|
nacos = None
|
|
|
|
from config import Config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class NacosConfig:
|
|
"""Nacos 配置"""
|
|
enabled: bool
|
|
server_addresses: str
|
|
namespace: str
|
|
group_name: str
|
|
cluster_name: str
|
|
username: Optional[str]
|
|
password: Optional[str]
|
|
heartbeat_interval: int
|
|
weight: float
|
|
ephemeral: bool
|
|
|
|
|
|
@dataclass
|
|
class ServiceConfig:
|
|
"""服务配置"""
|
|
service_name: str
|
|
host: str
|
|
port: int
|
|
ip: str
|
|
metadata: Dict[str, Any]
|
|
|
|
|
|
def _get_local_ip() -> str:
|
|
"""获取本地 IP 地址"""
|
|
try:
|
|
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
|
|
s.connect(("8.8.8.8", 80))
|
|
ip = s.getsockname()[0]
|
|
s.close()
|
|
return ip
|
|
except Exception as e:
|
|
logger.warning(f"获取本地 IP 失败，使用 127.0.0.1: {e}")
|
|
return "127.0.0.1"
|
|
|
|
|
|
def load_nacos_config() -> NacosConfig:
|
|
"""从 config.ini 读取 Nacos 配置"""
|
|
section = "nacos"
|
|
enabled = Config._config.getboolean(section, "enabled", fallback=False)
|
|
return NacosConfig(
|
|
enabled=enabled,
|
|
server_addresses=Config._config.get(section, "server", fallback="localhost:8848"),
|
|
namespace=Config._config.get(section, "namespace", fallback="public"),
|
|
group_name=Config._config.get(section, "group_name", fallback="DEFAULT_GROUP"),
|
|
cluster_name=Config._config.get(section, "cluster_name", fallback="DEFAULT"),
|
|
username=Config._config.get(section, "username", fallback="") or None,
|
|
password=Config._config.get(section, "password", fallback="") or None,
|
|
heartbeat_interval=Config._config.getint(section, "heartbeat_interval", fallback=5),
|
|
weight=Config._config.getfloat(section, "weight", fallback=1.0),
|
|
ephemeral=Config._config.getboolean(section, "ephemeral", fallback=True),
|
|
)
|
|
|
|
|
|
def load_service_config() -> ServiceConfig:
|
|
"""从 config.ini 读取服务配置"""
|
|
section = "app"
|
|
service_name = Config._config.get(section, "service_name", fallback="more-dots-api")
|
|
host = Config._config.get(section, "host", fallback="0.0.0.0")
|
|
port = Config._config.getint(section, "port", fallback=8000)
|
|
ip = host if host != "0.0.0.0" else _get_local_ip()
|
|
|
|
metadata = {
|
|
"version": Config._config.get(section, "version", fallback="1.0.0"),
|
|
"service_type": "fastapi",
|
|
"api_paths": "/health,/api/workflows,/api/workflows/stream,/nacos/status",
|
|
"streaming": "false",
|
|
"model_section": Config._config.get(section, "model_section", fallback=Config.DEFAULT_MODEL_SECTION),
|
|
}
|
|
|
|
extra_meta = Config.get_section("metadata")
|
|
if extra_meta:
|
|
metadata.update(extra_meta)
|
|
|
|
return ServiceConfig(
|
|
service_name=service_name,
|
|
host=host,
|
|
port=port,
|
|
ip=ip,
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
class NacosManager:
|
|
"""Nacos 服务注册与心跳管理"""
|
|
|
|
def __init__(self, nacos_config: NacosConfig, service_config: ServiceConfig):
|
|
self.nacos_config = nacos_config
|
|
self.service_config = service_config
|
|
self.client = None
|
|
self._heartbeat_task: Optional[asyncio.Task] = None
|
|
self._stop_event = asyncio.Event()
|
|
self.is_registered = False
|
|
|
|
def _init_client(self) -> bool:
|
|
"""初始化 Nacos 客户端"""
|
|
if nacos is None:
|
|
raise ImportError("未安装 nacos-sdk-python，请先安装依赖")
|
|
|
|
try:
|
|
self.client = nacos.NacosClient(
|
|
server_addresses=self.nacos_config.server_addresses,
|
|
namespace=self.nacos_config.namespace,
|
|
username=self.nacos_config.username,
|
|
password=self.nacos_config.password,
|
|
)
|
|
logger.info(f"Nacos 客户端初始化成功: {self.nacos_config.server_addresses}")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Nacos 客户端初始化失败: {e}")
|
|
return False
|
|
|
|
def register_service(self) -> bool:
|
|
"""注册服务到 Nacos"""
|
|
if not self.client and not self._init_client():
|
|
return False
|
|
|
|
try:
|
|
self.client.add_naming_instance(
|
|
service_name=self.service_config.service_name,
|
|
ip=self.service_config.ip,
|
|
port=self.service_config.port,
|
|
cluster_name=self.nacos_config.cluster_name,
|
|
group_name=self.nacos_config.group_name,
|
|
weight=self.nacos_config.weight,
|
|
metadata=self.service_config.metadata,
|
|
ephemeral=self.nacos_config.ephemeral,
|
|
)
|
|
self.is_registered = True
|
|
logger.info(
|
|
"✅ 服务注册成功: %s (%s:%s)",
|
|
self.service_config.service_name,
|
|
self.service_config.ip,
|
|
self.service_config.port,
|
|
)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"❌ 服务注册失败: {e}")
|
|
self.is_registered = False
|
|
return False
|
|
|
|
def deregister_service(self) -> bool:
|
|
"""从 Nacos 注销服务"""
|
|
if not self.client or not self.is_registered:
|
|
return True
|
|
|
|
try:
|
|
self.client.remove_naming_instance(
|
|
service_name=self.service_config.service_name,
|
|
ip=self.service_config.ip,
|
|
port=self.service_config.port,
|
|
cluster_name=self.nacos_config.cluster_name,
|
|
group_name=self.nacos_config.group_name,
|
|
)
|
|
self.is_registered = False
|
|
logger.info("✅ 服务注销成功: %s", self.service_config.service_name)
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"❌ 服务注销失败: {e}")
|
|
return False
|
|
|
|
def _send_heartbeat(self) -> None:
|
|
"""发送心跳"""
|
|
if not self.client or not self.is_registered:
|
|
return
|
|
|
|
self.client.send_heartbeat(
|
|
service_name=self.service_config.service_name,
|
|
ip=self.service_config.ip,
|
|
port=self.service_config.port,
|
|
cluster_name=self.nacos_config.cluster_name,
|
|
group_name=self.nacos_config.group_name,
|
|
)
|
|
|
|
async def _heartbeat_loop(self) -> None:
|
|
"""心跳循环"""
|
|
interval = max(1, self.nacos_config.heartbeat_interval)
|
|
self._stop_event.clear()
|
|
|
|
while not self._stop_event.is_set():
|
|
try:
|
|
self._send_heartbeat()
|
|
logger.debug("心跳发送成功: %s", self.service_config.service_name)
|
|
except Exception as e:
|
|
logger.warning(f"心跳发送失败: {e}")
|
|
# 尝试重新注册
|
|
try:
|
|
self.register_service()
|
|
except Exception as re:
|
|
logger.error(f"重新注册失败: {re}")
|
|
|
|
try:
|
|
await asyncio.wait_for(self._stop_event.wait(), timeout=interval)
|
|
except asyncio.TimeoutError:
|
|
continue
|
|
|
|
async def start(self) -> None:
|
|
"""启动注册与心跳"""
|
|
if not self.nacos_config.enabled:
|
|
logger.info("Nacos 未启用，跳过注册")
|
|
return
|
|
|
|
if self.register_service():
|
|
self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
|
|
logger.info("✅ Nacos 心跳任务已启动")
|
|
else:
|
|
logger.warning("⚠️ Nacos 注册失败，服务继续运行")
|
|
|
|
async def stop(self) -> None:
|
|
"""停止心跳并注销"""
|
|
self._stop_event.set()
|
|
|
|
if self._heartbeat_task:
|
|
self._heartbeat_task.cancel()
|
|
try:
|
|
await self._heartbeat_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
self.deregister_service()
|
|
|
|
def status(self) -> Dict[str, Any]:
|
|
"""获取当前状态"""
|
|
return {
|
|
"service_name": self.service_config.service_name,
|
|
"ip": self.service_config.ip,
|
|
"port": self.service_config.port,
|
|
"namespace": self.nacos_config.namespace,
|
|
"group": self.nacos_config.group_name,
|
|
"cluster": self.nacos_config.cluster_name,
|
|
"registered": self.is_registered,
|
|
"heartbeat_running": self._heartbeat_task is not None and not self._heartbeat_task.done(),
|
|
}
|