72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""Agent 池服务：注册/注销/心跳/负载/离线剔除。"""
|
|
import logging
|
|
import time
|
|
|
|
import redis.asyncio as aioredis
|
|
|
|
from app.constants import AgentStatus
|
|
from app.models.schemas import AgentHeartbeat, AgentInfo, AgentRegister
|
|
from app.repository.agent_repo import AgentRepo
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AgentService:
|
|
def __init__(self, redis: aioredis.Redis, repo: AgentRepo | None = None):
|
|
self.redis = redis
|
|
self.repo = repo or AgentRepo(redis)
|
|
|
|
async def register(self, reg: AgentRegister) -> AgentInfo:
|
|
agent = AgentInfo(
|
|
agent_id=reg.agent_id,
|
|
endpoint=reg.endpoint,
|
|
agent_tags=reg.agent_tags,
|
|
max_concurrent=reg.max_concurrent,
|
|
current_load=reg.current_load,
|
|
status=AgentStatus.ONLINE,
|
|
create_time=time.time(),
|
|
last_heartbeat=time.time(),
|
|
)
|
|
created = await self.repo.upsert(agent)
|
|
logger.info("agent registered agent_id=%s created=%s", agent.agent_id, created)
|
|
return agent
|
|
|
|
async def unregister(self, agent_id: str) -> bool:
|
|
existed = await self.repo.get(agent_id)
|
|
if not existed:
|
|
return False
|
|
await self.repo.remove(agent_id)
|
|
logger.info("agent unregistered agent_id=%s", agent_id)
|
|
return True
|
|
|
|
async def heartbeat(self, hb: AgentHeartbeat) -> bool:
|
|
ok = await self.repo.beat(hb.agent_id, hb.current_load)
|
|
if not ok:
|
|
logger.warning("heartbeat from unknown agent agent_id=%s", hb.agent_id)
|
|
return ok
|
|
|
|
async def online(self) -> list[AgentInfo]:
|
|
return await self.repo.online()
|
|
|
|
async def all(self) -> list[AgentInfo]:
|
|
return await self.repo.all()
|
|
|
|
async def get(self, agent_id: str) -> AgentInfo | None:
|
|
return await self.repo.get(agent_id)
|
|
|
|
async def take_offline(self, agent_id: str) -> bool:
|
|
"""手动下线 Agent。"""
|
|
a = await self.repo.get(agent_id)
|
|
if not a:
|
|
return False
|
|
await self.repo.update_status(agent_id, AgentStatus.OFFLINE)
|
|
logger.info("agent taken offline agent_id=%s", agent_id)
|
|
return True
|
|
|
|
async def purge_stale(self, timeout: float) -> list[str]:
|
|
"""剔除超时未心跳的 Agent。"""
|
|
stale = await self.repo.stale_agents(timeout)
|
|
for aid in stale:
|
|
await self.repo.update_status(aid, AgentStatus.OFFLINE)
|
|
logger.info("agent purged (stale) agent_id=%s", aid)
|
|
return stale |