62 lines
2.1 KiB
Rust
62 lines
2.1 KiB
Rust
use std::collections::HashMap;
|
||
use std::sync::Arc;
|
||
use tokio::sync::{Mutex, watch};
|
||
|
||
/// 共享的 Agent 取消注册表。
|
||
///
|
||
/// 每个正在执行的 Agent 在启动前按 topic_id 注册一个 watch::Sender，
|
||
/// 外部（如 /stop 命令）通过 cancel_by_topic() 发送取消信号。
|
||
/// Agent 循环内部通过 watch::Receiver::has_changed() 检测取消。
|
||
///
|
||
/// key 使用 topic_id（UUID），全局唯一，精确到话题级别。
|
||
#[derive(Clone)]
|
||
pub struct CancelManager {
|
||
tokens: Arc<Mutex<HashMap<String, watch::Sender<()>>>>,
|
||
}
|
||
|
||
impl CancelManager {
|
||
pub fn new() -> Self {
|
||
Self {
|
||
tokens: Arc::new(Mutex::new(HashMap::new())),
|
||
}
|
||
}
|
||
|
||
/// 按 topic_id 注册一个取消通道，返回 receiver 供 Agent 持有。
|
||
///
|
||
/// 如果同 topic_id 已有注册，旧 sender 被覆盖并 drop，
|
||
/// 旧 receiver 将收到通道关闭信号。
|
||
pub async fn register(&self, topic_id: &str) -> watch::Receiver<()> {
|
||
let (tx, rx) = watch::channel(());
|
||
self.tokens.lock().await.insert(topic_id.to_string(), tx);
|
||
rx
|
||
}
|
||
|
||
/// 按 topic_id 发送取消信号并移除注册条目。
|
||
///
|
||
/// 返回 `true` 表示找到了对应的任务并发送了取消信号，
|
||
/// 返回 `false` 表示没有找到对应的任务（可能已经完成或从未注册）。
|
||
pub async fn cancel_by_topic(&self, topic_id: &str) -> bool {
|
||
if let Some(tx) = self.tokens.lock().await.remove(topic_id) {
|
||
// send 可能失败（receiver 已被 drop），这不影响语义
|
||
let _ = tx.send(());
|
||
true
|
||
} else {
|
||
false
|
||
}
|
||
}
|
||
|
||
/// 按 topic_id 正常完成后清理注册条目（幂等）。
|
||
///
|
||
/// 与 cancel_by_topic() 不同，此方法不发送取消信号，仅移除条目。
|
||
/// 如果条目已被 cancel_by_topic() 移除，此调用为 no-op。
|
||
pub async fn remove_by_topic(&self, topic_id: &str) {
|
||
self.tokens.lock().await.remove(topic_id);
|
||
}
|
||
}
|
||
|
||
impl Default for CancelManager {
|
||
fn default() -> Self {
|
||
Self::new()
|
||
}
|
||
}
|