Files
rust-imager/src/queue.rs
2026-03-25 18:31:11 +08:00

723 lines
25 KiB
Rust
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

use serde::{Deserialize, Serialize};
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicUsize, AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::sync::{Mutex, Notify};
use tokio::time::sleep;
use tracing::{info, warn, error, debug};
use uuid::Uuid;
use image::GenericImageView;
use crate::models::{AlbumIndex, CreateAlbum, CreateImage};
/// 任务类型
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum TaskType {
ThumbnailGeneration, // 缩略图生成
AlbumImport, // 相册导入
}
/// 任务优先级
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub enum TaskPriority {
High = 0, // 前端请求的缩略图
Medium = 1, // 相册导入
Low = 2, // 后台质量升级
}
/// 通用任务结构
#[derive(Debug, Clone)]
pub struct Task {
pub id: Uuid,
pub task_type: TaskType,
pub priority: TaskPriority,
pub created_at: Instant,
pub data: TaskData,
}
/// 任务数据
#[derive(Debug, Clone)]
pub enum TaskData {
ThumbnailGeneration {
image_id: Uuid,
original_data: Vec<u8>,
},
AlbumImport {
album_path: String,
album_data: Option<Vec<u8>>,
},
}
/// 队列统计信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct QueueStats {
pub pending_tasks: usize,
pub active_tasks: usize,
pub completed_tasks: usize,
pub failed_tasks: usize,
pub queue_size_limit: usize,
pub is_paused: bool,
pub task_type_stats: Vec<TaskTypeStats>,
// HTTP请求频率信息
pub request_frequency: f64,
pub request_count: u64,
}
/// 任务类型统计
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskTypeStats {
pub task_type: TaskType,
pub pending: usize,
pub completed: usize,
pub failed: usize,
}
/// 智能任务队列
pub struct TaskQueue {
queue: Mutex<VecDeque<Task>>,
active_tasks: AtomicUsize,
completed_tasks: AtomicUsize,
failed_tasks: AtomicUsize,
max_concurrent_tasks: usize,
queue_size_limit: usize,
is_paused: AtomicBool,
notify: Notify,
last_activity: Mutex<Instant>,
// HTTP请求频率监控
request_count: AtomicU64,
last_request_time: Mutex<Instant>,
request_window: Duration,
// 配置阈值
high_request_rate_threshold: f64,
idle_request_rate_threshold: f64,
}
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 {
queue: Mutex::new(VecDeque::new()),
active_tasks: AtomicUsize::new(0),
completed_tasks: AtomicUsize::new(0),
failed_tasks: AtomicUsize::new(0),
max_concurrent_tasks,
queue_size_limit,
is_paused: AtomicBool::new(false),
notify: Notify::new(),
last_activity: Mutex::new(Instant::now()),
// HTTP请求频率监控
request_count: AtomicU64::new(0),
last_request_time: Mutex::new(Instant::now()),
request_window: Duration::from_secs(10),
// 配置阈值
high_request_rate_threshold,
idle_request_rate_threshold,
}
}
/// 添加任务到队列
pub async fn add_task(&self, task: Task) -> Result<(), String> {
let mut queue = self.queue.lock().await;
// 检查队列是否已满
if queue.len() >= self.queue_size_limit {
return Err("队列已满,无法添加新任务".to_string());
}
// 检查是否已有相同ID的任务
if queue.iter().any(|t| t.id == task.id) {
return Err("该任务已在队列中".to_string());
}
// 根据优先级插入任务(高优先级在前)
let insert_position = queue
.iter()
.position(|t| t.priority > task.priority)
.unwrap_or(queue.len());
queue.insert(insert_position, task);
// 更新最后活动时间
*self.last_activity.lock().await = Instant::now();
// 通知工作线程
self.notify.notify_one();
debug!("任务已添加到队列,当前队列大小: {}", queue.len());
Ok(())
}
/// 从队列获取下一个任务
pub async fn next_task(&self) -> Option<Task> {
// 检查是否暂停
if self.is_paused.load(Ordering::Acquire) {
return None;
}
// 检查并发任务数量
if self.active_tasks.load(Ordering::Acquire) >= self.max_concurrent_tasks {
return None;
}
let mut queue = self.queue.lock().await;
if let Some(task) = queue.pop_front() {
self.active_tasks.fetch_add(1, Ordering::Release);
Some(task)
} else {
None
}
}
/// 标记任务完成
pub fn task_completed(&self, success: bool) {
self.active_tasks.fetch_sub(1, Ordering::Release);
if success {
self.completed_tasks.fetch_add(1, Ordering::Release);
} else {
self.failed_tasks.fetch_add(1, Ordering::Release);
}
}
/// 暂停队列处理
pub fn pause(&self) {
self.is_paused.store(true, Ordering::Release);
info!("任务队列已暂停");
}
/// 恢复队列处理
pub fn resume(&self) {
self.is_paused.store(false, Ordering::Release);
self.notify.notify_one();
info!("任务队列已恢复");
}
/// 记录HTTP请求(在每次HTTP请求处理时调用)
pub async fn record_http_request(&self) {
let now = Instant::now();
let mut last_request_time = self.last_request_time.lock().await;
// 如果距离上次请求时间超过了窗口时间,重置计数器
if now.duration_since(*last_request_time) > self.request_window {
self.request_count.store(0, Ordering::Release);
}
// 增加请求计数
self.request_count.fetch_add(1, Ordering::Release);
*last_request_time = now;
}
/// 获取当前请求频率(请求数/秒)
pub fn get_request_frequency(&self) -> f64 {
let request_count = self.request_count.load(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
}
}
/// 检查服务器负载是否过高(基于HTTP请求频率)
pub async fn is_server_busy(&self) -> bool {
let request_frequency = self.get_request_frequency();
// 基于HTTP请求频率判断业务繁忙度
// 高请求频率(超过配置阈值)时认为服务器繁忙
let is_high_request_rate = request_frequency > self.high_request_rate_threshold;
// 如果请求频率高,或者并发任务/队列大小超过阈值,认为服务器繁忙
is_high_request_rate
}
/// 获取队列统计信息
pub async fn get_stats(&self) -> QueueStats {
let queue = self.queue.lock().await;
// 按任务类型统计
let mut type_stats = std::collections::HashMap::new();
for task in queue.iter() {
let stats = type_stats
.entry(task.task_type.clone())
.or_insert((0, 0, 0));
stats.0 += 1; // pending
}
let task_type_stats = type_stats
.into_iter()
.map(|(task_type, (pending, completed, failed))| TaskTypeStats {
task_type,
pending,
completed,
failed,
})
.collect();
QueueStats {
pending_tasks: queue.len(),
active_tasks: self.active_tasks.load(Ordering::Acquire),
completed_tasks: self.completed_tasks.load(Ordering::Acquire),
failed_tasks: self.failed_tasks.load(Ordering::Acquire),
queue_size_limit: self.queue_size_limit,
is_paused: self.is_paused.load(Ordering::Acquire),
task_type_stats,
// 添加请求频率信息
request_frequency: self.get_request_frequency(),
request_count: self.request_count.load(Ordering::Acquire),
}
}
/// 清理过期的任务(超过一定时间的任务)
pub async fn cleanup_expired_tasks(&self, max_age: Duration) {
let mut queue = self.queue.lock().await;
let now = Instant::now();
let initial_len = queue.len();
queue.retain(|task| now.duration_since(task.created_at) <= max_age);
let removed_count = initial_len - queue.len();
if removed_count > 0 {
warn!("清理了 {} 个过期任务", removed_count);
}
}
/// 等待队列变为空闲状态
pub async fn wait_for_idle(&self, timeout: Duration) -> bool {
let start = Instant::now();
while start.elapsed() < timeout {
sleep(Duration::from_millis(100)).await;
}
false
}
}
/// 缩略图生成Worker
pub struct ThumbnailWorker {
queue: Arc<TaskQueue>,
image_processor: Arc<crate::image_processor::ImageProcessor>,
}
impl ThumbnailWorker {
pub fn new(
queue: Arc<TaskQueue>,
image_processor: Arc<crate::image_processor::ImageProcessor>,
) -> Self {
Self {
queue,
image_processor,
}
}
/// 启动缩略图Worker
pub async fn start(self) {
info!("启动缩略图生成Worker");
loop {
// 检查队列中是否有任务
if let Some(task) = self.queue.next_task().await {
if matches!(task.task_type, TaskType::ThumbnailGeneration) {
debug!("缩略图Worker开始处理任务: {:?}", task.id);
let success = self.process_thumbnail_task(task).await;
self.queue.task_completed(success);
// 在处理完每个任务后检查服务器负载
if self.queue.is_server_busy().await {
info!("服务器负载较高,暂停队列处理");
self.queue.pause();
// 每30秒检查一次服务器状态,直到服务器不忙时恢复队列
loop {
tokio::time::sleep(Duration::from_secs(30)).await;
if !self.queue.is_server_busy().await {
info!("服务器负载恢复正常,恢复队列处理");
self.queue.resume();
break;
}
info!("服务器仍然繁忙,继续等待...");
}
}
} else {
// 如果不是缩略图任务,放回队列
let mut queue = self.queue.queue.lock().await;
queue.push_front(task);
}
} else {
// 队列为空,等待新任务通知
self.queue.notify.notified().await;
}
// 短暂休眠避免CPU占用过高
sleep(Duration::from_millis(100)).await;
}
}
/// 处理缩略图任务
async fn process_thumbnail_task(&self, task: Task) -> bool {
if let TaskData::ThumbnailGeneration {
image_id,
original_data,
} = task.data
{
// 使用具体的错误处理,避免 Send trait 问题
let result = self
.image_processor
.generate_thumbnail_from_memory(&original_data, &image_id.to_string())
.await;
match result {
Ok(thumbnail_data) => {
// 保存缩略图到文件
let thumbnail_path = self
.image_processor
.config
.thumbnails_dir
.join(image_id.to_string());
match tokio::fs::write(&thumbnail_path, &thumbnail_data).await {
Ok(_) => {
info!("✅ 缩略图生成成功: {:?}", image_id);
true
}
Err(e) => {
warn!("保存缩略图失败: {:?} - {}", image_id, e);
false
}
}
}
Err(e) => {
warn!("缩略图生成失败: {:?} - {:?}", image_id, e);
false
}
}
} else {
warn!("任务类型不匹配,期望缩略图生成任务");
false
}
}
}
/// 相册导入Worker
pub struct AlbumImportWorker {
queue: Arc<TaskQueue>,
db: Arc<crate::database::Database>,
config: crate::config::StorageConfig,
}
impl AlbumImportWorker {
pub fn new(
queue: Arc<TaskQueue>,
db: Arc<crate::database::Database>,
config: crate::config::StorageConfig,
) -> Self {
Self { queue, db, config }
}
/// 启动相册导入Worker
pub async fn start(self) {
info!("启动相册导入Worker");
loop {
// 检查队列中是否有任务
if let Some(task) = self.queue.next_task().await {
if matches!(task.task_type, TaskType::AlbumImport) {
debug!("相册导入Worker开始处理任务: {:?}", task.id);
let success = self.process_album_import_task(task).await;
self.queue.task_completed(success);
// 在处理完每个任务后检查服务器负载
if self.queue.is_server_busy().await {
info!("服务器负载较高,暂停队列处理");
self.queue.pause();
// 每30秒检查一次服务器状态,直到服务器不忙时恢复队列
loop {
tokio::time::sleep(Duration::from_secs(30)).await;
if !self.queue.is_server_busy().await {
info!("服务器负载恢复正常,恢复队列处理");
self.queue.resume();
break;
}
info!("服务器仍然繁忙,继续等待...");
}
}
} else {
// 如果不是相册导入任务,放回队列并减少active_tasks计数器
let mut queue = self.queue.queue.lock().await;
queue.push_front(task);
// 手动减少active_tasks计数器,因为任务没有被处理
self.queue.active_tasks.fetch_sub(1, Ordering::Release);
}
} else {
// 队列为空,等待新任务通知
self.queue.notify.notified().await;
}
// 短暂休眠避免CPU占用过高
sleep(Duration::from_millis(100)).await;
}
}
/// 处理相册导入任务
async fn process_album_import_task(&self, task: Task) -> bool {
if let TaskData::AlbumImport {
album_path,
album_data: _,
} = task.data
{
info!("开始处理相册导入任务: {}", album_path);
// 将相对路径与images_dir组合成完整路径
let full_path = std::path::Path::new(&self.config.images_dir).join(&album_path);
// 检查相册文件夹是否存在
if !full_path.exists() {
error!("相册文件夹不存在: {}", full_path.display());
return false;
}
// 导入单个相册
match self.import_single_album(&album_path, &full_path).await {
Ok(count) => {
info!("相册导入成功: {} (导入 {} 张图片)", album_path, count);
true
}
Err(e) => {
error!("相册导入失败: {} - {}", album_path, e);
false
}
}
} else {
warn!("任务类型不匹配,期望相册导入任务");
false
}
}
/// 导入单个相册(按照import.rs中的逻辑实现)
async fn import_single_album(&self, folder_name: &str, path: &std::path::Path) -> Result<usize, Box<dyn std::error::Error + Send + Sync>> {
info!("开始导入相册: {}", folder_name);
// 读取index.json文件
let index_file_path = path.join("index.json");
let album_title = if index_file_path.exists() {
match std::fs::read_to_string(&index_file_path) {
Ok(content) => {
match serde_json::from_str::<AlbumIndex>(&content) {
Ok(index) => {
info!("从index.json读取相册标题: {}", index.title);
index.title
}
Err(e) => {
warn!("解析index.json失败: {}, 使用文件夹名作为标题", e);
folder_name.to_string()
}
}
}
Err(e) => {
warn!("读取index.json失败: {}, 使用文件夹名作为标题", e);
folder_name.to_string()
}
}
} else {
info!("index.json不存在,使用文件夹名作为标题");
folder_name.to_string()
};
// 创建新相册
let new_album = CreateAlbum {
title: album_title.clone(),
document: folder_name.to_string(),
description: Some(format!("从文件夹 {} 自动导入的相册", folder_name)),
};
let album_id = self.db.create_album(new_album).await?;
info!("创建相册: {} (ID: {})", album_title, album_id);
// 读取图片URL列表(如果有index.json)
let image_urls = if index_file_path.exists() {
match std::fs::read_to_string(&index_file_path) {
Ok(content) => {
match serde_json::from_str::<AlbumIndex>(&content) {
Ok(index) => {
info!("从index.json读取到 {} 个图片URL", index.images.len());
Some(index.images)
}
Err(_) => None
}
}
Err(_) => None
}
} else {
None
};
// 快速扫描图片文件
let image_files = self.scan_image_files(path, folder_name)?;
if image_files.is_empty() {
info!("相册 {} 中没有找到图片文件", folder_name);
return Ok(0);
}
info!("发现 {} 张图片,开始处理...", image_files.len());
// 并行处理图片信息获取
let mut image_tasks = Vec::new();
for (index, (filename, image_path)) in image_files.iter().enumerate() {
let image_url = image_urls.as_ref().and_then(|urls| {
if urls.len() > index {
Some(urls[index].clone())
} else {
None
}
});
let filename_clone = filename.clone();
let image_path_clone = image_path.clone();
let album_id_clone = album_id;
let folder_name_clone = folder_name.to_string();
let task = tokio::task::spawn_blocking(move || {
// 在阻塞线程中处理图片信息
let (file_size, width, height) = match AlbumImportWorker::get_image_info_static(&image_path_clone) {
Ok(info) => info,
Err(e) => {
eprintln!("获取图片信息失败: {} - {}", image_path_clone.display(), e);
(0, 0, 0)
}
};
CreateImage {
album_id: album_id_clone,
filename: format!("{}/{}", folder_name_clone, filename_clone),
original_url: image_url,
title: Some(format!("图片 {}", index + 1)),
description: Some(format!("从文件夹 {} 导入的图片", folder_name_clone)),
file_size,
width,
height,
mime_type: match image_path_clone.extension().and_then(|s| s.to_str()).unwrap_or("jpg").to_lowercase().as_str() {
"jpg" | "jpeg" => "image/jpeg".to_string(),
"png" => "image/png".to_string(),
"gif" => "image/gif".to_string(),
"webp" => "image/webp".to_string(),
"bmp" => "image/bmp".to_string(),
_ => "image/jpeg".to_string(),
},
}
});
image_tasks.push(task);
}
// 收集所有图片信息
let mut images_to_insert = Vec::new();
for task in image_tasks {
match task.await {
Ok(image) => images_to_insert.push(image),
Err(e) => error!("图片处理任务失败: {}", e),
}
}
// 批量插入图片
if !images_to_insert.is_empty() {
match self.db.add_images(images_to_insert).await {
Ok(ids) => {
info!("成功导入 {} 张图片到相册 {}", ids.len(), folder_name);
// 更新相册图片数量
if let Err(e) = self.db.update_album_image_count(album_id).await {
error!("更新相册图片数量失败: {}", e);
}
Ok(ids.len())
}
Err(e) => {
error!("批量插入图片失败: {}", e);
Ok(0)
}
}
} else {
info!("没有图片需要插入到相册 {}", folder_name);
Ok(0)
}
}
/// 扫描图片文件(从import.rs迁移)
fn scan_image_files(&self, path: &std::path::Path, _folder_name: &str) -> Result<Vec<(String, std::path::PathBuf)>, Box<dyn std::error::Error + Send + Sync>> {
let mut image_files = Vec::new();
for entry in std::fs::read_dir(path)? {
let entry = entry?;
let file_path = entry.path();
if file_path.is_file() {
if let Some(extension) = file_path.extension().and_then(|s| s.to_str()) {
let ext_lower = extension.to_lowercase();
if ext_lower == "jpg" || ext_lower == "jpeg" || ext_lower == "png" ||
ext_lower == "gif" || ext_lower == "webp" || ext_lower == "bmp" {
if let Some(filename) = file_path.file_name().and_then(|s| s.to_str()) {
image_files.push((filename.to_string(), file_path));
}
}
}
}
}
Ok(image_files)
}
/// 静态方法:获取图片信息(用于spawn_blocking)
fn get_image_info_static(image_path: &std::path::Path) -> Result<(i64, i32, i32), Box<dyn std::error::Error + Send + Sync>> {
let metadata = std::fs::metadata(image_path)?;
let file_size = metadata.len() as i64;
// 使用image库获取图片尺寸
let img = image::open(image_path)?;
let (width, height) = img.dimensions();
Ok((file_size, width as i32, height as i32))
}
}
/// 任务工厂 - 用于创建不同类型的任务
pub struct TaskFactory;
impl TaskFactory {
/// 创建缩略图生成任务
pub fn create_thumbnail_task(
image_id: Uuid,
original_data: Vec<u8>,
priority: TaskPriority,
) -> Task {
Task {
id: Uuid::new_v4(),
task_type: TaskType::ThumbnailGeneration,
priority,
created_at: Instant::now(),
data: TaskData::ThumbnailGeneration {
image_id,
original_data,
},
}
}
/// 创建相册导入任务
pub fn create_album_import_task(album_path: String, album_data: Option<Vec<u8>>) -> Task {
Task {
id: Uuid::new_v4(),
task_type: TaskType::AlbumImport,
priority: TaskPriority::Medium,
created_at: Instant::now(),
data: TaskData::AlbumImport {
album_path,
album_data,
},
}
}
}