File size: 1,339 Bytes
bbb1195 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | // Rate Limiter
// 确保 API 调用间隔 ≥ 500ms
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::time::{sleep, Duration, Instant};
pub struct RateLimiter {
min_interval: Duration,
last_call: Arc<Mutex<Option<Instant>>>,
}
impl RateLimiter {
pub fn new(min_interval_ms: u64) -> Self {
Self {
min_interval: Duration::from_millis(min_interval_ms),
last_call: Arc::new(Mutex::new(None)),
}
}
pub async fn wait(&self) {
let mut last = self.last_call.lock().await;
if let Some(last_time) = *last {
let elapsed = last_time.elapsed();
if elapsed < self.min_interval {
sleep(self.min_interval - elapsed).await;
}
}
*last = Some(Instant::now());
}
}
#[cfg(test)]
mod tests {
use super::*;
use tokio::time::Instant;
#[tokio::test]
async fn test_rate_limiter() {
let limiter = RateLimiter::new(500);
let start = Instant::now();
limiter.wait().await; // 第一次调用，立即返回
let elapsed1 = start.elapsed().as_millis();
assert!(elapsed1 < 50);
limiter.wait().await; // 第二次调用，等待 500ms
let elapsed2 = start.elapsed().as_millis();
assert!(elapsed2 >= 500 && elapsed2 < 600);
}
}
|