- 纯内存子会话隔离上下文,一轮回传结论报告 - 白名单钳制低风险并恒剔 spawn_agent 防递归 - 子代理策略下高风险自动拒绝,不挂审批防卡死 - 轮数令牌挂钟三重熔断 - 父停止与删会话级联取消子代理 - AiAgentFinished 事件挂父会话
437 lines
25 KiB
Rust
437 lines
25 KiB
Rust
use std::collections::HashMap;
|
||
use std::sync::atomic::AtomicUsize;
|
||
use std::sync::Arc;
|
||
use tokio::sync::{Mutex, Semaphore};
|
||
|
||
// ============================================================
|
||
// LLM 调用并发控制(双层 Semaphore + 可选 per-provider 层)
|
||
// ============================================================
|
||
|
||
/// LLM 调用并发控制 — 全局 LLM 调用限流 global + 单对话内 per_conv 双层 Semaphore + 可选 per-provider 层
|
||
///
|
||
/// 限流对象:所有真实 LLM 调用(主循环 stream_llm / 标题生成 / 知识提炼)。
|
||
/// 不限流本地工具执行(tools.execute)——本地操作无外部成本、不受 RPM 约束。
|
||
///
|
||
/// ## global(permits 默认 3,「LLM 调用并发限流」原义)
|
||
/// F-09 batch5 曾把 global 改「并发会话数上限」(loop 入口持整 loop),用户决策修正「不设并发会话上限」后
|
||
/// 已删除 loop 入口 acquire。global 回归原义:由各单次 LLM 调用点(stream_llm 重试循环 / 标题 / 压缩 /
|
||
/// 提炼 / 项目分析扫描)各自 `acquire_global()` 拿 permit、调用结束 Drop 释放,防 provider 429。
|
||
/// 多对话 loop 并发不限数(用户接受 token 暴增)。
|
||
///
|
||
/// ## per_conv 改 HashMap<conv_id, Semaphore>
|
||
/// 原应用级单信号量(AiSession 单例 + generating 互斥下退化)改为
|
||
/// `HashMap<String, Arc<Semaphore>>`:每对话一份(permits=2,主循环 + 标题 + 提炼各自限流)。
|
||
/// `run_agentic_loop` 入口 `acquire_per_conv(conv_id)` 拿 1 permit 持整个 loop 生命周期(含工具执行/审批
|
||
/// 等待/重试),防单对话内并发 LLM 调用失控。这是单对话内限流,非会话数限制,符合用户「不设上限」。
|
||
/// `acquire_per_conv(conv_id)`:lock HashMap → 无则建(permits=2)→ clone Arc → 释放 lock → acquire_owned。
|
||
/// conv 退出清理(`release_conv(conv_id)`):conv 删除时 remove 条目(防 HashMap 无限增长)。
|
||
///
|
||
/// 运行时调整:tokio Semaphore 的 permits 数构造时固定、不可增减,
|
||
/// 故 `global` 用 `Arc<Mutex<Arc<Semaphore>>>` 双层包装——替换内层 Arc 即重建 Semaphore。
|
||
/// 已持有旧 permit 的任务不受影响(permit 绑定旧 Semaphore,软收敛),
|
||
/// 新请求 lock 后克隆到最新 Arc、自动走新限制。旧 Semaphore 随最后 permit 释放而 drop。
|
||
/// per_conv HashMap 内每条 Arc<Semaphore> 不需运行时改 permits(对话内并发上限 2 固定),故无重建需求。
|
||
///
|
||
/// ## per-provider 层(可选)
|
||
/// `per_provider` 为 HashMap<provider_id, Arc<Semaphore>>。调用方(agentic loop)经
|
||
/// `try_acquire_for_provider(pid)` 取额外 permit,防单 provider 被打满(限流 429)。
|
||
/// **单 provider 场景**:若未调 `set_provider_caps`,HashMap 为空,
|
||
/// `try_acquire_for_provider` 返 `NotConfigured`(无限流,行为同 F-01 前)。零变化保证。
|
||
/// **三态语义(G4.1,2026-08-05)**:`NotConfigured`(未配置无限流) / `Acquired(permit)`
|
||
/// (拿到) / `Exhausted`(占满,非阻塞跳下一 candidate)。修复旧阻塞 `acquire_for_provider`
|
||
/// 在主 provider 信号量占满时整条 fallback 链卡死(降级失效)。
|
||
/// 全局容量 = min(sum(各 provider 上限), global_cap):由调用方在配置时约束
|
||
/// (set_provider_caps 传 min(sum, global_cap)),非运行时强约束。
|
||
///
|
||
/// ## Phase3 预留: per_sub_flow 层(批2-B,占位未接入)
|
||
/// `per_sub_flow` 为 HashMap<sub_flow_id, Arc<Semaphore>>,用于 Phase3 单对话并行多轮的
|
||
/// **子流并发上限**(单对话内并行 spawn 多条子流时,防子流数失控)。本批仅加层 + 占位,
|
||
/// **不接入调用点**(接入待 Phase3 子流 spawn 落地)。
|
||
///
|
||
/// **key 约定(文档化,非强约束)**:`sub_flow_id` 采用 `"{conv_id}::{sub_id}"` 格式,
|
||
/// 唯一标识某对话下的某条子流,便于 release_sub_flow 精确清理。
|
||
///
|
||
/// **三层语义**:
|
||
/// - per_sub_flow = 单对话内**子流**并发上限(每条子流一份 Semaphore,permits 默认 3,预留)
|
||
/// - per_conv = 单对话内**总**并发上限(permits 默认 2)
|
||
/// - global = 全应用**总**并发上限(permits 默认 3)
|
||
///
|
||
/// **理想配额(用户配置建议,非硬限)**: per_sub_permits ≤ per_conv_permits ≤ global_permits。
|
||
/// 实际限流由各层 Semaphore 自然保证:global Semaphore 硬限全应用并发不超 global permits
|
||
/// (跨对话 sum(per_conv) 可 > global,但 global acquire_owned 自然阻塞,不超卖)。
|
||
///
|
||
/// **acquire 语义(非阻塞)**:`acquire_per_sub_flow` 用 `try_acquire_owned`(非阻塞),
|
||
/// 耗尽返 None 降级串行——子流并行不阻塞主 loop(Phase3 子流 spawn 失败不拖垮整对话)。
|
||
#[derive(Clone)]
|
||
pub struct LlmConcurrency {
|
||
/// 全局 LLM 调用并发上限(permits 默认 3):F-09 batch5 修正后回归原义,由各单次 LLM 调用点
|
||
/// (stream_llm/标题/压缩/提炼/项目分析)各自 acquire/drop 防 429,不再由 loop 入口持整 loop。
|
||
global: Arc<Mutex<Arc<Semaphore>>>,
|
||
/// 单对话内并发上限(F-09 B 批5):HashMap<conv_id, Semaphore>,每对话 permits=2。
|
||
/// acquire 时按 conv_id 取/建;release_conv 在 loop 结束 + 无 pending 时 remove。
|
||
per_conv: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
|
||
/// 每对话内并发上限(permits 默认 2,构造时传入 new(3, 2) 的第二参)。
|
||
/// F-09 B 批5:AtomicUsize 支持运行时热改(ai_set_concurrency_config 的 per_conv_limit)。
|
||
/// 热改后**已建对话**的旧 Semaphore 不变(permits 构造时固定),**新建对话**用新值;
|
||
/// 为使热改立即全量生效,set_per_conv 同时清空 HashMap 强制重建(软收敛:旧 permit 随 Drop 释放)。
|
||
per_conv_permits: Arc<AtomicUsize>,
|
||
/// per-provider 信号量表。空 = 无 per-provider 限流(单 provider 路径零变化)。
|
||
/// Arc<Mutex<HashMap>>:运行时增删 provider 配置时替换/插入,acquire 时 clone Arc。
|
||
per_provider: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
|
||
/// Phase3 预留(批2-B): per-sub_flow 信号量表。key = sub_flow_id(约定 "{conv_id}::{sub_id}")。
|
||
/// 空 = 无子流并发限流(Phase3 未接入时零变化)。acquire 用 try_acquire_owned(非阻塞),
|
||
/// 耗尽返 None 降级串行,子流并行不阻塞主 loop。
|
||
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
|
||
per_sub_flow: Arc<Mutex<HashMap<String, Arc<Semaphore>>>>,
|
||
/// Phase3 预留(批2-B): 每子流并发上限(permits 默认 3,内部初始化,预留)。
|
||
/// AtomicUsize 支持运行时热改;set_per_sub_permits 同时清空 HashMap(软收敛,对齐 set_per_conv)。
|
||
/// new(global, per_conv) 签名保持向后兼容,per_sub_permits 内部 AtomicUsize::new(3)。
|
||
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
|
||
per_sub_permits: Arc<AtomicUsize>,
|
||
}
|
||
|
||
/// F-260614-04(2026-08-05 G4.1 裁决): per-provider 限流 acquire 的**显式三态**结果。
|
||
///
|
||
/// **None 双关陷阱(本枚举的动机)**:旧 `acquire_for_provider` 用 `Option` 返回值,
|
||
/// `None` 既表示「未配置无限流」又可能被 naive 调用方误读为「耗尽」——耗尽语义下
|
||
/// naive 用 None 表耗尽会静默绕过限流打爆 429。三态把两种语义显式拆分:
|
||
///
|
||
/// - `NotConfigured`:未配置(per_provider 表无该 provider → 无限流,permits 跳过不落表)。
|
||
/// - `Acquired(permit)`:成功拿到 permit(绑 Semaphore,Drop 自动释放)。
|
||
/// - `Exhausted`:信号量占满(非阻塞 try_acquire 失败),调用方跳下一 candidate。
|
||
pub enum ProviderAcquire {
|
||
/// 未配置(per_provider 表无该 provider → 无限流,permits 跳过不落表)。
|
||
/// 与耗尽语义完全区分:调用方无 permit 直接 proceed(单 provider 零变化)。
|
||
NotConfigured,
|
||
/// 成功拿到 permit(绑 Semaphore,Drop 自动释放)。
|
||
Acquired(tokio::sync::OwnedSemaphorePermit),
|
||
/// 该 provider 信号量占满(try_acquire_owned 失败,非阻塞)。
|
||
/// 调用方跳下一 candidate(候选链降级),不阻塞等待。
|
||
Exhausted,
|
||
}
|
||
|
||
impl LlmConcurrency {
|
||
pub fn new(global: usize, per_conv: usize) -> Self {
|
||
Self {
|
||
global: Arc::new(Mutex::new(Arc::new(Semaphore::new(global)))),
|
||
per_conv: Arc::new(Mutex::new(HashMap::new())),
|
||
per_conv_permits: Arc::new(AtomicUsize::new(per_conv)),
|
||
per_provider: Arc::new(Mutex::new(HashMap::new())),
|
||
// Phase3 预留(批2-B): per_sub_flow 默认 permits=3,内部初始化。
|
||
// 签名保持 new(global, per_conv) 向后兼容;Phase3 接入后用户可经
|
||
// set_per_sub_permits 热改(配置化待后续批,本批仅占位)。
|
||
per_sub_flow: Arc::new(Mutex::new(HashMap::new())),
|
||
per_sub_permits: Arc::new(AtomicUsize::new(3)),
|
||
}
|
||
}
|
||
|
||
/// 取全局 LLM 调用并发 permit:由各单次 LLM 调用点(stream_llm/标题/压缩/提炼/项目分析)
|
||
/// 各自调用、调用结束 Drop 释放。F-09 batch5 修正后不再由 run_agentic_loop 入口持整 loop。
|
||
/// 重建后(set_global)新请求自动走最新 Semaphore。
|
||
pub async fn acquire_global(&self) -> tokio::sync::OwnedSemaphorePermit {
|
||
let sema = self.global.lock().await.clone();
|
||
sema.acquire_owned().await.expect("llm global semaphore closed")
|
||
}
|
||
|
||
/// 取全局 LLM 调用并发 permit 的**非阻塞**版本(对齐 try_acquire_per_conv 语义)。
|
||
///
|
||
/// 用途:spawn_agent 子代理循环入口一次性取 global + per_conv(child) 双份额,
|
||
/// 任一耗尽即 failed(子代理不阻塞等待,防占满父 loop 的执行槽)。
|
||
/// 锁持有短(不含 try_acquire)。
|
||
pub async fn try_acquire_global(&self) -> Option<tokio::sync::OwnedSemaphorePermit> {
|
||
let sema = self.global.lock().await.clone();
|
||
// try_acquire_owned 非阻塞:Err(NoPermits) 返 None,不阻塞调用方。
|
||
// Semaphore 不会 close(无 close() 调用),Err 仅可能为 NoPermits。
|
||
sema.try_acquire_owned().ok()
|
||
}
|
||
|
||
/// 取单对话内并发 permit(F-09 B 批5):按 conv_id 取/建 Semaphore,permits=当前 per_conv_permits(默认 2)。
|
||
/// lock HashMap → 无则建 → clone Arc → 释放 lock → acquire_owned。锁持有短(不含 await acquire)。
|
||
pub async fn acquire_per_conv(&self, conv_id: &str) -> tokio::sync::OwnedSemaphorePermit {
|
||
let sema = {
|
||
let mut map = self.per_conv.lock().await;
|
||
let permits = self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst);
|
||
map.entry(conv_id.to_string())
|
||
.or_insert_with(|| Arc::new(Semaphore::new(permits)))
|
||
.clone()
|
||
};
|
||
sema.acquire_owned().await.expect("llm per_conv semaphore closed")
|
||
}
|
||
|
||
/// 取单对话内并发 permit 的**非阻塞**版本(对齐 `acquire_per_sub_flow` 语义)。
|
||
///
|
||
/// 与 `acquire_per_conv` 同样按 conv_id 取/建 Semaphore(permits=当前 per_conv_permits),
|
||
/// 区别仅在最后一步用 `try_acquire_owned`(非阻塞):耗尽返 `None`,调用方降级,不阻塞。
|
||
///
|
||
/// **用途**:压缩(`compress_via_llm`)在 agentic loop 内被调用,loop 入口已持 1 个 per_conv
|
||
/// permit;当 `set_per_conv(1)` 后 per_conv Semaphore 仅 1 permit,loop 入口占满,压缩再
|
||
/// `acquire_per_conv`(阻塞)将永久等待 → `is_compressing` 永真 → 卡死。改用本非阻塞版本:
|
||
/// 拿不到(None)调用方提前返回降级(下轮再试),非崩溃。
|
||
///
|
||
/// - 返 `Some(permit)`:成功获取,permit 随 Drop 自动释放。
|
||
/// - 返 `None`:本对话并发已耗尽,调用方降级(非错误)。锁持有短(不含 try_acquire)。
|
||
pub async fn try_acquire_per_conv(
|
||
&self,
|
||
conv_id: &str,
|
||
) -> Option<tokio::sync::OwnedSemaphorePermit> {
|
||
let sema = {
|
||
let mut map = self.per_conv.lock().await;
|
||
let permits = self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst);
|
||
map.entry(conv_id.to_string())
|
||
.or_insert_with(|| Arc::new(Semaphore::new(permits)))
|
||
.clone()
|
||
};
|
||
// try_acquire_owned 非阻塞:Err(NoPermits) 返 None 降级,不阻塞调用方。
|
||
// 对齐 acquire_global/per_conv 的 expect 前提:Semaphore 不会 close(无 close() 调用)。
|
||
sema.try_acquire_owned().ok()
|
||
}
|
||
|
||
/// F-09 B 批5:conv 退出清理。loop 结束 + 无 pending 审批时 remove 该 conv 的 Semaphore 条目。
|
||
/// **时机由调用方判断**(agentic loop 退出点):仅在确信无后续 acquire 时调用,否则误删会致
|
||
/// 该 conv 下次 acquire 重建 Semaphore(限流计数清零,非致命,但语义偏离)。
|
||
/// remove 后已持 permit 不受影响(permit 绑旧 Arc,随 Drop 释放),仅阻止新条目累积。
|
||
pub async fn release_conv(&self, conv_id: &str) {
|
||
self.per_conv.lock().await.remove(conv_id);
|
||
}
|
||
|
||
/// 重建全局 Semaphore(软收敛:旧 permit 不回收,待其释放后新限制完全生效)
|
||
pub async fn set_global(&self, permits: usize) {
|
||
*self.global.lock().await = Arc::new(Semaphore::new(permits));
|
||
}
|
||
|
||
/// 重建单对话 Semaphore(F-09 B 批5 后 per_conv 为 HashMap)。
|
||
/// 行为:更新 per_conv_permits(AtomicUsize) + 清空 HashMap(软收敛:旧 permit 随 Drop 释放,
|
||
/// 新对话 acquire 用新 permits 值重建 Semaphore)。已建对话若仍在跑,旧 Semaphore 不变;
|
||
/// 下次该 conv 新 acquire 时因 HashMap 已清空会重建为新 permits。
|
||
pub async fn set_per_conv(&self, permits: usize) {
|
||
self.per_conv_permits
|
||
.store(permits, std::sync::atomic::Ordering::SeqCst);
|
||
self.per_conv.lock().await.clear();
|
||
}
|
||
|
||
/// 读当前 per_conv permits 生效值(对齐 set_per_conv 的写,补足读侧 getter)。
|
||
///
|
||
/// 用途:get_app_config AI 工具查「当前 LLM 并发配置」时返给 LLM(global 无对称 getter:
|
||
/// Semaphore permits 封装在 Arc<Mutex<Arc<Semaphore>>> 内层无读取接口,故此处只返 per_conv,
|
||
/// global 不返避免误导)。热改(set_per_conv)后立即反映,无需重启。
|
||
pub fn current_per_conv_permits(&self) -> usize {
|
||
self.per_conv_permits.load(std::sync::atomic::Ordering::SeqCst)
|
||
}
|
||
|
||
// ============================================================
|
||
// Phase3 预留(批2-B): per_sub_flow 层 — 占位未接入调用点
|
||
// ============================================================
|
||
|
||
/// Phase3 预留(批2-B): 取单子流内并发 permit(非阻塞)。
|
||
///
|
||
/// 按 `sub_flow_id` 取/建 Semaphore(permits=当前 per_sub_permits,默认 3)。
|
||
/// lock HashMap → 无则建(or_insert_with, permits 取 per_sub_permits 当前值)
|
||
/// → clone Arc → 释放 lock → **try_acquire_owned**(非阻塞)。
|
||
///
|
||
/// **非阻塞语义(关键)**:用 try_acquire_owned 而非 acquire_owned——耗尽返 None
|
||
/// 调用方降级串行处理,子流并行不阻塞主 loop(Phase3 子流 spawn 失败不拖垮整对话)。
|
||
///
|
||
/// - 返 `Ok(Some(permit))`:成功获取,permit 随 Drop 自动释放。
|
||
/// - 返 `Ok(None)`:子流并发已耗尽,调用方降级串行(非错误)。
|
||
///
|
||
/// **key 约定**:`sub_flow_id` 采用 `"{conv_id}::{sub_id}"` 格式(文档化,调用方组装)。
|
||
/// 锁持有短(不含 try_acquire),对齐 acquire_per_conv 模式。
|
||
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
|
||
pub async fn acquire_per_sub_flow(
|
||
&self,
|
||
sub_flow_id: &str,
|
||
) -> Option<tokio::sync::OwnedSemaphorePermit> {
|
||
let sema = {
|
||
let mut map = self.per_sub_flow.lock().await;
|
||
let permits = self.per_sub_permits.load(std::sync::atomic::Ordering::SeqCst);
|
||
map.entry(sub_flow_id.to_string())
|
||
.or_insert_with(|| Arc::new(Semaphore::new(permits)))
|
||
.clone()
|
||
};
|
||
// try_acquire_owned 非阻塞:Err(NoPermits) 返 None 降级串行,不阻塞主 loop。
|
||
// 对齐 acquire_global/per_conv 的 expect 前提:Semaphore 不会 close(无 close() 调用)。
|
||
sema.try_acquire_owned().ok()
|
||
}
|
||
|
||
/// Phase3 预留(批2-B): 子流完成清理。remove 该 sub_flow 的 Semaphore 条目(防 HashMap 无限增长)。
|
||
///
|
||
/// **对齐 release_conv 模式**:remove 后已持 permit 不受影响(permit 绑旧 Arc,随 Drop 释放),
|
||
/// 仅阻止新条目累积。时机由调用方判断(Phase3 子流退出点)。
|
||
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
|
||
pub async fn release_sub_flow(&self, sub_flow_id: &str) {
|
||
self.per_sub_flow.lock().await.remove(sub_flow_id);
|
||
}
|
||
|
||
/// Phase3 预留(批2-B): 热改 per_sub_flow permits 上限。
|
||
///
|
||
/// 行为(对齐 set_per_conv):更新 per_sub_permits(AtomicUsize) + 清空 HashMap(软收敛:
|
||
/// 旧 permit 随 Drop 释放,新子流 acquire 用新 permits 值重建 Semaphore)。已建子流若仍在跑,
|
||
/// 旧 Semaphore 不变;下次该 sub_flow 新 acquire 时因 HashMap 已清空会重建为新 permits。
|
||
#[allow(dead_code)] // Phase3 子流 spawn 落地时接入调用点,本批仅占位。
|
||
pub async fn set_per_sub_permits(&self, permits: usize) {
|
||
self.per_sub_permits
|
||
.store(permits, std::sync::atomic::Ordering::SeqCst);
|
||
self.per_sub_flow.lock().await.clear();
|
||
}
|
||
|
||
/// F-260614-04(2026-08-05 G4.1 裁决): per-provider 并发 permit 的**非阻塞三态**版本。
|
||
///
|
||
/// **背景**:旧 `acquire_for_provider` 用阻塞 `acquire_owned()`,主 provider 信号量占满时
|
||
/// 整条 fallback 候选链卡死(降级失效)。compress/title 已改 try_acquire,候选链是唯一漏网。
|
||
/// 本方法用 `try_acquire_owned`(非阻塞),三态返回见 [`ProviderAcquire`]:
|
||
///
|
||
/// - `NotConfigured`:per_provider 表无该 provider(未配置 → 无限流,单 provider 零变化)。
|
||
/// - `Acquired(permit)`:成功拿到,permit 随 Drop 自动释放。
|
||
/// - `Exhausted`:信号量占满,调用方跳下一 candidate(候选链降级),不阻塞。
|
||
///
|
||
/// **仅 per-provider 层非阻塞**:global/per_conv 仍由 `acquire_global`/`acquire_per_conv`
|
||
/// 阻塞限流(调用方各自持有),不受本方法影响。
|
||
///
|
||
/// 锁持有短(不含 try_acquire),对齐 `acquire_per_sub_flow` 模式。
|
||
///
|
||
/// **调用方用法**(agentic 候选循环):
|
||
/// ```ignore
|
||
/// match llm_concurrency.try_acquire_for_provider(&candidate.id).await {
|
||
/// ProviderAcquire::NotConfigured => { /* 无 permit 直接 proceed */ }
|
||
/// ProviderAcquire::Acquired(permit) => { /* 持 permit stream,Drop 自动释放 */ }
|
||
/// ProviderAcquire::Exhausted => { /* 跳下一 candidate */ }
|
||
/// }
|
||
/// ```
|
||
pub async fn try_acquire_for_provider(
|
||
&self,
|
||
provider_id: &str,
|
||
) -> ProviderAcquire {
|
||
let sema = {
|
||
let map = self.per_provider.lock().await;
|
||
map.get(provider_id).cloned()
|
||
};
|
||
match sema {
|
||
// 无配置:无限流(permits 跳过不落表),显式区分于耗尽。
|
||
None => ProviderAcquire::NotConfigured,
|
||
Some(sema) => match sema.try_acquire_owned() {
|
||
// 成功拿到:permit 绑 Semaphore,Drop 自动释放。
|
||
Ok(permit) => ProviderAcquire::Acquired(permit),
|
||
// 占满:非阻塞返 Exhausted,调用方跳下一 candidate(降级)。
|
||
// try_acquire_owned 不会 Err 于 closed(Semaphore 无 close() 调用)。
|
||
Err(_) => ProviderAcquire::Exhausted,
|
||
},
|
||
}
|
||
}
|
||
|
||
/// 批量设置 per-provider 并发上限(替换整表)。
|
||
///
|
||
/// 调用方(Settings 配置热改 / 启动初始化)传入 `{ provider_id: permits }` map,
|
||
/// 替换整张 per_provider 表(软收敛:已持 permit 不回收)。全局容量约束 min(sum, global_cap)
|
||
/// 由调用方在构造 map 时应用(本函数不做强约束,只落表)。
|
||
///
|
||
/// 传空 map → 清空 per_provider 表(所有 provider 回退到无 per-provider 限流)。
|
||
pub async fn set_provider_caps(&self, caps: HashMap<String, usize>) {
|
||
let mut map = self.per_provider.lock().await;
|
||
map.clear();
|
||
for (pid, permits) in caps {
|
||
// permits=0 等同无配置(Semaphore::new(0) 永远 acquire 不到 → 死锁),
|
||
// 故 permits=0 跳过(不落表 → try_acquire_for_provider 返 NotConfigured → 无限流)。
|
||
if permits > 0 {
|
||
map.insert(pid, Arc::new(Semaphore::new(permits)));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
// ============================================================
|
||
// F-260614-04(2026-08-05 G4.1): try_acquire_for_provider 三态测试
|
||
// NotConfigured / Acquired / Exhausted 显式区分(修 None 双关:None='未配置无限流',
|
||
// naive 用 None 表耗尽会静默绕过限流)。仅验证 per-provider 层三态语义,
|
||
// global/per_conv 仍阻塞限流(不经本方法)。
|
||
// ============================================================
|
||
|
||
/// 未配置(空 per_provider 表)→ NotConfigured:无限流,单 provider 零变化。
|
||
#[tokio::test]
|
||
async fn try_acquire_for_provider_not_configured() {
|
||
let lc = LlmConcurrency::new(3, 2);
|
||
let r = lc.try_acquire_for_provider("p1").await;
|
||
assert!(matches!(r, ProviderAcquire::NotConfigured));
|
||
}
|
||
|
||
/// 配置且有空闲 permit → Acquired(permit 绑 Semaphore,Drop 自动释放)。
|
||
#[tokio::test]
|
||
async fn try_acquire_for_provider_acquired() {
|
||
let lc = LlmConcurrency::new(3, 2);
|
||
let mut caps = HashMap::new();
|
||
caps.insert("p1".to_string(), 2usize);
|
||
lc.set_provider_caps(caps).await;
|
||
|
||
let r = lc.try_acquire_for_provider("p1").await;
|
||
let permit = match r {
|
||
ProviderAcquire::Acquired(p) => p,
|
||
other => panic!("配置后应返回 Acquired,实际 {:?}", discriminant(&other)),
|
||
};
|
||
// Drop 释放后仍可再拿(permit 绑旧 Semaphore Arc,set 后计数守恒)。
|
||
drop(permit);
|
||
let r2 = lc.try_acquire_for_provider("p1").await;
|
||
assert!(matches!(r2, ProviderAcquire::Acquired(_)));
|
||
}
|
||
|
||
/// 配置但信号量占满 → Exhausted(非阻塞,不等待)。
|
||
#[tokio::test]
|
||
async fn try_acquire_for_provider_exhausted() {
|
||
let lc = LlmConcurrency::new(3, 2);
|
||
let mut caps = HashMap::new();
|
||
caps.insert("p1".to_string(), 1usize); // 仅 1 permit
|
||
lc.set_provider_caps(caps).await;
|
||
|
||
// 第一次拿到(permit 仍持于作用域内,不 Drop)。
|
||
let first = lc.try_acquire_for_provider("p1").await;
|
||
assert!(matches!(first, ProviderAcquire::Acquired(_)));
|
||
// 第二次耗尽:信号量满,非阻塞返 Exhausted。
|
||
let second = lc.try_acquire_for_provider("p1").await;
|
||
assert!(matches!(second, ProviderAcquire::Exhausted));
|
||
|
||
// Drop 首 permit 后恢复可再拿。
|
||
drop(first);
|
||
let third = lc.try_acquire_for_provider("p1").await;
|
||
assert!(matches!(third, ProviderAcquire::Acquired(_)));
|
||
}
|
||
|
||
/// 配置后清空(传空 map)→ 回退 NotConfigured(permits 跳过不落表)。
|
||
#[tokio::test]
|
||
async fn try_acquire_for_provider_not_configured_after_caps_cleared() {
|
||
let lc = LlmConcurrency::new(3, 2);
|
||
let mut caps = HashMap::new();
|
||
caps.insert("p1".to_string(), 1usize);
|
||
lc.set_provider_caps(caps).await;
|
||
lc.set_provider_caps(HashMap::new()).await; // 清空整表
|
||
|
||
let r = lc.try_acquire_for_provider("p1").await;
|
||
assert!(matches!(r, ProviderAcquire::NotConfigured));
|
||
}
|
||
|
||
/// 配置了 p1,未配置 p2 → p1 Acquired / p2 NotConfigured(按 provider 区分,不互相污染)。
|
||
#[tokio::test]
|
||
async fn try_acquire_for_provider_distinct_providers() {
|
||
let lc = LlmConcurrency::new(3, 2);
|
||
let mut caps = HashMap::new();
|
||
caps.insert("p1".to_string(), 1usize);
|
||
lc.set_provider_caps(caps).await;
|
||
|
||
// p1 配置 → Acquired;p2 未配置 → NotConfigured。
|
||
assert!(matches!(lc.try_acquire_for_provider("p1").await, ProviderAcquire::Acquired(_)));
|
||
assert!(matches!(lc.try_acquire_for_provider("p2").await, ProviderAcquire::NotConfigured));
|
||
// p1 仅 1 permit 已满 → p2 不受影响。
|
||
assert!(matches!(lc.try_acquire_for_provider("p2").await, ProviderAcquire::NotConfigured));
|
||
}
|
||
|
||
/// 辅助:仅用于 panic 信息里给枚举变体一个 Debug 标签(无需求时保持简单)。
|
||
fn discriminant(r: &ProviderAcquire) -> &'static str {
|
||
match r {
|
||
ProviderAcquire::NotConfigured => "NotConfigured",
|
||
ProviderAcquire::Acquired(_) => "Acquired",
|
||
ProviderAcquire::Exhausted => "Exhausted",
|
||
}
|
||
}
|
||
}
|