perf: improve worker and storage throughput
Some checks failed
CI / verify (push) Has been cancelled

This commit is contained in:
237899745
2026-07-26 00:57:29 +08:00
parent da9b273253
commit de5f451cd1
14 changed files with 611 additions and 101 deletions

View File

@@ -9,12 +9,15 @@ RUST_LOG=info,tower_http=info,imageforge=debug
# 数据库 # 数据库
DATABASE_URL=postgres://imageforge:devpassword@localhost:5432/imageforge DATABASE_URL=postgres://imageforge:devpassword@localhost:5432/imageforge
DATABASE_MAX_CONNECTIONS=10 DATABASE_MAX_CONNECTIONS=16
# Redis # Redis
REDIS_URL=redis://localhost:6379 REDIS_URL=redis://localhost:6379
# Worker 并发(每个批量任务内同时处理的文件数) # Worker 同时处理的批量任务数
WORKER_TASK_CONCURRENCY=4
# 每个批量任务内同时处理的文件数
WORKER_CONCURRENCY=4 WORKER_CONCURRENCY=4
# 单进程图片处理并发上限API 与 Worker 均生效,默认等于 CPU 线程数) # 单进程图片处理并发上限API 与 Worker 均生效,默认等于 CPU 线程数)

View File

@@ -21,10 +21,11 @@ ADMIN_EMAIL=admin@example.com
ADMIN_USERNAME=admin ADMIN_USERNAME=admin
ADMIN_PASSWORD=replace-with-a-strong-admin-password ADMIN_PASSWORD=replace-with-a-strong-admin-password
# A four-core host should start with two image jobs per process. # 8-core starting point: four tasks, two files per task, four CPU-bound image jobs.
DATABASE_MAX_CONNECTIONS=10 DATABASE_MAX_CONNECTIONS=16
WORKER_TASK_CONCURRENCY=4
WORKER_CONCURRENCY=2 WORKER_CONCURRENCY=2
IMAGE_PROCESSING_CONCURRENCY=2 IMAGE_PROCESSING_CONCURRENCY=4
# Resource ceilings tuned for an 8-core / 16 GB application host. # Resource ceilings tuned for an 8-core / 16 GB application host.
POSTGRES_MEMORY_LIMIT=2g POSTGRES_MEMORY_LIMIT=2g

View File

@@ -2,7 +2,7 @@ name: imageforge
x-imageforge-environment: &imageforge-environment x-imageforge-environment: &imageforge-environment
DATABASE_URL: postgres://imageforge:${POSTGRES_PASSWORD:?POSTGRES_PASSWORD is required}@postgres:5432/imageforge 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 REDIS_URL: redis://redis:6379
JWT_SECRET: ${JWT_SECRET:?JWT_SECRET is required} JWT_SECRET: ${JWT_SECRET:?JWT_SECRET is required}
JWT_EXPIRY_HOURS: ${JWT_EXPIRY_HOURS:-168} JWT_EXPIRY_HOURS: ${JWT_EXPIRY_HOURS:-168}
@@ -10,6 +10,7 @@ x-imageforge-environment: &imageforge-environment
BILLING_PROVIDER: ${BILLING_PROVIDER:-stripe} BILLING_PROVIDER: ${BILLING_PROVIDER:-stripe}
STORAGE_PATH: /app/uploads STORAGE_PATH: /app/uploads
PUBLIC_BASE_URL: ${PUBLIC_BASE_URL:-http://localhost:8080} PUBLIC_BASE_URL: ${PUBLIC_BASE_URL:-http://localhost:8080}
WORKER_TASK_CONCURRENCY: ${WORKER_TASK_CONCURRENCY:-4}
WORKER_CONCURRENCY: ${WORKER_CONCURRENCY:-2} WORKER_CONCURRENCY: ${WORKER_CONCURRENCY:-2}
IMAGE_PROCESSING_CONCURRENCY: ${IMAGE_PROCESSING_CONCURRENCY:-2} IMAGE_PROCESSING_CONCURRENCY: ${IMAGE_PROCESSING_CONCURRENCY:-2}
ALLOW_ANONYMOUS_UPLOAD: ${ALLOW_ANONYMOUS_UPLOAD:-true} ALLOW_ANONYMOUS_UPLOAD: ${ALLOW_ANONYMOUS_UPLOAD:-true}

View File

@@ -276,7 +276,8 @@ dotenvy = "0.15"
### 1. 并发处理 ### 1. 并发处理
- 使用 Tokio 异步运行时 - 使用 Tokio 异步运行时
- 图片压缩使用 `spawn_blocking` 避免阻塞异步线程 - 图片压缩使用 `spawn_blocking` 避免阻塞异步线程
- 可配置 Worker 线程数 - `WORKER_TASK_CONCURRENCY` 控制任务级并发,避免大批量任务独占 Worker
- `WORKER_CONCURRENCY` 控制单任务内文件并发,`IMAGE_PROCESSING_CONCURRENCY` 作为进程级 CPU 闸门
```rust ```rust
// 在独立线程池中执行 CPU 密集型压缩 // 在独立线程池中执行 CPU 密集型压缩
@@ -292,6 +293,7 @@ let result = tokio::task::spawn_blocking(move || {
### 3. 缓存策略 ### 3. 缓存策略
- Redis 缓存用户会话 - Redis 缓存用户会话
- S3 活动端点与历史端点缓存 5 秒AWS SDK Client 按端点和内外网地址复用连接池
- 可选:相同图片哈希缓存结果(去重) - 可选:相同图片哈希缓存结果(去重)
## 安全考虑 ## 安全考虑

View File

@@ -540,7 +540,7 @@ Dead-letter stream: stream:compress_jobs:dead
Dead-letter max length: 10000 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 任务进度(可选) ### 5.4 任务进度(可选)
``` ```

View File

@@ -11,7 +11,7 @@
- 最低 2 核 CPU、4GB 内存;启用 AVIF 和独立 Worker 时建议 4 核、8GB 内存 - 最低 2 核 CPU、4GB 内存;启用 AVIF 和独立 Worker 时建议 4 核、8GB 内存
- 首次构建可访问 Docker Hub 与 crates.io - 首次构建可访问 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达到上限后返回写入错误而不是继续挤占宿主机内存。 生产 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 df -h
``` ```
若图片压缩长时间排队,先检查 CPU`IMAGE_PROCESSING_CONCURRENCY`;若 API 健康但批量任务不推进,检查 Worker 日志Redis 状态 若图片压缩长时间排队,先检查 CPU再调 `IMAGE_PROCESSING_CONCURRENCY`大任务阻塞小任务时提高 `WORKER_TASK_CONCURRENCY`,单个批量任务推进过慢时再评估 `WORKER_CONCURRENCY`若 API 健康但批量任务不推进,检查 Worker 日志Redis pending 数和数据库连接池等待情况

View File

@@ -59,6 +59,8 @@ curl --fail http://127.0.0.1:18180/metrics
- `imageforge_storage_fallbacks_total` - `imageforge_storage_fallbacks_total`
- `imageforge_dead_letters_total` - `imageforge_dead_letters_total`
`imageforge_storage_fallbacks_total` 是存储容量风险信号,而不仅是普通降级统计。建议对 `increase(imageforge_storage_fallbacks_total[5m]) > 0` 持续 5 分钟设置高优先级告警,并同步监控应用服务器 `uploads` 卷使用率,避免 S3 长时间不可用时本地回退写满磁盘。
Prometheus 抓取示例: Prometheus 抓取示例:
```yaml ```yaml

View File

@@ -0,0 +1,2 @@
CREATE INDEX IF NOT EXISTS idx_usage_events_occurred_at
ON usage_events(occurred_at DESC);

View File

@@ -176,6 +176,7 @@ async fn create_storage_endpoint(
.fetch_one(&state.db) .fetch_one(&state.db)
.await .await
.map_err(|err| AppError::new(ErrorCode::Internal, "创建存储端点失败").with_source(err))?; .map_err(|err| AppError::new(ErrorCode::Internal, "创建存储端点失败").with_source(err))?;
state.storage_cache.invalidate();
audit_storage_action( audit_storage_action(
&state, &state,
@@ -336,6 +337,7 @@ async fn update_storage_endpoint(
"端点配置在测试期间发生变化,请重新编辑", "端点配置在测试期间发生变化,请重新编辑",
) )
})?; })?;
state.storage_cache.invalidate();
audit_storage_action( audit_storage_action(
&state, &state,
@@ -364,9 +366,18 @@ async fn test_storage_endpoint(
let (_jar, admin_id) = require_admin(&state, jar, &headers, ip).await?; let (_jar, admin_id) = require_admin(&state, jar, &headers, ip).await?;
let endpoint = storage::get_endpoint(&state, endpoint_id).await?; let endpoint = storage::get_endpoint(&state, endpoint_id).await?;
ensure_configurable(&endpoint)?; ensure_configurable(&endpoint)?;
let tested_config_updated_at = endpoint.updated_at;
if let Err(err) = storage::test_endpoint(&state, &endpoint).await { if let Err(err) = storage::test_endpoint(&state, &endpoint).await {
let detail = storage_test_message(&err); 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( audit_storage_action(
&state, &state,
admin_id, admin_id,
@@ -379,7 +390,15 @@ async fn test_storage_endpoint(
return Err(AppError::new(ErrorCode::StorageUnavailable, detail)); 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?; let endpoint = storage::get_endpoint(&state, endpoint_id).await?;
ensure_configurable(&endpoint)?; ensure_configurable(&endpoint)?;
audit_storage_action( audit_storage_action(
@@ -416,7 +435,15 @@ async fn activate_storage_endpoint(
if let Err(err) = storage::test_endpoint(&state, &endpoint).await { if let Err(err) = storage::test_endpoint(&state, &endpoint).await {
let detail = storage_test_message(&err); 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)); return Err(AppError::new(ErrorCode::StorageUnavailable, detail));
} }
@@ -456,6 +483,7 @@ async fn activate_storage_endpoint(
tx.commit().await.map_err(|err| { tx.commit().await.map_err(|err| {
AppError::new(ErrorCode::Internal, "提交存储切换事务失败").with_source(err) AppError::new(ErrorCode::Internal, "提交存储切换事务失败").with_source(err)
})?; })?;
state.storage_cache.invalidate();
let endpoint = storage::get_endpoint(&state, endpoint_id).await?; let endpoint = storage::get_endpoint(&state, endpoint_id).await?;
audit_storage_action( audit_storage_action(
@@ -505,6 +533,7 @@ async fn delete_storage_endpoint(
"端点配置已发生变化,请刷新后重试", "端点配置已发生变化,请刷新后重试",
)); ));
} }
state.storage_cache.invalidate();
audit_storage_action( audit_storage_action(
&state, &state,
admin_id, admin_id,
@@ -624,8 +653,9 @@ async fn record_test_result(
ok: bool, ok: bool,
error: Option<&str>, error: Option<&str>,
admin_id: Uuid, admin_id: Uuid,
expected_updated_at: DateTime<Utc>,
) -> Result<(), AppError> { ) -> Result<(), AppError> {
sqlx::query( let updated = sqlx::query(
r#" r#"
UPDATE storage_endpoints UPDATE storage_endpoints
SET last_test_at = NOW(), SET last_test_at = NOW(),
@@ -633,16 +663,25 @@ async fn record_test_result(
last_test_error = $3, last_test_error = $3,
updated_at = NOW(), updated_at = NOW(),
updated_by = $4 updated_by = $4
WHERE id = $1 WHERE id = $1 AND updated_at = $5 AND deleted_at IS NULL
"#, "#,
) )
.bind(endpoint_id) .bind(endpoint_id)
.bind(ok) .bind(ok)
.bind(error) .bind(error)
.bind(admin_id) .bind(admin_id)
.bind(expected_updated_at)
.execute(&state.db) .execute(&state.db)
.await .await
.map_err(|err| AppError::new(ErrorCode::Internal, "记录存储测试结果失败").with_source(err))?; .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(()) Ok(())
} }

View File

@@ -14,6 +14,7 @@ pub struct Config {
pub redis_url: String, pub redis_url: String,
pub worker_task_concurrency: u32,
pub worker_concurrency: u32, pub worker_concurrency: u32,
pub image_processing_concurrency: u32, pub image_processing_concurrency: u32,
@@ -62,11 +63,16 @@ impl Config {
let redis_url = env_string("REDIS_URL") let redis_url = env_string("REDIS_URL")
.ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 REDIS_URL"))?; .ok_or_else(|| AppError::new(ErrorCode::InvalidRequest, "缺少环境变量 REDIS_URL"))?;
let worker_concurrency = env_u32("WORKER_CONCURRENCY").unwrap_or_else(|| { let worker_task_concurrency = env_u32("WORKER_TASK_CONCURRENCY")
std::thread::available_parallelism() .filter(|value| *value > 0)
.map(|v| v.get() as u32) .unwrap_or(4);
.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") let image_processing_concurrency = env_u32("IMAGE_PROCESSING_CONCURRENCY")
.filter(|value| *value > 0) .filter(|value| *value > 0)
.unwrap_or_else(|| { .unwrap_or_else(|| {
@@ -126,6 +132,7 @@ impl Config {
database_url, database_url,
database_max_connections, database_max_connections,
redis_url, redis_url,
worker_task_concurrency,
worker_concurrency, worker_concurrency,
image_processing_concurrency, image_processing_concurrency,
jwt_secret, jwt_secret,

View File

@@ -44,6 +44,7 @@ async fn main() -> Result<(), AppError> {
mailer: std::sync::Arc::new(mailer), mailer: std::sync::Arc::new(mailer),
image_processing_semaphore, image_processing_semaphore,
runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(), runtime_policy_cache: crate::services::settings::RuntimePolicyCache::new(),
storage_cache: crate::services::storage::StorageCache::new(),
}; };
match state.config.role.as_str() { match state.config.role.as_str() {

View File

@@ -13,14 +13,17 @@ use aws_sdk_s3::Client;
use bytes::Bytes; use bytes::Bytes;
use chrono::{DateTime, Datelike, Utc}; use chrono::{DateTime, Datelike, Utc};
use sqlx::FromRow; use sqlx::FromRow;
use std::collections::HashMap;
use std::path::{Path, PathBuf}; 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 tokio::io::{AsyncReadExt, AsyncWriteExt};
use url::Url; use url::Url;
use uuid::Uuid; use uuid::Uuid;
const MULTIPART_THRESHOLD: u64 = 64 * 1024 * 1024; const MULTIPART_THRESHOLD: u64 = 64 * 1024 * 1024;
const MULTIPART_PART_SIZE: usize = 16 * 1024 * 1024; const MULTIPART_PART_SIZE: usize = 16 * 1024 * 1024;
const STORAGE_ENDPOINT_CACHE_TTL: Duration = Duration::from_secs(5);
#[derive(Debug, Clone, FromRow)] #[derive(Debug, Clone, FromRow)]
pub struct StorageEndpoint { pub struct StorageEndpoint {
@@ -60,6 +63,153 @@ pub struct ObjectLocator {
pub key: String, pub key: String,
} }
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum EndpointKind {
Internal,
Public,
}
#[derive(Clone, Default)]
pub struct StorageCache {
inner: Arc<RwLock<StorageCacheState>>,
}
#[derive(Default)]
struct StorageCacheState {
active_endpoint: Option<CachedActiveEndpoint>,
endpoints: HashMap<Uuid, CachedEndpoint>,
clients: HashMap<(Uuid, EndpointKind), CachedClient>,
}
struct CachedActiveEndpoint {
loaded_at: Instant,
endpoint: Option<StorageEndpoint>,
}
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<Option<StorageEndpoint>> {
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<StorageEndpoint>) {
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<StorageEndpoint> {
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<Client> {
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<Vec<StorageEndpoint>, AppError> { pub async fn list_endpoints(state: &AppState) -> Result<Vec<StorageEndpoint>, AppError> {
sqlx::query_as::<_, StorageEndpoint>( sqlx::query_as::<_, StorageEndpoint>(
r#" r#"
@@ -82,7 +232,11 @@ pub async fn get_endpoint(
state: &AppState, state: &AppState,
endpoint_id: Uuid, endpoint_id: Uuid,
) -> Result<StorageEndpoint, AppError> { ) -> Result<StorageEndpoint, AppError> {
sqlx::query_as::<_, StorageEndpoint>( if let Some(endpoint) = state.storage_cache.endpoint(endpoint_id) {
return Ok(endpoint);
}
let endpoint = sqlx::query_as::<_, StorageEndpoint>(
r#" r#"
SELECT id, name, internal_endpoint, public_endpoint, bucket, region, SELECT id, name, internal_endpoint, public_endpoint, bucket, region,
access_key_encrypted, secret_key_encrypted, access_key_hint, access_key_encrypted, secret_key_encrypted, access_key_hint,
@@ -97,11 +251,17 @@ pub async fn get_endpoint(
.fetch_optional(&state.db) .fetch_optional(&state.db)
.await .await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询存储端点失败").with_source(err))? .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<Option<StorageEndpoint>, AppError> { pub async fn active_endpoint(state: &AppState) -> Result<Option<StorageEndpoint>, AppError> {
sqlx::query_as::<_, StorageEndpoint>( if let Some(endpoint) = state.storage_cache.active_endpoint() {
return Ok(endpoint);
}
let endpoint = sqlx::query_as::<_, StorageEndpoint>(
r#" r#"
SELECT id, name, internal_endpoint, public_endpoint, bucket, region, SELECT id, name, internal_endpoint, public_endpoint, bucket, region,
access_key_encrypted, secret_key_encrypted, access_key_hint, access_key_encrypted, secret_key_encrypted, access_key_hint,
@@ -115,7 +275,9 @@ pub async fn active_endpoint(state: &AppState) -> Result<Option<StorageEndpoint>
) )
.fetch_optional(&state.db) .fetch_optional(&state.db)
.await .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 { 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, endpoint: &StorageEndpoint,
kind: EndpointKind, kind: EndpointKind,
) -> Result<Client, AppError> { ) -> Result<Client, AppError> {
if let Some(client) = state.storage_cache.client(endpoint, kind) {
return Ok(client);
}
let access_key = let access_key =
settings::decrypt_secret(state, &endpoint.access_key_encrypted).map_err(|err| { settings::decrypt_secret(state, &endpoint.access_key_encrypted).map_err(|err| {
AppError::new(ErrorCode::StorageUnavailable, "存储端点凭据不可用").with_source(err) AppError::new(ErrorCode::StorageUnavailable, "存储端点凭据不可用").with_source(err)
@@ -576,13 +742,17 @@ fn client_for(
EndpointKind::Internal => &endpoint.internal_endpoint, EndpointKind::Internal => &endpoint.internal_endpoint,
EndpointKind::Public => &endpoint.public_endpoint, EndpointKind::Public => &endpoint.public_endpoint,
}; };
Ok(build_client( let client = build_client(
access_key, access_key,
secret_key, secret_key,
endpoint_url, endpoint_url,
&endpoint.region, &endpoint.region,
endpoint.force_path_style, endpoint.force_path_style,
)) );
state
.storage_cache
.store_client(endpoint, kind, client.clone());
Ok(client)
} }
fn build_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) allowed_origin.is_some_and(|value| value == "*" || value == expected_origin)
} }
#[derive(Clone, Copy)]
enum EndpointKind {
Internal,
Public,
}
#[cfg(test)] #[cfg(test)]
mod tests { mod tests {
use super::*; use super::*;
@@ -860,4 +1024,61 @@ mod tests {
"https://tp.workyai.cn" "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,
}
}
} }

View File

@@ -1,6 +1,7 @@
use crate::config::Config; use crate::config::Config;
use crate::services::mail::Mailer; use crate::services::mail::Mailer;
use crate::services::settings::RuntimePolicyCache; use crate::services::settings::RuntimePolicyCache;
use crate::services::storage::StorageCache;
#[derive(Clone)] #[derive(Clone)]
pub struct AppState { pub struct AppState {
@@ -10,4 +11,5 @@ pub struct AppState {
pub mailer: std::sync::Arc<Mailer>, pub mailer: std::sync::Arc<Mailer>,
pub image_processing_semaphore: std::sync::Arc<tokio::sync::Semaphore>, pub image_processing_semaphore: std::sync::Arc<tokio::sync::Semaphore>,
pub runtime_policy_cache: RuntimePolicyCache, pub runtime_policy_cache: RuntimePolicyCache,
pub storage_cache: StorageCache,
} }

View File

@@ -11,7 +11,7 @@ use redis::AsyncCommands;
use sqlx::FromRow; use sqlx::FromRow;
use std::net::IpAddr; use std::net::IpAddr;
use std::sync::Arc; use std::sync::Arc;
use std::time::Instant; use std::time::Duration;
use tokio::sync::Semaphore; use tokio::sync::Semaphore;
use tokio::task::JoinSet; use tokio::task::JoinSet;
use uuid::Uuid; 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 DEAD_STREAM_KEY: &str = metrics::DEAD_STREAM_KEY;
const MAX_DELIVERIES: usize = 3; const MAX_DELIVERIES: usize = 3;
const STALE_MESSAGE_IDLE_MS: usize = 5 * 60 * 1000; 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> { pub async fn run(state: AppState) -> Result<(), AppError> {
tracing::info!("Worker started"); tracing::info!("Worker started");
@@ -29,21 +36,80 @@ pub async fn run(state: AppState) -> Result<(), AppError> {
let consumer = format!("worker_{}", Uuid::new_v4()); let consumer = format!("worker_{}", Uuid::new_v4());
ensure_group(&state).await?; 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 { loop {
if let Err(err) = poll_once(&state, &consumer).await { while let Some(result) = inflight.try_join_next() {
tracing::error!(error = ?err, "worker poll error"); log_message_task_result(result);
tokio::time::sleep(std::time::Duration::from_secs(2)).await;
} }
if last_maintenance.elapsed().as_secs() >= 300 { let available = task_concurrency.saturating_sub(inflight.len());
if let Err(err) = maintenance(&state).await { if available > 0 {
tracing::error!(error = ?err, "maintenance failed"); 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<Result<(), AppError>, 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<Vec<StreamId>, AppError> {
let mut conn = state.redis.clone(); let mut conn = state.redis.clone();
if let Some(msg) = claim_stale_message(&mut conn, consumer).await? { 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 opts = StreamReadOptions::default()
let pending_opts = StreamReadOptions::default()
.group(GROUP_NAME, consumer) .group(GROUP_NAME, consumer)
.count(1); .count(count)
let mut reply: redis::streams::StreamReadReply = conn .block(QUEUE_BLOCK_MS);
.xread_options(&[STREAM_KEY], &["0"], &pending_opts) let reply: redis::streams::StreamReadReply = conn
.xread_options(&[STREAM_KEY], &[">"], &opts)
.await .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()) { Ok(reply
let opts = StreamReadOptions::default() .keys
.group(GROUP_NAME, consumer) .into_iter()
.count(1) .flat_map(|key| key.ids.into_iter())
.block(5000); .collect())
}
reply = conn async fn process_message_with_retries(
.xread_options(&[STREAM_KEY], &[">"], &opts) state: AppState,
.await consumer: String,
.map_err(|err| AppError::new(ErrorCode::Internal, "读取队列失败").with_source(err))?; 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() { let mut reclaim_backoff = Duration::from_secs(INITIAL_RETRY_SECONDS);
return Ok(()); loop {
} match redeliver_message(&mut conn, &consumer, &message.id).await {
Ok(Some(redelivered)) => {
for key in reply.keys { message = redelivered;
for msg in key.ids { break;
handle_message(state, &mut conn, msg).await?; }
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(()) Ok(())
} }
async fn handle_message( async fn handle_message(
state: &AppState, state: &AppState,
conn: &mut redis::aio::ConnectionManager, conn: &mut redis::aio::ConnectionManager,
msg: StreamId, msg: &StreamId,
) -> Result<(), AppError> { ) -> Result<(), AppError> {
let Some(task_id_str) = msg.get::<String>("task_id") else { let Some(task_id_str) = msg.get::<String>("task_id") else {
ack_message(conn, &msg.id).await?; ack_message(conn, &msg.id).await?;
@@ -147,17 +306,25 @@ async fn handle_message(
return Ok(()); return Ok(());
} }
tracing::warn!(
task_id = %task_id,
deliveries,
error = %err,
"task processing failed; message left pending for retry"
);
Err(err) Err(err)
} }
} }
} }
async fn redeliver_message(
conn: &mut redis::aio::ConnectionManager,
consumer: &str,
message_id: &str,
) -> Result<Option<StreamId>, 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( async fn claim_stale_message(
conn: &mut redis::aio::ConnectionManager, conn: &mut redis::aio::ConnectionManager,
consumer: &str, consumer: &str,
@@ -212,6 +379,21 @@ fn should_dead_letter(deliveries: usize) -> bool {
deliveries >= MAX_DELIVERIES 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( async fn write_dead_letter(
conn: &mut redis::aio::ConnectionManager, conn: &mut redis::aio::ConnectionManager,
message_id: &str, message_id: &str,
@@ -1026,22 +1208,38 @@ async fn maintenance(state: &AppState) -> Result<(), AppError> {
} }
async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> { async fn settle_finished_anonymous_reservations(state: &AppState) -> Result<(), AppError> {
let task_ids: Vec<Uuid> = sqlx::query_scalar( for _ in 0..MAX_MAINTENANCE_BATCHES {
r#" let task_ids: Vec<Uuid> = sqlx::query_scalar(
SELECT id r#"
FROM tasks SELECT id
WHERE anonymous_units_reserved > 0 FROM tasks
AND status IN ('completed', 'failed', 'cancelled') WHERE anonymous_units_reserved > 0
ORDER BY completed_at ASC NULLS FIRST AND status IN ('completed', 'failed', 'cancelled')
LIMIT 200 ORDER BY completed_at ASC NULLS FIRST
"#, LIMIT $1
) "#,
.fetch_all(&state.db) )
.await .bind(MAINTENANCE_BATCH_SIZE)
.map_err(|err| AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err))?; .fetch_all(&state.db)
.await
.map_err(|err| {
AppError::new(ErrorCode::Internal, "查询待结算匿名任务失败").with_source(err)
})?;
for task_id in task_ids { let batch_len = task_ids.len();
quota::settle_anonymous_task_reservation(state, task_id).await?; 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(()) Ok(())
} }
@@ -1117,20 +1315,34 @@ async fn cleanup_expired_records(state: &AppState) -> Result<(), AppError> {
} }
async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> { async fn cleanup_expired_tasks(state: &AppState) -> Result<(), AppError> {
let task_ids: Vec<Uuid> = for _ in 0..MAX_MAINTENANCE_BATCHES {
sqlx::query_scalar("SELECT id FROM tasks WHERE expires_at < NOW() LIMIT 200") let task_ids: Vec<Uuid> = sqlx::query_scalar(
.fetch_all(&state.db) "SELECT id FROM tasks WHERE expires_at < NOW() ORDER BY expires_at ASC LIMIT $1",
.await )
.unwrap_or_default(); .bind(MAINTENANCE_BATCH_SIZE)
.fetch_all(&state.db)
.await
.map_err(|err| AppError::new(ErrorCode::Internal, "查询过期任务失败").with_source(err))?;
if task_ids.is_empty() { if task_ids.is_empty() {
return Ok(()); break;
}
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");
} }
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(()) Ok(())
@@ -1232,4 +1444,21 @@ mod tests {
assert!(should_dead_letter(3)); assert!(should_dead_letter(3));
assert!(should_dead_letter(10)); 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);
}
} }