feat: add local fallback for S3 writes

2026-07-25 16:29:05 +08:00
parent 63744c6d2f
commit 0fe3d4ce8e
11 changed files with 169 additions and 58 deletions
--- a/src/api/admin_storage.rs
+++ b/src/api/admin_storage.rs
@@ -60,6 +60,8 @@ struct StorageEndpointView {
 #[derive(Debug, Serialize)]
 struct StorageEndpointsResponse {
    active_backend: String,
+    local_object_count: i64,
+    local_stored_bytes: i64,
    endpoints: Vec<StorageEndpointView>,
 }

@@ -109,6 +111,7 @@ async fn list_storage_endpoints(
    } else {
        "local"
    };
+    let (local_object_count, local_stored_bytes) = local_usage(&state).await?;
    let mut views = Vec::with_capacity(endpoints.len());
    for endpoint in endpoints {
        views.push(endpoint_view(&state, endpoint).await?);
@@ -118,6 +121,8 @@ async fn list_storage_endpoints(
        success: true,
        data: StorageEndpointsResponse {
            active_backend: active_backend.to_string(),
+            local_object_count,
+            local_stored_bytes,
            endpoints: views,
        },
    }))
@@ -536,6 +541,28 @@ async fn endpoint_usage(state: &AppState, endpoint_id: Uuid) -> Result<(i64, i64
    .map_err(|err| AppError::new(ErrorCode::Internal, "统计存储使用量失败").with_source(err))
 }

+async fn local_usage(state: &AppState) -> Result<(i64, i64), AppError> {
+    sqlx::query_as::<_, (i64, i64)>(
+        r#"
+        SELECT
+          (SELECT COUNT(*) FROM task_files
+             WHERE status = 'completed'
+               AND storage_backend = 'local'
+               AND COALESCE(storage_key, storage_path) IS NOT NULL)
+            + (SELECT COUNT(*) FROM tasks
+                 WHERE zip_storage_backend = 'local'
+                   AND zip_storage_key IS NOT NULL) AS object_count,
+          COALESCE((SELECT SUM(compressed_size)::BIGINT FROM task_files
+                      WHERE status = 'completed' AND storage_backend = 'local'), 0::BIGINT)
+            + COALESCE((SELECT SUM(zip_size)::BIGINT FROM tasks
+                          WHERE zip_storage_backend = 'local'), 0::BIGINT) AS stored_bytes
+        "#,
+    )
+    .fetch_one(&state.db)
+    .await
+    .map_err(|err| AppError::new(ErrorCode::Internal, "统计本地存储使用量失败").with_source(err))
+}
+
 async fn record_test_result(
    state: &AppState,
    endpoint_id: Uuid,
--- a/src/services/storage.rs
+++ b/src/services/storage.rs
@@ -3,12 +3,14 @@ use crate::services::settings;
 use crate::state::AppState;

 use aws_sdk_s3::config::{
-    BehaviorVersion, Credentials, Region, RequestChecksumCalculation, ResponseChecksumValidation,
+    retry::RetryConfig, timeout::TimeoutConfig, BehaviorVersion, Credentials, Region,
+    RequestChecksumCalculation, ResponseChecksumValidation,
 };
 use aws_sdk_s3::presigning::PresigningConfig;
 use aws_sdk_s3::primitives::ByteStream;
 use aws_sdk_s3::types::{CompletedMultipartUpload, CompletedPart};
 use aws_sdk_s3::Client;
+use bytes::Bytes;
 use chrono::{DateTime, Datelike, Utc};
 use sqlx::FromRow;
 use std::path::{Path, PathBuf};
@@ -151,28 +153,50 @@ pub async fn store_bytes(
    bytes: Vec<u8>,
    content_type: &str,
 ) -> Result<StoredObject, AppError> {
+    let bytes = Bytes::from(bytes);
    if let Some(endpoint) = active_endpoint(state).await? {
-        let client = client_for(state, &endpoint, EndpointKind::Internal)?;
-        let size = bytes.len() as u64;
-        let output = client
-            .put_object()
-            .bucket(&endpoint.bucket)
-            .key(key)
-            .content_type(content_type)
-            .body(ByteStream::from(bytes))
-            .send()
-            .await
-            .map_err(|err| storage_error("上传 S3 对象失败", err))?;
-
-        return Ok(StoredObject {
-            backend: "s3".to_string(),
-            endpoint_id: Some(endpoint.id),
-            key: key.to_string(),
-            etag: output.e_tag().map(ToOwned::to_owned),
-            size,
-        });
+        match store_bytes_s3(state, &endpoint, key, bytes.clone(), content_type).await {
+            Ok(stored) => return Ok(stored),
+            Err(err) => log_local_fallback(&endpoint, key, &err),
+        }
    }

+    store_bytes_local(state, key, bytes.as_ref()).await
+}
+
+async fn store_bytes_s3(
+    state: &AppState,
+    endpoint: &StorageEndpoint,
+    key: &str,
+    bytes: Bytes,
+    content_type: &str,
+) -> Result<StoredObject, AppError> {
+    let client = client_for(state, endpoint, EndpointKind::Internal)?;
+    let size = bytes.len() as u64;
+    let output = client
+        .put_object()
+        .bucket(&endpoint.bucket)
+        .key(key)
+        .content_type(content_type)
+        .body(ByteStream::from(bytes))
+        .send()
+        .await
+        .map_err(|err| storage_error("上传 S3 对象失败", err))?;
+
+    Ok(StoredObject {
+        backend: "s3".to_string(),
+        endpoint_id: Some(endpoint.id),
+        key: key.to_string(),
+        etag: output.e_tag().map(ToOwned::to_owned),
+        size,
+    })
+}
+
+async fn store_bytes_local(
+    state: &AppState,
+    key: &str,
+    bytes: &[u8],
+) -> Result<StoredObject, AppError> {
    let path = local_path(state, key)?;
    if let Some(parent) = path.parent() {
        tokio::fs::create_dir_all(parent).await.map_err(|err| {
@@ -203,35 +227,58 @@ pub async fn store_file(
    })?;

    if let Some(endpoint) = active_endpoint(state).await? {
-        let client = client_for(state, &endpoint, EndpointKind::Internal)?;
-        let etag = if metadata.len() >= MULTIPART_THRESHOLD {
-            multipart_upload(&client, &endpoint.bucket, key, path, content_type).await?
-        } else {
-            let body = ByteStream::from_path(path).await.map_err(|err| {
-                AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err)
-            })?;
-            client
-                .put_object()
-                .bucket(&endpoint.bucket)
-                .key(key)
-                .content_type(content_type)
-                .body(body)
-                .send()
-                .await
-                .map_err(|err| storage_error("上传 S3 对象失败", err))?
-                .e_tag()
-                .map(ToOwned::to_owned)
-        };
-
-        return Ok(StoredObject {
-            backend: "s3".to_string(),
-            endpoint_id: Some(endpoint.id),
-            key: key.to_string(),
-            etag,
-            size: metadata.len(),
-        });
+        match store_file_s3(state, &endpoint, key, path, content_type, metadata.len()).await {
+            Ok(stored) => return Ok(stored),
+            Err(err) => log_local_fallback(&endpoint, key, &err),
+        }
    }

+    store_file_local(state, key, path, metadata.len()).await
+}
+
+async fn store_file_s3(
+    state: &AppState,
+    endpoint: &StorageEndpoint,
+    key: &str,
+    path: &Path,
+    content_type: &str,
+    size: u64,
+) -> Result<StoredObject, AppError> {
+    let client = client_for(state, endpoint, EndpointKind::Internal)?;
+    let etag = if size >= MULTIPART_THRESHOLD {
+        multipart_upload(&client, &endpoint.bucket, key, path, content_type).await?
+    } else {
+        let body = ByteStream::from_path(path).await.map_err(|err| {
+            AppError::new(ErrorCode::StorageUnavailable, "读取待上传文件失败").with_source(err)
+        })?;
+        client
+            .put_object()
+            .bucket(&endpoint.bucket)
+            .key(key)
+            .content_type(content_type)
+            .body(body)
+            .send()
+            .await
+            .map_err(|err| storage_error("上传 S3 对象失败", err))?
+            .e_tag()
+            .map(ToOwned::to_owned)
+    };
+
+    Ok(StoredObject {
+        backend: "s3".to_string(),
+        endpoint_id: Some(endpoint.id),
+        key: key.to_string(),
+        etag,
+        size,
+    })
+}
+
+async fn store_file_local(
+    state: &AppState,
+    key: &str,
+    path: &Path,
+    size: u64,
+) -> Result<StoredObject, AppError> {
    let destination = local_path(state, key)?;
    if let Some(parent) = destination.parent() {
        tokio::fs::create_dir_all(parent).await.map_err(|err| {
@@ -247,10 +294,21 @@ pub async fn store_file(
        endpoint_id: None,
        key: destination.to_string_lossy().to_string(),
        etag: None,
-        size: metadata.len(),
+        size,
    })
 }

+fn log_local_fallback(endpoint: &StorageEndpoint, key: &str, err: &AppError) {
+    tracing::warn!(
+        storage_endpoint_id = %endpoint.id,
+        storage_endpoint = %endpoint.name,
+        object_key = %key,
+        error = %err,
+        fallback_backend = "local",
+        "S3 write failed; storing object on local disk"
+    );
+}
+
 pub async fn read_bytes(state: &AppState, object: &ObjectLocator) -> Result<Vec<u8>, AppError> {
    if object.backend == "local" {
        return tokio::fs::read(&object.key).await.map_err(|err| {
@@ -509,6 +567,12 @@ fn build_client(
        .region(Region::new(region.to_string()))
        .endpoint_url(endpoint_url)
        .force_path_style(force_path_style)
+        .retry_config(RetryConfig::standard().with_max_attempts(2))
+        .timeout_config(
+            TimeoutConfig::builder()
+                .connect_timeout(Duration::from_secs(3))
+                .build(),
+        )
        .request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
        .response_checksum_validation(ResponseChecksumValidation::WhenRequired)
        .build();