146 lines
5.9 KiB
Python
146 lines
5.9 KiB
Python
"""Agent 池 Redis 存储层。"""
|
||
import time
|
||
from typing import Any
|
||
|
||
import redis.asyncio as aioredis
|
||
|
||
from app import constants as C
|
||
from app.constants import AgentStatus
|
||
from app.models.schemas import AgentInfo
|
||
|
||
|
||
class AgentRepo:
|
||
def __init__(self, redis: aioredis.Redis):
|
||
self.redis = redis
|
||
|
||
@staticmethod
|
||
def _info_key(agent_id: str) -> str:
|
||
return C.AGENT_INFO_KEY.format(agent_id=agent_id)
|
||
|
||
@staticmethod
|
||
def _tag_key(tag: str) -> str:
|
||
return C.AGENT_TAG_KEY.format(tag=tag)
|
||
|
||
# ---------- 写入 ----------
|
||
async def upsert(self, agent: AgentInfo, *, heartbeat: bool = False) -> bool:
|
||
"""写入 Agent 信息，返回是否新建。非心跳时为全量注册。"""
|
||
key = self._info_key(agent.agent_id)
|
||
existed = await self.redis.exists(key)
|
||
if not existed:
|
||
await self.redis.sadd(C.AGENT_ALL_SET, agent.agent_id)
|
||
# 建立标签索引
|
||
for tag in agent.agent_tags:
|
||
await self.redis.sadd(self._tag_key(tag), agent.agent_id)
|
||
mapping: dict[str, Any] = {
|
||
"agent_id": agent.agent_id,
|
||
"endpoint": agent.endpoint,
|
||
"agent_tags": ",".join(agent.agent_tags),
|
||
"max_concurrent": str(agent.max_concurrent),
|
||
"current_load": str(agent.current_load),
|
||
"last_heartbeat": str(agent.last_heartbeat),
|
||
"status": agent.status.value,
|
||
"create_time": str(agent.create_time),
|
||
}
|
||
await self.redis.hset(key, mapping=mapping)
|
||
# 心跳 ZSet 记录时间戳
|
||
await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent.agent_id: agent.last_heartbeat})
|
||
return not existed
|
||
|
||
async def update_status(self, agent_id: str, status: AgentStatus) -> bool:
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
await self.redis.hset(key, "status", status.value)
|
||
if status == AgentStatus.OFFLINE:
|
||
await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
|
||
return True
|
||
|
||
async def beat(self, agent_id: str, current_load: int) -> bool:
|
||
"""更新心跳时间与负载。"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
now = time.time()
|
||
await self.redis.hset(
|
||
key,
|
||
mapping={"last_heartbeat": str(now), "current_load": str(current_load), "status": AgentStatus.ONLINE.value},
|
||
)
|
||
await self.redis.zadd(C.AGENT_HEARTBEAT_ZSET, {agent_id: now})
|
||
return True
|
||
|
||
async def adjust_load(self, agent_id: str, delta: int) -> bool:
|
||
"""原子调整 Agent 负载（delta 可为正/负），负载下限为 0。"""
|
||
key = self._info_key(agent_id)
|
||
if not await self.redis.exists(key):
|
||
return False
|
||
# hincrby 为原子增量操作，兼容 fakeredis（不支持 Lua eval）
|
||
new = await self.redis.hincrby(key, "current_load", delta)
|
||
if new < 0:
|
||
await self.redis.hset(key, "current_load", "0")
|
||
return True
|
||
|
||
# ---------- 读取 ----------
|
||
async def get(self, agent_id: str) -> AgentInfo | None:
|
||
raw = await self.redis.hgetall(self._info_key(agent_id))
|
||
if not raw:
|
||
return None
|
||
return AgentInfo(
|
||
agent_id=raw.get("agent_id", ""),
|
||
endpoint=raw.get("endpoint", ""),
|
||
agent_tags=[t for t in raw.get("agent_tags", "").split(",") if t],
|
||
max_concurrent=int(raw.get("max_concurrent", 1) or 1),
|
||
current_load=int(raw.get("current_load", 0) or 0),
|
||
last_heartbeat=float(raw.get("last_heartbeat", 0) or 0),
|
||
status=AgentStatus(raw.get("status", AgentStatus.ONLINE.value)),
|
||
create_time=float(raw.get("create_time", 0) or 0),
|
||
)
|
||
|
||
async def online(self) -> list[AgentInfo]:
|
||
return [a for a in await self.all() if a.status == AgentStatus.ONLINE]
|
||
|
||
async def all(self) -> list[AgentInfo]:
|
||
ids = list(await self.redis.smembers(C.AGENT_ALL_SET))
|
||
out = []
|
||
for aid in ids:
|
||
a = await self.get(aid)
|
||
if a:
|
||
out.append(a)
|
||
return out
|
||
|
||
async def by_tags(self, tags: list[str]) -> list[AgentInfo]:
|
||
"""按标签索引取 Agent 池（取所有标签的交集，无标签则返回全部在线）。"""
|
||
if not tags:
|
||
return await self.online()
|
||
keys = [self._tag_key(t) for t in tags]
|
||
if len(keys) == 1:
|
||
ids = list(await self.redis.smembers(keys[0]))
|
||
else:
|
||
await self.redis.sinterstore("agent:tmp:intersect", keys)
|
||
ids = list(await self.redis.smembers("agent:tmp:intersect"))
|
||
await self.redis.delete("agent:tmp:intersect")
|
||
out = []
|
||
for aid in ids:
|
||
a = await self.get(aid)
|
||
if a and a.status == AgentStatus.ONLINE:
|
||
out.append(a)
|
||
return out
|
||
|
||
async def remove(self, agent_id: str) -> None:
|
||
"""注销：删除信息、心跳 ZSet、标签索引、全量索引。"""
|
||
a = await self.get(agent_id)
|
||
if a:
|
||
for tag in a.agent_tags:
|
||
await self.redis.srem(self._tag_key(tag), agent_id)
|
||
await self.redis.delete(self._info_key(agent_id))
|
||
await self.redis.zrem(C.AGENT_HEARTBEAT_ZSET, agent_id)
|
||
await self.redis.srem(C.AGENT_ALL_SET, agent_id)
|
||
|
||
async def stale_agents(self, timeout: float) -> list[str]:
|
||
"""返回超过 timeout 秒未心跳的 Agent ID（基于 ZSet score）。"""
|
||
cutoff = time.time() - timeout
|
||
scored = await self.redis.zrangebyscore(C.AGENT_HEARTBEAT_ZSET, 0, cutoff)
|
||
return list(scored)
|
||
|
||
async def heartbeat_map(self) -> dict[str, float]:
|
||
"""返回 {agent_id: last_heartbeat}。"""
|
||
return {k: float(v) for k, v in await self.redis.zrange(C.AGENT_HEARTBEAT_ZSET, 0, -1, withscores=True)} |