96 lines
3.3 KiB
Rust
96 lines
3.3 KiB
Rust
use serde::{Deserialize, Serialize};
|
||
use std::sync::Arc;
|
||
use std::sync::atomic::{AtomicU64, AtomicBool};
|
||
use std::time::{Duration, Instant};
|
||
use tokio::sync::Notify;
|
||
use tracing::info;
|
||
|
||
/// 队列统计信息
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct QueueStats {
|
||
pub request_frequency: f64,
|
||
pub request_count: u64,
|
||
}
|
||
|
||
/// 简化的任务队列（仅用于记录HTTP请求频率）
|
||
pub struct TaskQueue {
|
||
// HTTP请求频率监控
|
||
request_count: AtomicU64,
|
||
last_request_time: std::sync::Mutex<Instant>,
|
||
request_window: Duration,
|
||
is_paused: AtomicBool,
|
||
notify: Notify,
|
||
}
|
||
|
||
impl TaskQueue {
|
||
pub fn new(_max_concurrent_tasks: usize, _queue_size_limit: usize, _high_request_rate_threshold: f64, _idle_request_rate_threshold: f64) -> Self {
|
||
Self {
|
||
// HTTP请求频率监控
|
||
request_count: AtomicU64::new(0),
|
||
last_request_time: std::sync::Mutex::new(Instant::now()),
|
||
request_window: Duration::from_secs(10),
|
||
is_paused: AtomicBool::new(false),
|
||
notify: Notify::new(),
|
||
}
|
||
}
|
||
|
||
pub fn new_with_db(_max_concurrent_tasks: usize, _queue_size_limit: usize, _high_request_rate_threshold: f64, _idle_request_rate_threshold: f64, _db: Arc<crate::db::Database>) -> Self {
|
||
Self {
|
||
// HTTP请求频率监控
|
||
request_count: AtomicU64::new(0),
|
||
last_request_time: std::sync::Mutex::new(Instant::now()),
|
||
request_window: Duration::from_secs(10),
|
||
is_paused: AtomicBool::new(false),
|
||
notify: Notify::new(),
|
||
}
|
||
}
|
||
|
||
/// 记录HTTP请求（在每次HTTP请求处理时调用）
|
||
pub async fn record_http_request(&self) {
|
||
let now = Instant::now();
|
||
let mut last_request_time = self.last_request_time.lock().unwrap();
|
||
|
||
// 如果距离上次请求时间超过了窗口时间，重置计数器
|
||
if now.duration_since(*last_request_time) > self.request_window {
|
||
self.request_count.store(0, std::sync::atomic::Ordering::Release);
|
||
}
|
||
|
||
// 增加请求计数
|
||
self.request_count.fetch_add(1, std::sync::atomic::Ordering::Release);
|
||
*last_request_time = now;
|
||
}
|
||
|
||
/// 获取当前请求频率（请求数/秒）
|
||
pub fn get_request_frequency(&self) -> f64 {
|
||
let request_count = self.request_count.load(std::sync::atomic::Ordering::Acquire) as f64;
|
||
let window_seconds = self.request_window.as_secs_f64();
|
||
|
||
if window_seconds > 0.0 {
|
||
request_count / window_seconds
|
||
} else {
|
||
0.0
|
||
}
|
||
}
|
||
|
||
/// 获取队列统计信息
|
||
pub async fn get_stats(&self) -> QueueStats {
|
||
QueueStats {
|
||
// 添加请求频率信息
|
||
request_frequency: self.get_request_frequency(),
|
||
request_count: self.request_count.load(std::sync::atomic::Ordering::Acquire),
|
||
}
|
||
}
|
||
|
||
/// 暂停队列处理
|
||
pub fn pause(&self) {
|
||
self.is_paused.store(true, std::sync::atomic::Ordering::Release);
|
||
info!("任务队列已暂停");
|
||
}
|
||
|
||
/// 恢复队列处理
|
||
pub fn resume(&self) {
|
||
self.is_paused.store(false, std::sync::atomic::Ordering::Release);
|
||
self.notify.notify_one();
|
||
info!("任务队列已恢复");
|
||
}
|
||
} |