Agent 并发控制：从单线程到分布式锁

┌──────────────────────────────────────────────────┐
│              Agent 并发控制层次                    │
├──────────────────┬───────────────────────────────┤
│  请求级           │ 多用户并发请求                 │
│  (Request)       │ → 限流 + 队列 + 优先级         │
├──────────────────┼───────────────────────────────┤
│  Agent级          │ 单用户多Agent实例              │
│  (Agent)         │ → 信号量 + 状态隔离            │
├──────────────────┼───────────────────────────────┤
│  工具级           │ 多工具并行调用                 │
│  (Tool)          │ → 并发限制 + 超时控制          │
├──────────────────┼───────────────────────────────┤
│  资源级           │ 共享资源访问                   │
│  (Resource)      │ → 分布式锁 + 乐观锁            │
└──────────────────┴───────────────────────────────┘

import asyncio
import time
from collections import defaultdict

class MultiDimensionRateLimiter:
    """多维度限流器"""
    
    def __init__(self, redis_client):
        self.redis = redis_client
    
    async def check(
        self,
        user_id: str,
        agent_name: str,
        limits: dict
    ) -> bool:
        """多维度限流检查
        
        limits = {
            "qps": 10,           # 每秒请求数
            "concurrent": 5,     # 最大并发
            "daily_tokens": 500000,  # 日Token上限
            "monthly_cost": 100,     # 月成本上限
        }
        """
        checks = []
        
        # QPS 限流（滑动窗口）
        if "qps" in limits:
            checks.append(self._check_sliding_window(
                f"qps:{user_id}:{agent_name}",
                limits["qps"],
                window=1
            ))
        
        # 并发限流
        if "concurrent" in limits:
            checks.append(self._check_concurrent(
                f"conc:{user_id}:{agent_name}",
                limits["concurrent"]
            ))
        
        # Token 配额
        if "daily_tokens" in limits:
            checks.append(self._check_quota(
                f"tokens:{user_id}:{date.today()}",
                limits["daily_tokens"]
            ))
        
        results = await asyncio.gather(*checks)
        return all(results)
    
    async def _check_sliding_window(
        self,
        key: str,
        limit: int,
        window: int
    ) -> bool:
        """滑动窗口限流"""
        now = time.time()
        pipe = self.redis.pipeline()
        
        # 移除窗口外的记录
        pipe.zremrangebyscore(key, 0, now - window)
        # 添加当前请求
        pipe.zadd(key, {str(uuid.uuid4()): now})
        # 统计窗口内请求数
        pipe.zcard(key)
        # 设置过期时间
        pipe.expire(key, window * 2)
        
        results = await pipe.execute()
        count = results[2]
        
        if count > limit:
            # 移除刚添加的记录
            pipe.zrem(key, str(uuid.uuid4()))
            return False
        return True
    
    async def _check_concurrent(self, key: str, limit: int) -> bool:
        """并发数限流"""
        current = await self.redis.incr(key)
        if current == 1:
            await self.redis.expire(key, 300)  # 5分钟自动过期
        
        if current > limit:
            await self.redis.decr(key)
            return False
        return True
    
    async def release_concurrent(self, key: str):
        """释放并发槽"""
        await self.redis.decr(key)


class PriorityRequestQueue:
    """优先级请求队列"""
    
    def __init__(self, max_concurrent: int = 100):
        self.semaphore = asyncio.Semaphore(max_concurrent)
        self.queues = {
            "high": asyncio.Queue(maxsize=50),
            "normal": asyncio.Queue(maxsize=500),
            "low": asyncio.Queue(maxsize=1000),
        }
        self._running = True
    
    async def submit(
        self,
        request: Request,
        priority: str = "normal"
    ) -> Response:
        """提交请求"""
        queue = self.queues.get(priority, self.queues["normal"])
        
        if queue.full():
            raise QueueFullError(f"{priority} queue is full")
        
        future = asyncio.Future()
        await queue.put((request, future))
        
        # 启动消费者（如果未启动）
        asyncio.create_task(self._consume())
        
        return await future
    
    async def _consume(self):
        """消费请求：按优先级调度"""
        async with self.semaphore:
            # 按优先级获取请求
            for priority in ["high", "normal", "low"]:
                queue = self.queues[priority]
                if not queue.empty():
                    request, future = await queue.get()
                    try:
                        result = await self._process(request)
                        future.set_result(result)
                    except Exception as e:
                        future.set_exception(e)
                    return
            
            # 所有队列为空
            await asyncio.sleep(0.01)

class ParallelToolExecutor:
    """并行工具调用管理器"""
    
    def __init__(self):
        self.tool_limits = {}  # {tool_name: Semaphore}
        self.global_limit = asyncio.Semaphore(10)
    
    def register_tool(
        self,
        name: str,
        max_concurrent: int = 5,
        timeout: float = 30.0
    ):
        """注册工具并发限制"""
        self.tool_limits[name] = {
            "semaphore": asyncio.Semaphore(max_concurrent),
            "timeout": timeout
        }
    
    async def execute_parallel(
        self,
        tool_calls: list[ToolCall]
    ) -> list[ToolResult]:
        """并行执行多个工具调用"""
        
        tasks = []
        for call in tool_calls:
            task = asyncio.create_task(
                self._execute_single(call)
            )
            tasks.append(task)
        
        # 等待所有完成，设置全局超时
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        return [
            result if not isinstance(result, Exception) 
            else ToolResult(success=False, error=str(result))
            for result in results
        ]
    
    async def _execute_single(self, call: ToolCall) -> ToolResult:
        """执行单个工具调用（带并发限制和超时）"""
        tool_limit = self.tool_limits.get(call.tool)
        
        if not tool_limit:
            raise UnknownToolError(call.tool)
        
        # 双重信号量：全局 + 工具级
        async with self.global_limit:
            async with tool_limit["semaphore"]:
                try:
                    result = await asyncio.wait_for(
                        self._call_tool(call),
                        timeout=tool_limit["timeout"]
                    )
                    return ToolResult(success=True, data=result)
                except asyncio.TimeoutError:
                    return ToolResult(
                        success=False,
                        error=f"Tool {call.tool} timed out"
                    )


# 注册工具并发限制
executor = ParallelToolExecutor()
executor.register_tool("web_search", max_concurrent=5, timeout=10)
executor.register_tool("database_query", max_concurrent=3, timeout=15)
executor.register_tool("file_read", max_concurrent=20, timeout=5)
executor.register_tool("file_write", max_concurrent=2, timeout=10)
executor.register_tool("send_email", max_concurrent=1, timeout=30)

class DistributedLock:
    """基于 Redis 的分布式锁"""
    
    def __init__(self, redis_client):
        self.redis = redis_client
        self._local_locks = {}  # 本地锁缓存
    
    async def acquire(
        self,
        resource: str,
        holder: str,
        ttl: int = 30,
        timeout: float = 10.0
    ) -> bool:
        """获取分布式锁
        
        Args:
            resource: 锁定的资源名
            holder: 持有者标识（如 agent_id + session_id）
            ttl: 锁的生存时间（秒）
            timeout: 等待获取锁的超时时间
        """
        lock_key = f"lock:{resource}"
        start = time.time()
        
        while time.time() - start < timeout:
            # 尝试获取锁（原子操作）
            acquired = await self.redis.set(
                lock_key, holder, nx=True, ex=ttl
            )
            
            if acquired:
                # 启动看门狗续期
                asyncio.create_task(self._watchdog(resource, holder, ttl))
                return True
            
            # 检查持有者是否是自己（重入）
            current = await self.redis.get(lock_key)
            if current == holder:
                await self.redis.expire(lock_key, ttl)
                return True
            
            # 等待重试
            await asyncio.sleep(0.1)
        
        return False
    
    async def release(self, resource: str, holder: str) -> bool:
        """释放锁（使用 Lua 脚本保证原子性）"""
        lua_script = """
        if redis.call("GET", KEYS[1]) == ARGV[1] then
            return redis.call("DEL", KEYS[1])
        else
            return 0
        end
        """
        
        result = await self.redis.eval(
            lua_script, 1,
            f"lock:{resource}", holder
        )
        return bool(result)
    
    async def _watchdog(self, resource: str, holder: str, ttl: int):
        """看门狗：自动续期锁"""
        lock_key = f"lock:{resource}"
        renew_interval = ttl * 0.3  # 在30% TTL时续期
        
        while True:
            await asyncio.sleep(renew_interval)
            
            # 检查是否仍持有锁
            current = await self.redis.get(lock_key)
            if current != holder:
                break  # 锁已被释放或被抢
            
            # 续期
            await self.redis.expire(lock_key, ttl)


class AgentResourceLock:
    """Agent 资源锁管理"""
    
    def __init__(self, lock_manager: DistributedLock):
        self.locks = lock_manager
    
    async def with_resource_lock(
        self,
        resource: str,
        agent_id: str,
        action: callable,
        ttl: int = 30
    ):
        """在锁保护下执行操作"""
        acquired = await self.locks.acquire(
            resource, agent_id, ttl=ttl
        )
        
        if not acquired:
            raise LockAcquisitionError(
                f"Could not acquire lock for {resource}"
            )
        
        try:
            return await action()
        finally:
            await self.locks.release(resource, agent_id)

class OptimisticLock:
    """乐观锁：适用于低冲突场景"""
    
    async def update_with_version(
        self,
        resource_id: str,
        update_fn: callable,
        max_retries: int = 3
    ) -> any:
        """带版本控制的更新"""
        
        for attempt in range(max_retries):
            # 1. 读取当前版本
            current = await self.db.get(resource_id)
            if not current:
                raise NotFoundError(resource_id)
            
            version = current.version
            
            # 2. 计算更新
            updated = await update_fn(current.data)
            
            # 3. 乐观写入（带版本检查）
            result = await self.db.update_with_version(
                resource_id,
                data=updated,
                expected_version=version
            )
            
            if result.success:
                return updated
            
            # 版本冲突，重试
            logger.warning(
                f"Optimistic lock conflict on {resource_id}, "
                f"attempt {attempt+1}"
            )
            await asyncio.sleep(0.05 * attempt)  # 短暂退避
        
        raise ConcurrencyError(
            f"Failed after {max_retries} retries due to conflicts"
        )

class AgentInstanceManager:
    """Agent 实例管理——控制单用户的 Agent 并发"""
    
    def __init__(self, redis_client):
        self.redis = redis_client
        self.limits = {
            "free": 1,       # 免费用户：1个并发Agent
            "pro": 5,        # Pro用户：5个
            "enterprise": 20, # 企业用户：20个
        }
    
    async def acquire_slot(
        self,
        user_id: str,
        tier: str,
        agent_type: str
    ) -> AgentSlot:
        """获取 Agent 执行槽"""
        
        limit = self.limits.get(tier, 1)
        key = f"agent_slots:{user_id}"
        
        # 检查当前活跃数
        active = await self.redis.scard(key)
        if active >= limit:
            raise ConcurrencyLimitError(
                f"User {user_id} has {active}/{limit} active agents"
            )
        
        # 分配槽位
        slot_id = str(uuid.uuid4())
        await self.redis.sadd(key, slot_id)
        await self.redis.expire(key, 3600)  # 1小时过期
        
        return AgentSlot(
            slot_id=slot_id,
            user_id=user_id,
            agent_type=agent_type,
            acquired_at=time.time()
        )
    
    async def release_slot(self, slot: AgentSlot):
        """释放 Agent 执行槽"""
        key = f"agent_slots:{slot.user_id}"
        await self.redis.srem(key, slot.slot_id)


class AgentCoordinator:
    """Agent 协调器——管理多Agent并发执行"""
    
    async def run_agents_concurrent(
        self,
        agents: list[Agent],
        max_parallel: int = 5
    ) -> list[AgentResult]:
        """并发运行多个 Agent"""
        
        semaphore = asyncio.Semaphore(max_parallel)
        
        async def run_one(agent: Agent) -> AgentResult:
            async with semaphore:
                try:
                    result = await agent.run()
                    return AgentResult(
                        agent_id=agent.id,
                        success=True,
                        result=result
                    )
                except Exception as e:
                    return AgentResult(
                        agent_id=agent.id,
                        success=False,
                        error=str(e)
                    )
        
        return await asyncio.gather(*[run_one(a) for a in agents])
    
    async def run_agents_pipeline(
        self,
        agents: list[Agent],
        dependencies: dict  # {agent_id: [depends_on_ids]}
    ) -> dict[str, AgentResult]:
        """按依赖关系运行 Agent（拓扑排序）"""
        
        results = {}
        completed = set()
        
        while len(completed) < len(agents):
            # 找到可执行的 Agent（依赖已完成）
            ready = [
                agent for agent in agents
                if agent.id not in completed
                and all(
                    dep in completed 
                    for dep in dependencies.get(agent.id, [])
                )
            ]
            
            if not ready:
                # 检测死锁
                raise DeadlockError("Dependency cycle detected")
            
            # 并发执行就绪的 Agent
            batch_results = await self.run_agents_concurrent(ready)
            
            for result in batch_results:
                results[result.agent_id] = result
                completed.add(result.agent_id)
        
        return results

class ConcurrencyMonitor:
    """并发控制监控"""
    
    METRICS = {
        "active_requests": "当前活跃请求数",
        "queue_depth": "各优先级队列深度",
        "concurrent_agents": "各用户活跃Agent数",
        "tool_concurrent_usage": "各工具并发使用数",
        "lock_wait_time": "锁等待时间分布",
        "lock_contention_rate": "锁竞争率",
        "rate_limit_hits": "限流触发次数",
    }
    
    async def health_check(self) -> dict:
        return {
            "healthy": await self._check_health(),
            "active_connections": await self._count_active(),
            "queue_health": await self._check_queues(),
            "lock_health": await self._check_locks(),
            "alerts": await self._get_alerts(),
        }
    
    async def detect_issues(self) -> list[Issue]:
        issues = []
        
        # 检测队列堆积
        for priority, queue in self.queues.items():
            if queue.qsize() > queue.maxsize * 0.8:
                issues.append(Issue(
                    severity="warning",
                    type="queue_backlog",
                    detail=f"{priority} queue at {queue.qsize()}/{queue.maxsize}"
                ))
        
        # 检测锁饥饿
        lock_waits = await self._get_lock_wait_times()
        for resource, wait_time in lock_waits.items():
            if wait_time > 5.0:
                issues.append(Issue(
                    severity="high",
                    type="lock_starvation",
                    detail=f"{resource} lock wait: {wait_time:.1f}s"
                ))
        
        # 检测限流频繁触发
        rate_limit_hits = await self._get_rate_limit_hits()
        for user_id, hits in rate_limit_hits.items():
            if hits > 100:  # 每分钟超过100次
                issues.append(Issue(
                    severity="medium",
                    type="excessive_rate_limiting",
                    detail=f"User {user_id} hit rate limit {hits} times"
                ))
        
        return issues

□ 多维限流（QPS + 并发数 + Token + 成本）
□ 优先级队列（高/中/低优先级请求隔离）
□ 工具级并发限制（不同工具有不同并发上限）
□ 全局并发限制（总并发不超过系统容量）
□ 分布式锁保护共享资源（带看门狗续期）
□ 乐观锁用于低冲突场景
□ Agent 实例数按用户等级限制
□ 依赖感知的 Agent 调度（拓扑排序）
□ 死锁检测和预防
□ 并发监控和告警
□ 队列堆积自动扩容
□ 压测验证并发上限