Kris
ecd3c90e80
新增设备身份管理、认证限流、并发通道、Webhook路由、RBAC权限控制、SSE/轮询降级等全套网关通道功能，包含： 1. 设备身份生成与签名验证 2. 设备令牌认证与速率限制 3. 内存+数据库双重设备注册表 4. 并发通道限流管理 5. Webhook安全处理与路由 6. RBAC权限校验系统 7. OpenAI API兼容适配层 8. Tailscale认证支持 9. HTTP轮询降级机制
115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
import asyncio
|
||
import logging
|
||
from collections.abc import AsyncIterator
|
||
from contextlib import asynccontextmanager
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
DEFAULT_MAX_LANES = 5
|
||
DEFAULT_GLOBAL_MAX_LANES = 50
|
||
|
||
|
||
class ChannelLane:
|
||
"""单个渠道账户的并发通道。
|
||
|
||
每个账户一个独立 lane，控制该账户下同时运行的 Agent 数量。
|
||
"""
|
||
|
||
def __init__(self, max_lanes: int = DEFAULT_MAX_LANES):
|
||
self._semaphore = asyncio.Semaphore(max_lanes)
|
||
self._active_runs: int = 0
|
||
self._busy: bool = False
|
||
self._max_lanes = max_lanes
|
||
|
||
@asynccontextmanager
|
||
async def acquire(self) -> AsyncIterator[None]:
|
||
async with self._semaphore:
|
||
self._active_runs += 1
|
||
if self._active_runs >= self._max_lanes:
|
||
self._busy = True
|
||
try:
|
||
yield
|
||
finally:
|
||
self._active_runs = max(0, self._active_runs - 1)
|
||
if self._active_runs < self._max_lanes:
|
||
self._busy = False
|
||
|
||
@property
|
||
def active_runs(self) -> int:
|
||
return self._active_runs
|
||
|
||
@property
|
||
def is_busy(self) -> bool:
|
||
return self._busy
|
||
|
||
@property
|
||
def available_slots(self) -> int:
|
||
return max(0, self._max_lanes - self._active_runs)
|
||
|
||
|
||
class LaneManager:
|
||
"""并发通道管理器。
|
||
|
||
按账户维度管理并发通道，同时提供全局上限保护。
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
per_account_max: int = DEFAULT_MAX_LANES,
|
||
global_max: int = DEFAULT_GLOBAL_MAX_LANES,
|
||
):
|
||
self._per_account_max = per_account_max
|
||
self._lanes: dict[str, ChannelLane] = {}
|
||
self._global_semaphore = asyncio.Semaphore(global_max)
|
||
|
||
def _key(self, channel_type: str, account_id: str) -> str:
|
||
return f"{channel_type}:{account_id}"
|
||
|
||
def get_lane(self, channel_type: str, account_id: str) -> ChannelLane:
|
||
key = self._key(channel_type, account_id)
|
||
if key not in self._lanes:
|
||
self._lanes[key] = ChannelLane(max_lanes=self._per_account_max)
|
||
return self._lanes[key]
|
||
|
||
@asynccontextmanager
|
||
async def run_with_lane(self, channel_type: str, account_id: str) -> AsyncIterator[ChannelLane]:
|
||
lane = self.get_lane(channel_type, account_id)
|
||
if lane.is_busy:
|
||
logger.warning(
|
||
"Channel lane busy: %s/%s, active=%d",
|
||
channel_type,
|
||
account_id,
|
||
lane.active_runs,
|
||
)
|
||
async with lane.acquire():
|
||
async with self._global_semaphore:
|
||
yield lane
|
||
|
||
def remove_lane(self, channel_type: str, account_id: str) -> None:
|
||
key = self._key(channel_type, account_id)
|
||
self._lanes.pop(key, None)
|
||
|
||
def get_stats(self) -> dict[str, dict]:
|
||
return {
|
||
key: {
|
||
"active_runs": lane.active_runs,
|
||
"busy": lane.is_busy,
|
||
"available_slots": lane.available_slots,
|
||
}
|
||
for key, lane in self._lanes.items()
|
||
}
|
||
|
||
def get_total_active(self) -> int:
|
||
return sum(lane.active_runs for lane in self._lanes.values())
|
||
|
||
def cleanup_idle(self) -> int:
|
||
idle_keys = [key for key, lane in self._lanes.items() if lane.active_runs == 0]
|
||
for key in idle_keys:
|
||
del self._lanes[key]
|
||
if idle_keys:
|
||
logger.debug("LaneManager cleaned up %d idle lanes", len(idle_keys))
|
||
return len(idle_keys)
|
||
|
||
|
||
lane_manager = LaneManager()
|