From de5f451cd11bdcb1cb55eb9ee92a44d5d6d0994f Mon Sep 17 00:00:00 2001 From: 237899745 <237899745@users.noreply.git.workyai.cn> Date: Sun, 26 Jul 2026 00:57:29 +0800 Subject: [PATCH] perf: improve worker and storage throughput --- .env.example | 7 +- docker/.env.production.example | 7 +- docker/docker-compose.prod.yml | 3 +- docs/architecture.md | 4 +- docs/database.md | 2 +- docs/deployment.md | 4 +- docs/observability.md | 2 + migrations/013_usage_events_time_index.sql | 2 + src/api/admin_storage.rs | 49 ++- src/config.rs | 17 +- src/main.rs | 1 + src/services/storage.rs | 247 +++++++++++++- src/state.rs | 2 + src/worker/mod.rs | 365 +++++++++++++++++---- 14 files changed, 611 insertions(+), 101 deletions(-) create mode 100644 migrations/013_usage_events_time_index.sql diff --git a/.env.example b/.env.example index 73e2be3..9732e85 100644 --- a/.env.example +++ b/.env.example @@ -9,12 +9,15 @@ RUST_LOG=info,tower_http=info,imageforge=debug # 数据库 DATABASE_URL=postgres://imageforge:devpassword@localhost:5432/imageforge -DATABASE_MAX_CONNECTIONS=10 +DATABASE_MAX_CONNECTIONS=16 # Redis REDIS_URL=redis://localhost:6379 -# Worker 并发（每个批量任务内同时处理的文件数） +# Worker 同时处理的批量任务数 +WORKER_TASK_CONCURRENCY=4 + +# 每个批量任务内同时处理的文件数 WORKER_CONCURRENCY=4 # 单进程图片处理并发上限（API 与 Worker 均生效，默认等于 CPU 线程数） diff --git a/docker/.env.production.example b/docker/.env.production.example index 851d372..1328ffd 100644 --- a/docker/.env.production.example +++ b/docker/.env.production.example @@ -21,10 +21,11 @@ ADMIN_EMAIL=admin@example.com ADMIN_USERNAME=admin ADMIN_PASSWORD=replace-with-a-strong-admin-password -# A four-core host should start with two image jobs per process. -DATABASE_MAX_CONNECTIONS=10 +# 8-core starting point: four tasks, two files per task, four CPU-bound image jobs. +DATABASE_MAX_CONNECTIONS=16 +WORKER_TASK_CONCURRENCY=4 WORKER_CONCURRENCY=2 -IMAGE_PROCESSING_CONCURRENCY=2 +IMAGE_PROCESSING_CONCURRENCY=4 # Resource ceilings tuned for an 8-core / 16 GB application host. POSTGRES_MEMORY_LIMIT=2g diff --git a/docker/docker-compose.prod.yml b/docker/docker-compose.prod.yml index 994646e..e57a102 100644 --- a/docker/docker-compose.prod.yml +++ b/docker/docker-compose.prod.yml @@ -2,7 +2,7 @@ name: imageforge x-imageforge-environment: &imageforge-environment DATABASE_URL: postgres://imageforge:${POSTGRES_PASSWORD:?POSTGRES_PASSWORD is required}@postgres:5432/imageforge - DATABASE_MAX_CONNECTIONS: ${DATABASE_MAX_CONNECTIONS:-10} + DATABASE_MAX_CONNECTIONS: ${DATABASE_MAX_CONNECTIONS:-16} REDIS_URL: redis://redis:6379 JWT_SECRET: ${JWT_SECRET:?JWT_SECRET is required} JWT_EXPIRY_HOURS: ${JWT_EXPIRY_HOURS:-168} @@ -10,6 +10,7 @@ x-imageforge-environment: &imageforge-environment BILLING_PROVIDER: ${BILLING_PROVIDER:-stripe} STORAGE_PATH: /app/uploads PUBLIC_BASE_URL: ${PUBLIC_BASE_URL:-http://localhost:8080} + WORKER_TASK_CONCURRENCY: ${WORKER_TASK_CONCURRENCY:-4} WORKER_CONCURRENCY: ${WORKER_CONCURRENCY:-2} IMAGE_PROCESSING_CONCURRENCY: ${IMAGE_PROCESSING_CONCURRENCY:-2} ALLOW_ANONYMOUS_UPLOAD: ${ALLOW_ANONYMOUS_UPLOAD:-true} diff --git a/docs/architecture.md b/docs/architecture.md index 91711a2..c951b2e 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -276,7 +276,8 @@ dotenvy = "0.15" ### 1. 并发处理 - 使用 Tokio 异步运行时 - 图片压缩使用 `spawn_blocking` 避免阻塞异步线程 -- 可配置 Worker 线程数 +- `WORKER_TASK_CONCURRENCY` 控制任务级并发，避免大批量任务独占 Worker +- `WORKER_CONCURRENCY` 控制单任务内文件并发，`IMAGE_PROCESSING_CONCURRENCY` 作为进程级 CPU 闸门 ```rust // 在独立线程池中执行 CPU 密集型压缩 @@ -292,6 +293,7 @@ let result = tokio::task::spawn_blocking(move || { ### 3. 缓存策略 - Redis 缓存用户会话 +- S3 活动端点与历史端点缓存 5 秒，AWS SDK Client 按端点和内外网地址复用连接池 - 可选：相同图片哈希缓存结果（去重） ## 安全考虑 diff --git a/docs/database.md b/docs/database.md index 7266148..e5e13d7 100644 --- a/docs/database.md +++ b/docs/database.md @@ -540,7 +540,7 @@ Dead-letter stream: stream:compress_jobs:dead Dead-letter max length: 10000 ``` -Worker 每次优先处理本消费者 pending，并认领空闲超过 5 分钟的其他消费者消息。消息第 3 次投递仍失败时写入死信流、将任务标记为失败并 ACK 原消息；已删除或已进入终态的任务直接 ACK。 +Worker 最多并发处理 `WORKER_TASK_CONCURRENCY` 个消息。调度器只领取新消息或空闲超过 5 分钟的其他消费者消息，处理中的长任务每 60 秒通过 `XCLAIM JUSTID` 刷新 pending 空闲时间，避免被其他 Worker 重复认领且不增加投递次数；单条消息失败后在原任务槽内按 2 秒、4 秒指数退避并通过 `XCLAIM` 重新投递。第 3 次投递仍失败时写入死信流、将任务标记为失败并 ACK 原消息；已删除或已进入终态的任务直接 ACK。 ### 5.4 任务进度（可选） ``` diff --git a/docs/deployment.md b/docs/deployment.md index d320333..fd265fa 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -11,7 +11,7 @@ - 最低 2 核 CPU、4GB 内存；启用 AVIF 和独立 Worker 时建议 4 核、8GB 内存 - 首次构建可访问 Docker Hub 与 crates.io -Debian 13、4 核 CPU、8GB 内存的实测起始值为 `WORKER_CONCURRENCY=2` 和 `IMAGE_PROCESSING_CONCURRENCY=2`。AVIF 是 CPU 密集型编码，不要直接把并发设置为 CPU 核数的数倍。 +Debian 13、4 核 CPU、8GB 内存的起始值建议为 `WORKER_TASK_CONCURRENCY=2`、`WORKER_CONCURRENCY=2` 和 `IMAGE_PROCESSING_CONCURRENCY=2`；8 核应用服务器可从 `4/2/4` 开始。三者分别表示同时处理的任务数、单任务内文件数和单进程 CPU 图片处理上限。最后一项是全局 CPU 闸门，因此不要把它设置为 CPU 核数的数倍。任务并发提高后，数据库连接池建议至少为 `WORKER_TASK_CONCURRENCY * WORKER_CONCURRENCY + 4`，生产示例使用 16。 生产 Compose 默认按 8 核 16GB 主机设置可覆盖的资源上限：API 3GB、Worker 8GB、PostgreSQL 2GB、Redis 1GB；对应变量为 `API_MEMORY_LIMIT`、`WORKER_MEMORY_LIMIT`、`POSTGRES_MEMORY_LIMIT` 和 `REDIS_MEMORY_LIMIT`。Redis 的 `REDIS_MAXMEMORY` 默认 768MB，达到上限后返回写入错误而不是继续挤占宿主机内存。 @@ -170,4 +170,4 @@ docker stats df -h ``` -若图片压缩长时间排队，先检查 CPU，再下调 `IMAGE_PROCESSING_CONCURRENCY`；若 API 健康但批量任务不推进，检查 Worker 日志和 Redis 状态。 +若图片压缩长时间排队，先检查 CPU，再调整 `IMAGE_PROCESSING_CONCURRENCY`；大任务阻塞小任务时提高 `WORKER_TASK_CONCURRENCY`，单个批量任务推进过慢时再评估 `WORKER_CONCURRENCY`。若 API 健康但批量任务不推进，检查 Worker 日志、Redis pending 数和数据库连接池等待情况。 diff --git a/docs/observability.md b/docs/observability.md index dc84f3f..01a1dec 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -59,6 +59,8 @@ curl --fail http://127.0.0.1:18180/metrics - `imageforge_storage_fallbacks_total` - `imageforge_dead_letters_total` +`imageforge_storage_fallbacks_total` 是存储容量风险信号，而不仅是普通降级统计。建议对 `increase(imageforge_storage_fallbacks_total[5m]) > 0` 持续 5 分钟设置高优先级告警，并同步监控应用服务器 `uploads` 卷使用率，避免 S3 长时间不可用时本地回退写满磁盘。 + Prometheus 抓取示例： ```yaml diff --git a/migrations/013_usage_events_time_index.sql b/migrations/013_usage_events_time_index.sql new file mode 100644 index 0000000..87ee5e9 --- /dev/null +++ b/migrations/013_usage_events_time_index.sql @@ -0,0 +1,2 @@ +CREATE INDEX IF NOT EXISTS idx_usage_events_occurred_at + ON usage_events(occurred_at DESC); diff --git a/src/api/admin_storage.rs b/src/api/admin_storage.rs index 667289e..247c57f 100644 --- a/src/api/admin_storage.rs +++ b/src/api/admin_storage.rs @@ -176,6 +176,7 @@ async fn create_storage_endpoint( .fetch_one(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "创建存储端点失败").with_source(err))?; + state.storage_cache.invalidate(); audit_storage_action( &state, @@ -336,6 +337,7 @@ async fn update_storage_endpoint( "端点配置在测试期间发生变化，请重新编辑", ) })?; + state.storage_cache.invalidate(); audit_storage_action( &state, @@ -364,9 +366,18 @@ async fn test_storage_endpoint( let (_jar, admin_id) = require_admin(&state, jar, &headers, ip).await?; let endpoint = storage::get_endpoint(&state, endpoint_id).await?; ensure_configurable(&endpoint)?; + let tested_config_updated_at = endpoint.updated_at; if let Err(err) = storage::test_endpoint(&state, &endpoint).await { let detail = storage_test_message(&err); - record_test_result(&state, endpoint_id, false, Some(&detail), admin_id).await?; + record_test_result( + &state, + endpoint_id, + false, + Some(&detail), + admin_id, + tested_config_updated_at, + ) + .await?; audit_storage_action( &state, admin_id, @@ -379,7 +390,15 @@ async fn test_storage_endpoint( return Err(AppError::new(ErrorCode::StorageUnavailable, detail)); } - record_test_result(&state, endpoint_id, true, None, admin_id).await?; + record_test_result( + &state, + endpoint_id, + true, + None, + admin_id, + tested_config_updated_at, + ) + .await?; let endpoint = storage::get_endpoint(&state, endpoint_id).await?; ensure_configurable(&endpoint)?; audit_storage_action( @@ -416,7 +435,15 @@ async fn activate_storage_endpoint( if let Err(err) = storage::test_endpoint(&state, &endpoint).await { let detail = storage_test_message(&err); - record_test_result(&state, endpoint_id, false, Some(&detail), admin_id).await?; + record_test_result( + &state, + endpoint_id, + false, + Some(&detail), + admin_id, + tested_config_updated_at, + ) + .await?; return Err(AppError::new(ErrorCode::StorageUnavailable, detail)); } @@ -456,6 +483,7 @@ async fn activate_storage_endpoint( tx.commit().await.map_err(|err| { AppError::new(ErrorCode::Internal, "提交存储切换事务失败").with_source(err) })?; + state.storage_cache.invalidate(); let endpoint = storage::get_endpoint(&state, endpoint_id).await?; audit_storage_action( @@ -505,6 +533,7 @@ async fn delete_storage_endpoint( "端点配置已发生变化，请刷新后重试", )); } + state.storage_cache.invalidate(); audit_storage_action( &state, admin_id, @@ -624,8 +653,9 @@ async fn record_test_result( ok: bool, error: Option<&str>, admin_id: Uuid, + expected_updated_at: DateTime, ) -> Result<(), AppError> { - sqlx::query( + let updated = sqlx::query( r#" UPDATE storage_endpoints SET last_test_at = NOW(), @@ -633,16 +663,25 @@ async fn record_test_result( last_test_error = $3, updated_at = NOW(), updated_by = $4 - WHERE id = $1 + WHERE id = $1 AND updated_at = $5 AND deleted_at IS NULL "#, ) .bind(endpoint_id) .bind(ok) .bind(error) .bind(admin_id) + .bind(expected_updated_at) .execute(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "记录存储测试结果失败").with_source(err))?; + if updated.rows_affected() == 0 { + state.storage_cache.invalidate(); + return Err(AppError::new( + ErrorCode::InvalidRequest, + "端点配置在测试期间发生变化，请重新测试", + )); + } + state.storage_cache.invalidate(); Ok(()) } diff --git a/src/config.rs b/src/config.rs index 510625d..b86bbec 100644 --- a/src/config.rs +++ b/src/config.rs @@ -14,6 +14,7 @@ pub struct Config { pub redis_url: String, + pub worker_task_concurrency: u32, pub worker_concurrency: u32, pub image_processing_concurrency: u32, @@ -62,11 +63,16 @@ impl Config { let redis_url = env_string("REDIS_URL") .ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 REDIS_URL"))?; - let worker_concurrency = env_u32("WORKER_CONCURRENCY").unwrap_or_else(|| { - std::thread::available_parallelism() - .map(|v| v.get() as u32) - .unwrap_or(4) - }); + let worker_task_concurrency = env_u32("WORKER_TASK_CONCURRENCY") + .filter(|value| *value > 0) + .unwrap_or(4); + let worker_concurrency = env_u32("WORKER_CONCURRENCY") + .filter(|value| *value > 0) + .unwrap_or_else(|| { + std::thread::available_parallelism() + .map(|v| v.get() as u32) + .unwrap_or(4) + }); let image_processing_concurrency = env_u32("IMAGE_PROCESSING_CONCURRENCY") .filter(|value| *value > 0) .unwrap_or_else(|| { @@ -126,6 +132,7 @@ impl Config { database_url, database_max_connections, redis_url, + worker_task_concurrency, worker_concurrency, image_processing_concurrency, jwt_secret, diff --git a/src/main.rs b/src/main.rs index dc3332e..822a219 100644 --- a/src/main.rs +++ b/src/main.rs @@ -44,6 +44,7 @@ async fn main() -> Result<(), AppError> { mailer: std::sync::Arc::new(mailer), image_processing_semaphore, runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), + storage_cache: crate::services::storage::StorageCache::new(), }; match state.config.role.as_str() { diff --git a/src/services/storage.rs b/src/services/storage.rs index 3584bfa..a8fca68 100644 --- a/src/services/storage.rs +++ b/src/services/storage.rs @@ -13,14 +13,17 @@ use aws_sdk_s3::Client; use bytes::Bytes; use chrono::{DateTime, Datelike, Utc}; use sqlx::FromRow; +use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::time::Duration; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use url::Url; use uuid::Uuid; const MULTIPART_THRESHOLD: u64 = 64 * 1024 * 1024; const MULTIPART_PART_SIZE: usize = 16 * 1024 * 1024; +const STORAGE_ENDPOINT_CACHE_TTL: Duration = Duration::from_secs(5); #[derive(Debug, Clone, FromRow)] pub struct StorageEndpoint { @@ -60,6 +63,153 @@ pub struct ObjectLocator { pub key: String, } +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +enum EndpointKind { + Internal, + Public, +} + +#[derive(Clone, Default)] +pub struct StorageCache { + inner: Arc>, +} + +#[derive(Default)] +struct StorageCacheState { + active_endpoint: Option, + endpoints: HashMap, + clients: HashMap<(Uuid, EndpointKind), CachedClient>, +} + +struct CachedActiveEndpoint { + loaded_at: Instant, + endpoint: Option, +} + +struct CachedEndpoint { + loaded_at: Instant, + endpoint: StorageEndpoint, +} + +struct CachedClient { + identity: ClientIdentity, + client: Client, +} + +#[derive(PartialEq, Eq)] +struct ClientIdentity { + endpoint_url: String, + region: String, + force_path_style: bool, + access_key_encrypted: String, + secret_key_encrypted: String, +} + +impl ClientIdentity { + fn new(endpoint: &StorageEndpoint, kind: EndpointKind) -> Self { + Self { + endpoint_url: match kind { + EndpointKind::Internal => endpoint.internal_endpoint.clone(), + EndpointKind::Public => endpoint.public_endpoint.clone(), + }, + region: endpoint.region.clone(), + force_path_style: endpoint.force_path_style, + access_key_encrypted: endpoint.access_key_encrypted.clone(), + secret_key_encrypted: endpoint.secret_key_encrypted.clone(), + } + } +} + +impl StorageCache { + pub fn new() -> Self { + Self::default() + } + + pub fn invalidate(&self) { + *self + .inner + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) = StorageCacheState::default(); + } + + fn active_endpoint(&self) -> Option> { + let cache = self + .inner + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let cached = cache.active_endpoint.as_ref()?; + (cached.loaded_at.elapsed() < STORAGE_ENDPOINT_CACHE_TTL).then(|| cached.endpoint.clone()) + } + + fn store_active_endpoint(&self, endpoint: Option) { + let mut cache = self + .inner + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + if let Some(endpoint) = endpoint.as_ref() { + cache.endpoints.insert( + endpoint.id, + CachedEndpoint { + loaded_at: Instant::now(), + endpoint: endpoint.clone(), + }, + ); + } + cache.active_endpoint = Some(CachedActiveEndpoint { + loaded_at: Instant::now(), + endpoint, + }); + } + + fn endpoint(&self, endpoint_id: Uuid) -> Option { + let cache = self + .inner + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()); + let cached = cache.endpoints.get(&endpoint_id)?; + (cached.loaded_at.elapsed() < STORAGE_ENDPOINT_CACHE_TTL).then(|| cached.endpoint.clone()) + } + + fn store_endpoint(&self, endpoint: StorageEndpoint) { + self.inner + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .endpoints + .insert( + endpoint.id, + CachedEndpoint { + loaded_at: Instant::now(), + endpoint, + }, + ); + } + + fn client(&self, endpoint: &StorageEndpoint, kind: EndpointKind) -> Option { + let identity = ClientIdentity::new(endpoint, kind); + self.inner + .read() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clients + .get(&(endpoint.id, kind)) + .filter(|cached| cached.identity == identity) + .map(|cached| cached.client.clone()) + } + + fn store_client(&self, endpoint: &StorageEndpoint, kind: EndpointKind, client: Client) { + self.inner + .write() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + .clients + .insert( + (endpoint.id, kind), + CachedClient { + identity: ClientIdentity::new(endpoint, kind), + client, + }, + ); + } +} + pub async fn list_endpoints(state: &AppState) -> Result, AppError> { sqlx::query_as::<_, StorageEndpoint>( r#" @@ -82,7 +232,11 @@ pub async fn get_endpoint( state: &AppState, endpoint_id: Uuid, ) -> Result { - sqlx::query_as::<_, StorageEndpoint>( + if let Some(endpoint) = state.storage_cache.endpoint(endpoint_id) { + return Ok(endpoint); + } + + let endpoint = sqlx::query_as::<_, StorageEndpoint>( r#" SELECT id, name, internal_endpoint, public_endpoint, bucket, region, access_key_encrypted, secret_key_encrypted, access_key_hint, @@ -97,11 +251,17 @@ pub async fn get_endpoint( .fetch_optional(&state.db) .await .map_err(|err| AppError::new(ErrorCode::Internal, "查询存储端点失败").with_source(err))? - .ok_or_else(|| AppError::new(ErrorCode::NotFound, "存储端点不存在")) + .ok_or_else(|| AppError::new(ErrorCode::NotFound, "存储端点不存在"))?; + state.storage_cache.store_endpoint(endpoint.clone()); + Ok(endpoint) } pub async fn active_endpoint(state: &AppState) -> Result, AppError> { - sqlx::query_as::<_, StorageEndpoint>( + if let Some(endpoint) = state.storage_cache.active_endpoint() { + return Ok(endpoint); + } + + let endpoint = sqlx::query_as::<_, StorageEndpoint>( r#" SELECT id, name, internal_endpoint, public_endpoint, bucket, region, access_key_encrypted, secret_key_encrypted, access_key_hint, @@ -115,7 +275,9 @@ pub async fn active_endpoint(state: &AppState) -> Result ) .fetch_optional(&state.db) .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询活动存储端点失败").with_source(err)) + .map_err(|err| AppError::new(ErrorCode::Internal, "查询活动存储端点失败").with_source(err))?; + state.storage_cache.store_active_endpoint(endpoint.clone()); + Ok(endpoint) } pub fn result_key(retention_hours: i64, task_id: Uuid, file_id: Uuid, extension: &str) -> String { @@ -564,6 +726,10 @@ fn client_for( endpoint: &StorageEndpoint, kind: EndpointKind, ) -> Result { + if let Some(client) = state.storage_cache.client(endpoint, kind) { + return Ok(client); + } + let access_key = settings::decrypt_secret(state, &endpoint.access_key_encrypted).map_err(|err| { AppError::new(ErrorCode::StorageUnavailable, "存储端点凭据不可用").with_source(err) @@ -576,13 +742,17 @@ fn client_for( EndpointKind::Internal => &endpoint.internal_endpoint, EndpointKind::Public => &endpoint.public_endpoint, }; - Ok(build_client( + let client = build_client( access_key, secret_key, endpoint_url, &endpoint.region, endpoint.force_path_style, - )) + ); + state + .storage_cache + .store_client(endpoint, kind, client.clone()); + Ok(client) } fn build_client( @@ -774,12 +944,6 @@ fn cors_allows_origin(allowed_origin: Option<&str>, expected_origin: &str) -> bo allowed_origin.is_some_and(|value| value == "*" || value == expected_origin) } -#[derive(Clone, Copy)] -enum EndpointKind { - Internal, - Public, -} - #[cfg(test)] mod tests { use super::*; @@ -860,4 +1024,61 @@ mod tests { "https://tp.workyai.cn" )); } + + #[test] + fn storage_cache_reuses_matching_clients_and_invalidates_changes() { + let cache = StorageCache::new(); + let mut endpoint = test_storage_endpoint(); + let client = build_client( + "access-key".to_string(), + "secret-key".to_string(), + &endpoint.internal_endpoint, + &endpoint.region, + endpoint.force_path_style, + ); + + cache.store_client(&endpoint, EndpointKind::Internal, client); + assert!(cache.client(&endpoint, EndpointKind::Internal).is_some()); + assert!(cache.client(&endpoint, EndpointKind::Public).is_none()); + + endpoint.internal_endpoint = "https://new-storage.example.com".to_string(); + assert!(cache.client(&endpoint, EndpointKind::Internal).is_none()); + + cache.invalidate(); + assert!(cache.client(&endpoint, EndpointKind::Internal).is_none()); + } + + #[test] + fn active_endpoint_cache_distinguishes_cached_local_backend() { + let cache = StorageCache::new(); + assert!(cache.active_endpoint().is_none()); + cache.store_active_endpoint(None); + assert!(matches!(cache.active_endpoint(), Some(None))); + cache.invalidate(); + assert!(cache.active_endpoint().is_none()); + } + + fn test_storage_endpoint() -> StorageEndpoint { + let now = Utc::now(); + StorageEndpoint { + id: Uuid::new_v4(), + name: "test".to_string(), + internal_endpoint: "https://storage.example.com".to_string(), + public_endpoint: "https://files.example.com".to_string(), + bucket: "images".to_string(), + region: "garage".to_string(), + access_key_encrypted: "encrypted-access".to_string(), + secret_key_encrypted: "encrypted-secret".to_string(), + access_key_hint: "test".to_string(), + force_path_style: true, + presign_ttl_seconds: 300, + is_active: true, + last_test_at: None, + last_test_ok: None, + last_test_error: None, + created_at: now, + updated_at: now, + deleted_at: None, + } + } } diff --git a/src/state.rs b/src/state.rs index 4b444d5..33ba26c 100644 --- a/src/state.rs +++ b/src/state.rs @@ -1,6 +1,7 @@ use crate::config::Config; use crate::services::mail::Mailer; use crate::services::settings::RuntimePolicyCache; +use crate::services::storage::StorageCache; #[derive(Clone)] pub struct AppState { @@ -10,4 +11,5 @@ pub struct AppState { pub mailer: std::sync::Arc, pub image_processing_semaphore: std::sync::Arc, pub runtime_policy_cache: RuntimePolicyCache, + pub storage_cache: StorageCache, } diff --git a/src/worker/mod.rs b/src/worker/mod.rs index 318b4e1..475c8bc 100644 --- a/src/worker/mod.rs +++ b/src/worker/mod.rs @@ -11,7 +11,7 @@ use redis::AsyncCommands; use sqlx::FromRow; use std::net::IpAddr; use std::sync::Arc; -use std::time::Instant; +use std::time::Duration; use tokio::sync::Semaphore; use tokio::task::JoinSet; use uuid::Uuid; @@ -21,6 +21,13 @@ const GROUP_NAME: &str = metrics::QUEUE_GROUP_NAME; const DEAD_STREAM_KEY: &str = metrics::DEAD_STREAM_KEY; const MAX_DELIVERIES: usize = 3; const STALE_MESSAGE_IDLE_MS: usize = 5 * 60 * 1000; +const MESSAGE_HEARTBEAT_SECONDS: u64 = 60; +const QUEUE_BLOCK_MS: usize = 1_000; +const MAINTENANCE_INTERVAL_SECONDS: u64 = 300; +const MAINTENANCE_BATCH_SIZE: i64 = 1_000; +const MAX_MAINTENANCE_BATCHES: usize = 20; +const INITIAL_RETRY_SECONDS: u64 = 2; +const MAX_RETRY_SECONDS: u64 = 30; pub async fn run(state: AppState) -> Result<(), AppError> { tracing::info!("Worker started"); @@ -29,21 +36,80 @@ pub async fn run(state: AppState) -> Result<(), AppError> { let consumer = format!("worker_{}", Uuid::new_v4()); ensure_group(&state).await?; + tokio::spawn(maintenance_loop(state.clone())); - let mut last_maintenance = Instant::now(); + let task_concurrency = state.config.worker_task_concurrency.max(1) as usize; + let mut inflight = JoinSet::new(); + let mut poll_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS); + tracing::info!(task_concurrency, "Worker task scheduler ready"); loop { - if let Err(err) = poll_once(&state, &consumer).await { - tracing::error!(error = ?err, "worker poll error"); - tokio::time::sleep(std::time::Duration::from_secs(2)).await; + while let Some(result) = inflight.try_join_next() { + log_message_task_result(result); } - if last_maintenance.elapsed().as_secs() >= 300 { - if let Err(err) = maintenance(&state).await { - tracing::error!(error = ?err, "maintenance failed"); + let available = task_concurrency.saturating_sub(inflight.len()); + if available > 0 { + match read_messages(&state, &consumer, available).await { + Ok(messages) => { + poll_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS); + for message in messages { + let state = state.clone(); + let consumer = consumer.clone(); + inflight.spawn(async move { + process_message_with_retries(state, consumer, message).await + }); + } + } + Err(err) => { + tracing::error!( + error = ?err, + retry_in_seconds = poll_backoff.as_secs(), + "worker poll error" + ); + tokio::time::sleep(poll_backoff).await; + poll_backoff = next_backoff(poll_backoff); + continue; + } } - last_maintenance = Instant::now(); } + + if inflight.len() >= task_concurrency { + if let Some(result) = inflight.join_next().await { + log_message_task_result(result); + } + } else { + tokio::select! { + result = inflight.join_next(), if !inflight.is_empty() => { + if let Some(result) = result { + log_message_task_result(result); + } + } + _ = tokio::time::sleep(Duration::from_millis(25)) => {} + } + } + } +} + +async fn maintenance_loop(state: AppState) { + let mut interval = tokio::time::interval_at( + tokio::time::Instant::now() + Duration::from_secs(MAINTENANCE_INTERVAL_SECONDS), + Duration::from_secs(MAINTENANCE_INTERVAL_SECONDS), + ); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + interval.tick().await; + if let Err(err) = maintenance(&state).await { + tracing::error!(error = ?err, "maintenance failed"); + } + } +} + +fn log_message_task_result(result: Result, tokio::task::JoinError>) { + match result { + Ok(Ok(())) => {} + Ok(Err(err)) => tracing::error!(error = ?err, "worker message task stopped"), + Err(err) => tracing::error!(error = ?err, "worker message task panicked"), } } @@ -71,51 +137,144 @@ async fn ensure_group(state: &AppState) -> Result<(), AppError> { } } -async fn poll_once(state: &AppState, consumer: &str) -> Result<(), AppError> { +async fn read_messages( + state: &AppState, + consumer: &str, + count: usize, +) -> Result, AppError> { let mut conn = state.redis.clone(); if let Some(msg) = claim_stale_message(&mut conn, consumer).await? { - return handle_message(state, &mut conn, msg).await; + return Ok(vec![msg]); } - // Retry messages already delivered to this consumer before taking new work. - let pending_opts = StreamReadOptions::default() + let opts = StreamReadOptions::default() .group(GROUP_NAME, consumer) - .count(1); - let mut reply: redis::streams::StreamReadReply = conn - .xread_options(&[STREAM_KEY], &["0"], &pending_opts) + .count(count) + .block(QUEUE_BLOCK_MS); + let reply: redis::streams::StreamReadReply = conn + .xread_options(&[STREAM_KEY], &[">"], &opts) .await - .map_err(|err| AppError::new(ErrorCode::Internal, "读取待重试任务失败").with_source(err))?; + .map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?; - if !reply.keys.iter().any(|key| !key.ids.is_empty()) { - let opts = StreamReadOptions::default() - .group(GROUP_NAME, consumer) - .count(1) - .block(5000); + Ok(reply + .keys + .into_iter() + .flat_map(|key| key.ids.into_iter()) + .collect()) +} - reply = conn - .xread_options(&[STREAM_KEY], &[">"], &opts) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?; - } +async fn process_message_with_retries( + state: AppState, + consumer: String, + mut message: StreamId, +) -> Result<(), AppError> { + loop { + match handle_message_with_heartbeat(&state, &consumer, &message).await { + Ok(()) => return Ok(()), + Err(err) => { + let mut conn = state.redis.clone(); + let deliveries = match pending_delivery_count(&mut conn, &message.id).await { + Ok(deliveries) => deliveries, + Err(count_err) => { + tracing::warn!( + message_id = %message.id, + error = ?count_err, + "failed to read delivery count; using initial retry delay" + ); + 1 + } + }; + let delay = delivery_backoff(deliveries); + tracing::warn!( + message_id = %message.id, + deliveries, + retry_in_seconds = delay.as_secs(), + error = ?err, + "worker message handling failed; retry scheduled" + ); + tokio::time::sleep(delay).await; - if reply.keys.is_empty() { - return Ok(()); - } - - for key in reply.keys { - for msg in key.ids { - handle_message(state, &mut conn, msg).await?; + let mut reclaim_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS); + loop { + match redeliver_message(&mut conn, &consumer, &message.id).await { + Ok(Some(redelivered)) => { + message = redelivered; + break; + } + Ok(None) => return Ok(()), + Err(reclaim_err) => { + tracing::error!( + message_id = %message.id, + retry_in_seconds = reclaim_backoff.as_secs(), + error = ?reclaim_err, + "failed to reclaim pending worker message" + ); + tokio::time::sleep(reclaim_backoff).await; + reclaim_backoff = next_backoff(reclaim_backoff); + } + } + } + } } } +} +async fn handle_message_with_heartbeat( + state: &AppState, + consumer: &str, + message: &StreamId, +) -> Result<(), AppError> { + let mut conn = state.redis.clone(); + let handling = handle_message(state, &mut conn, message); + tokio::pin!(handling); + let mut heartbeat = tokio::time::interval_at( + tokio::time::Instant::now() + Duration::from_secs(MESSAGE_HEARTBEAT_SECONDS), + Duration::from_secs(MESSAGE_HEARTBEAT_SECONDS), + ); + heartbeat.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + result = &mut handling => return result, + _ = heartbeat.tick() => { + if let Err(err) = touch_pending_message(state, consumer, &message.id).await { + tracing::warn!( + message_id = %message.id, + error = ?err, + "failed to refresh worker message heartbeat" + ); + } + } + } + } +} + +async fn touch_pending_message( + state: &AppState, + consumer: &str, + message_id: &str, +) -> Result<(), AppError> { + let mut conn = state.redis.clone(); + redis::cmd("XCLAIM") + .arg(STREAM_KEY) + .arg(GROUP_NAME) + .arg(consumer) + .arg(0) + .arg(message_id) + .arg("JUSTID") + .query_async::<_, redis::Value>(&mut conn) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "刷新队列任务心跳失败").with_source(err) + })?; Ok(()) } async fn handle_message( state: &AppState, conn: &mut redis::aio::ConnectionManager, - msg: StreamId, + msg: &StreamId, ) -> Result<(), AppError> { let Some(task_id_str) = msg.get::("task_id") else { ack_message(conn, &msg.id).await?; @@ -147,17 +306,25 @@ async fn handle_message( return Ok(()); } - tracing::warn!( - task_id = %task_id, - deliveries, - error = %err, - "task processing failed; message left pending for retry" - ); Err(err) } } } +async fn redeliver_message( + conn: &mut redis::aio::ConnectionManager, + consumer: &str, + message_id: &str, +) -> Result, AppError> { + let claimed: StreamClaimReply = conn + .xclaim(STREAM_KEY, GROUP_NAME, consumer, 0, &[message_id]) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "重新认领待重试任务失败").with_source(err) + })?; + Ok(claimed.ids.into_iter().next()) +} + async fn claim_stale_message( conn: &mut redis::aio::ConnectionManager, consumer: &str, @@ -212,6 +379,21 @@ fn should_dead_letter(deliveries: usize) -> bool { deliveries >= MAX_DELIVERIES } +fn delivery_backoff(deliveries: usize) -> Duration { + let exponent = deliveries.saturating_sub(1).min(4) as u32; + Duration::from_secs( + INITIAL_RETRY_SECONDS + .saturating_mul(2_u64.saturating_pow(exponent)) + .min(MAX_RETRY_SECONDS), + ) +} + +fn next_backoff(current: Duration) -> Duration { + current + .saturating_mul(2) + .min(Duration::from_secs(MAX_RETRY_SECONDS)) +} + async fn write_dead_letter( conn: &mut redis::aio::ConnectionManager, message_id: &str, @@ -1026,22 +1208,38 @@ async fn maintenance(state: &AppState) -> Result<(), AppError> { } async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> { - let task_ids: Vec = sqlx::query_scalar( - r#" - SELECT id - FROM tasks - WHERE anonymous_units_reserved > 0 - AND status IN ('completed', 'failed', 'cancelled') - ORDER BY completed_at ASC NULLS FIRST - LIMIT 200 - "#, - ) - .fetch_all(&state.db) - .await - .map_err(|err| AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err))?; + for _ in 0..MAX_MAINTENANCE_BATCHES { + let task_ids: Vec = sqlx::query_scalar( + r#" + SELECT id + FROM tasks + WHERE anonymous_units_reserved > 0 + AND status IN ('completed', 'failed', 'cancelled') + ORDER BY completed_at ASC NULLS FIRST + LIMIT $1 + "#, + ) + .bind(MAINTENANCE_BATCH_SIZE) + .fetch_all(&state.db) + .await + .map_err(|err| { + AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err) + })?; - for task_id in task_ids { - quota::settle_anonymous_task_reservation(state, task_id).await?; + let batch_len = task_ids.len(); + let mut settled = 0usize; + for task_id in task_ids { + match quota::settle_anonymous_task_reservation(state, task_id).await { + Ok(_) => settled += 1, + Err(err) => { + tracing::warn!(task_id = %task_id, error = %err, "anonymous quota settlement deferred") + } + } + } + if settled == 0 || batch_len < MAINTENANCE_BATCH_SIZE as usize { + break; + } + tokio::task::yield_now().await; } Ok(()) } @@ -1117,20 +1315,34 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> { } async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> { - let task_ids: Vec = - sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200") - .fetch_all(&state.db) - .await - .unwrap_or_default(); + for _ in 0..MAX_MAINTENANCE_BATCHES { + let task_ids: Vec = sqlx::query_scalar( + "SELECT id FROM tasks WHERE expires_at < NOW() ORDER BY expires_at ASC LIMIT $1", + ) + .bind(MAINTENANCE_BATCH_SIZE) + .fetch_all(&state.db) + .await + .map_err(|err| AppError::new(ErrorCode::Internal, "查询过期任务失败").with_source(err))?; - if task_ids.is_empty() { - return Ok(()); - } - - for task_id in task_ids { - if let Err(err) = cleanup_expired_task(state, task_id).await { - tracing::warn!(task_id = %task_id, error = %err, "expired task cleanup deferred"); + if task_ids.is_empty() { + break; } + + let batch_len = task_ids.len(); + let mut cleaned = 0usize; + for task_id in task_ids { + match cleanup_expired_task(state, task_id).await { + Ok(()) => cleaned += 1, + Err(err) => { + tracing::warn!(task_id = %task_id, error = %err, "expired task cleanup deferred") + } + } + } + + if cleaned == 0 || batch_len < MAINTENANCE_BATCH_SIZE as usize { + break; + } + tokio::task::yield_now().await; } Ok(()) @@ -1232,4 +1444,21 @@ mod tests { assert!(should_dead_letter(3)); assert!(should_dead_letter(10)); } + + #[test] + fn retry_backoff_is_exponential_and_capped() { + assert_eq!(delivery_backoff(1), Duration::from_secs(2)); + assert_eq!(delivery_backoff(2), Duration::from_secs(4)); + assert_eq!(delivery_backoff(3), Duration::from_secs(8)); + assert_eq!(delivery_backoff(10), Duration::from_secs(30)); + assert_eq!( + next_backoff(Duration::from_secs(30)), + Duration::from_secs(30) + ); + } + + #[test] + fn message_heartbeat_precedes_stale_claim_threshold() { + assert!(MESSAGE_HEARTBEAT_SECONDS * 1_000 < STALE_MESSAGE_IDLE_MS as u64); + } }