104 lines
3.4 KiB
Rust
104 lines
3.4 KiB
Rust
//! TextDelta 节流聚合器(IM 频控防护)
|
|
//!
|
|
//! IM 平台消息 API 频控(飞书 ~50 QPS 应用级、企微 200 次/分钟),逐 delta 刷卡片必被限流。
|
|
//! 策略:per-conv 聚合缓冲 + 时间窗(默认 800ms)批量 flush;Completed 时强制冲刷尾包。
|
|
//!
|
|
//! 设计要点:
|
|
//! - 每个 key 一个缓冲,互不干扰
|
|
//! - 纯函数式:feed() 返回 Option<全量文本>(窗口到期才返回),不持有回调/不 await
|
|
//! - flush 语义:全量文本(IM 卡片更新是「全量内容」语义,非增量)
|
|
//! - finish() 冲刷尾包并清缓冲;abort() 丢弃缓冲(错误/停止,不输出残文)
|
|
|
|
use std::collections::HashMap;
|
|
use std::time::{Duration, Instant};
|
|
|
|
/// 单 key 缓冲
|
|
struct ConvBuffer {
|
|
/// 累积全量文本
|
|
full_text: String,
|
|
/// 上次 flush 时刻
|
|
last_flush: Instant,
|
|
}
|
|
|
|
/// 节流聚合器
|
|
pub struct DeltaThrottle {
|
|
window: Duration,
|
|
buffers: HashMap<String, ConvBuffer>,
|
|
}
|
|
|
|
impl DeltaThrottle {
|
|
pub fn new(window_ms: u64) -> Self {
|
|
Self { window: Duration::from_millis(window_ms), buffers: HashMap::new() }
|
|
}
|
|
|
|
/// 喂入 delta。窗口到期则返回 Some(全量文本)(调用方负责外发),否则 None。
|
|
pub fn feed(&mut self, key: &str, delta: &str) -> Option<String> {
|
|
let now = Instant::now();
|
|
let buf = self.buffers.entry(key.to_string()).or_insert_with(|| ConvBuffer {
|
|
full_text: String::new(),
|
|
last_flush: now,
|
|
});
|
|
buf.full_text.push_str(delta);
|
|
if buf.last_flush.elapsed() >= self.window {
|
|
buf.last_flush = now;
|
|
Some(buf.full_text.clone())
|
|
} else {
|
|
None
|
|
}
|
|
}
|
|
|
|
/// 会话完成:强制冲刷尾包 + 清缓冲。返回 Some(全量文本) 或(空会话)None。
|
|
pub fn finish(&mut self, key: &str) -> Option<String> {
|
|
self.buffers.remove(key).map(|buf| buf.full_text)
|
|
}
|
|
|
|
/// 丢弃缓冲(错误/停止,不输出残文)
|
|
///
|
|
/// 注:当前 TunnelEvent 无 Error/Stopped 变体,核心循环暂无调用点;
|
|
/// df-tunnel Phase2 扩展后接入。
|
|
#[allow(dead_code)]
|
|
pub fn abort(&mut self, key: &str) {
|
|
self.buffers.remove(key);
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn throttle_window_batch() {
|
|
let mut t = DeltaThrottle::new(60_000); // 60s 窗口:模拟窗口永不到期
|
|
assert_eq!(t.feed("k", "你"), None);
|
|
assert_eq!(t.feed("k", "好"), None);
|
|
assert_eq!(t.feed("k", "!"), None);
|
|
assert_eq!(t.finish("k"), Some("你好!".to_string()), "finish 冲刷全量");
|
|
}
|
|
|
|
#[test]
|
|
fn window_elapsed_flushes() {
|
|
let mut t = DeltaThrottle::new(0); // 0ms 窗口:每次 feed 都到期
|
|
assert_eq!(t.feed("k", "a"), Some("a".to_string()));
|
|
assert_eq!(t.feed("k", "b"), Some("ab".to_string()));
|
|
// 缓冲仍在 finish 前的最后 flush 版本;finish 拿走剩余
|
|
assert_eq!(t.finish("k"), Some("ab".to_string()));
|
|
}
|
|
|
|
#[test]
|
|
fn abort_drops() {
|
|
let mut t = DeltaThrottle::new(60_000);
|
|
t.feed("k", "残文");
|
|
t.abort("k");
|
|
assert_eq!(t.finish("k"), None, "abort 后缓冲不存在");
|
|
}
|
|
|
|
#[test]
|
|
fn keys_isolated() {
|
|
let mut t = DeltaThrottle::new(60_000);
|
|
t.feed("a", "1");
|
|
t.feed("b", "2");
|
|
assert_eq!(t.finish("a"), Some("1".to_string()));
|
|
assert_eq!(t.finish("b"), Some("2".to_string()));
|
|
}
|
|
}
|