From d9bb2c867ea94b37bfcd55b628a60bf8cf92ef3b Mon Sep 17 00:00:00 2001 From: Floating-Y <118035379+Floating-Y@users.noreply.github.com> Date: Mon, 26 Oct 2026 22:20:27 +0800 Subject: [PATCH] Keep sources schedulable when sync job insertion fails Signed-off-by: Floating-Y <118035379+Floating-Y@users.noreply.github.com> --- .../src/api/google_drive_routes.rs | 8 +- .../src/api/google_drive_routes_tests.rs | 265 ++++++++++++++++++ .../utopia-server/src/api/sources_routes.rs | 16 +- crates/utopia-server/src/main.rs | 18 +- crates/utopia-store/src/sources.rs | 40 ++- 5 files changed, 308 insertions(+), 39 deletions(-) diff --git a/crates/utopia-server/src/api/google_drive_routes.rs b/crates/utopia-server/src/api/google_drive_routes.rs index 12efe9740..717898445 100644 --- a/crates/utopia-server/src/api/google_drive_routes.rs +++ b/crates/utopia-server/src/api/google_drive_routes.rs @@ -469,13 +469,7 @@ pub async fn select( ) .await?; if !body.folders.is_empty() { - let _ = utopia_store::sources::mark_queued(&state.pool, source_id).await; - utopia_store::jobs::enqueue( - &state.pool, - "sync_source", - json!({ "source_id": source_id }), - ) - .await?; + utopia_store::sources::enqueue_sync_after_config_change(&state.pool, source_id).await?; } state.emit_source(kb_id); let _ = utopia_store::audit::record( diff --git a/crates/utopia-server/src/api/google_drive_routes_tests.rs b/crates/utopia-server/src/api/google_drive_routes_tests.rs index 06174863a..33793df77 100644 --- a/crates/utopia-server/src/api/google_drive_routes_tests.rs +++ b/crates/utopia-server/src/api/google_drive_routes_tests.rs @@ -1,7 +1,10 @@ //! 连接流程的几道关卡与「谁连上的谁来选」（0074 决定 3、5），连真库跑。 //! 测试环境没配 OAuth 客户端：过了所有关卡的回调停在「没配置」，不会去敲 Google use super::*; +use axum::body::{to_bytes, Body}; +use axum::http::{Request, StatusCode}; use std::sync::Arc; +use tower::ServiceExt; use utopia_core::models::User; struct Fixture { @@ -114,6 +117,268 @@ impl Fixture { .await? .config) } + + async fn sync_http(&self) -> anyhow::Result<(StatusCode, Value)> { + self.post_http( + &format!("/api/v1/kbs/{}/sources/{}/sync", self.kb, self.source), + json!({}), + ) + .await + } + + async fn post_http(&self, path: &str, body: Value) -> anyhow::Result<(StatusCode, Value)> { + let token = crate::auth::issue_token(&self.state, self.connector.id)?; + let response = super::super::router(self.state.clone(), &Default::default()) + .oneshot( + Request::post(path) + .header("Authorization", format!("Bearer {token}")) + .header("Content-Type", "application/json") + .body(Body::from(body.to_string()))?, + ) + .await?; + let status = response.status(); + let bytes = to_bytes(response.into_body(), 1 << 20).await?; + Ok((status, serde_json::from_slice(&bytes)?)) + } + + async fn sync_state(&self) -> anyhow::Result<(String, i64)> { + Ok(sqlx::query_as( + "SELECT last_sync_status, + (SELECT count(*) FROM jobs WHERE kind = 'sync_source' + AND payload->>'source_id' = sources.id::text) + FROM sources WHERE id = $1", + ) + .bind(self.source) + .fetch_one(&self.pool) + .await?) + } + + async fn choose_folder(&self, id: &str) -> ApiResult> { + select( + State(self.state.clone()), + AuthUser(self.connector.clone()), + Path((self.kb, self.source)), + Json(SelectBody { + folders: vec![FolderEntry { + id: id.into(), + name: id.into(), + }], + }), + ) + .await + } + + async fn reject_sync_jobs(&self, reject: bool) -> anyhow::Result<()> { + // Install once: dropping a table trigger between tests can deadlock parallel fixtures. + static INSTALLED: tokio::sync::OnceCell<()> = tokio::sync::OnceCell::const_new(); + INSTALLED + .get_or_try_init(|| async { + sqlx::raw_sql( + "SELECT pg_advisory_xact_lock(hashtext('test_reject_source_sync')); + CREATE OR REPLACE FUNCTION test_reject_source_sync() RETURNS trigger + LANGUAGE plpgsql AS $$ BEGIN + IF NEW.kind = 'sync_source' AND EXISTS( + SELECT 1 FROM sources WHERE id::text = NEW.payload->>'source_id' + AND name = 'reject-source-sync-test') + THEN RAISE EXCEPTION 'test source sync insert failure'; END IF; + RETURN NEW; + END $$; + DO $$ BEGIN + IF NOT EXISTS(SELECT 1 FROM pg_trigger WHERE tgname = 'test_reject_source_sync' + AND tgrelid = 'jobs'::regclass) THEN + CREATE TRIGGER test_reject_source_sync BEFORE INSERT ON jobs + FOR EACH ROW EXECUTE FUNCTION test_reject_source_sync(); + END IF; + END $$;", + ) + .execute(&self.pool) + .await + .map(drop) + }) + .await?; + sqlx::query("UPDATE sources SET name = $2 WHERE id = $1") + .bind(self.source) + .bind(if reject { + "reject-source-sync-test" + } else { + "Drive" + }) + .execute(&self.pool) + .await?; + Ok(()) + } +} + +#[tokio::test] +async fn source_sync_enqueue_creation_failure_keeps_source_schedulable() -> anyhow::Result<()> { + let Some(mut fx) = Fixture::new().await? else { + return Ok(()); + }; + fx.reject_sync_jobs(false).await?; + let config = json!({"url": "https://example.com/source-sync-test"}); + let (status, body) = fx + .post_http( + &format!("/api/v1/kbs/{}/sources", fx.kb), + json!({"kind": "url", "name": "reject-source-sync-test", "config": config}), + ) + .await?; + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR, "{body}"); + fx.source = sqlx::query_scalar( + "SELECT id FROM sources WHERE kb_id = $1 AND kind = 'url' + AND name = 'reject-source-sync-test'", + ) + .bind(fx.kb) + .fetch_one(&fx.pool) + .await?; + assert_eq!(fx.config().await?, config); + assert_eq!(fx.sync_state().await?, ("never".into(), 0)); + fx.reject_sync_jobs(false).await?; + let (status, body) = fx.sync_http().await?; + assert_eq!(status, StatusCode::OK, "{body}"); + assert_eq!(body["queued"], true); + assert_eq!(fx.sync_state().await?, ("queued".into(), 1)); + Ok(()) +} + +#[tokio::test] +async fn source_sync_enqueue_manual_failure_rolls_back_and_can_retry() -> anyhow::Result<()> { + let Some(fx) = Fixture::new().await? else { + return Ok(()); + }; + fx.reject_sync_jobs(true).await?; + let (status, body) = fx.sync_http().await?; + assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR, "{body}"); + assert_eq!(fx.sync_state().await?, ("never".into(), 0)); + + fx.reject_sync_jobs(false).await?; + let (status, body) = fx.sync_http().await?; + assert_eq!(status, StatusCode::OK, "{body}"); + assert_eq!(body["queued"], true); + assert_eq!(fx.sync_state().await?, ("queued".into(), 1)); + let (payload, max_attempts): (Value, i32) = sqlx::query_as( + "SELECT payload, max_attempts FROM jobs WHERE kind = 'sync_source' + AND payload->>'source_id' = $1", + ) + .bind(fx.source.to_string()) + .fetch_one(&fx.pool) + .await?; + assert_eq!(payload, json!({"source_id": fx.source})); + assert_eq!(max_attempts, 3); + + for pending in ["queued", "running"] { + sqlx::query("UPDATE sources SET last_sync_status = $2 WHERE id = $1") + .bind(fx.source) + .bind(pending) + .execute(&fx.pool) + .await?; + let (status, body) = fx.sync_http().await?; + assert_eq!(status, StatusCode::OK, "{body}"); + assert_eq!(body["queued"], false); + assert_eq!(fx.sync_state().await?, (pending.into(), 1)); + } + Ok(()) +} + +#[tokio::test] +async fn source_sync_enqueue_concurrent_manual_requests_enqueue_once() -> anyhow::Result<()> { + let Some(fx) = Fixture::new().await? else { + return Ok(()); + }; + // AppState and the router do not start a worker, so the first task stays queued. + let (first, second) = tokio::join!(fx.sync_http(), fx.sync_http()); + let mut queued = Vec::new(); + for result in [first, second] { + let (status, body) = result?; + assert_eq!(status, StatusCode::OK, "{body}"); + queued.push(body["queued"].as_bool().expect("queued is a boolean")); + } + queued.sort_unstable(); + assert_eq!(queued, [false, true]); + assert_eq!(fx.sync_state().await?, ("queued".into(), 1)); + Ok(()) +} + +#[tokio::test] +async fn source_sync_enqueue_drive_selection_keeps_follow_up_when_pending() -> anyhow::Result<()> { + for pending in ["queued", "running"] { + let Some(fx) = Fixture::new().await? else { + return Ok(()); + }; + let (status, body) = fx.sync_http().await?; + assert_eq!(status, StatusCode::OK, "{body}"); + sqlx::query("UPDATE sources SET last_sync_status = $2 WHERE id = $1") + .bind(fx.source) + .bind(pending) + .execute(&fx.pool) + .await?; + sqlx::query("UPDATE jobs SET status = $2 WHERE payload->>'source_id' = $1") + .bind(fx.source.to_string()) + .bind(pending) + .execute(&fx.pool) + .await?; + let _ = fx + .choose_folder("next-folder") + .await + .map_err(|ApiErr(error)| error)?; + assert_eq!(fx.sync_state().await?, (pending.into(), 2)); + let latest_status: String = sqlx::query_scalar( + "SELECT status FROM jobs WHERE kind = 'sync_source' + AND payload->>'source_id' = $1 ORDER BY id DESC LIMIT 1", + ) + .bind(fx.source.to_string()) + .fetch_one(&fx.pool) + .await?; + assert_eq!(latest_status, "queued"); + assert_eq!(fx.config().await?["folders"][0]["id"], "next-folder"); + } + Ok(()) +} + +#[tokio::test] +async fn source_sync_enqueue_drive_failure_keeps_config_but_rolls_back_queue() -> anyhow::Result<()> +{ + for initial in ["never", "running"] { + let Some(fx) = Fixture::new().await? else { + return Ok(()); + }; + let existing_jobs = if initial == "running" { + let (status, body) = fx.sync_http().await?; + assert_eq!(status, StatusCode::OK, "{body}"); + sqlx::query("UPDATE sources SET last_sync_status = 'running' WHERE id = $1") + .bind(fx.source) + .execute(&fx.pool) + .await?; + sqlx::query("UPDATE jobs SET status = 'running' WHERE payload->>'source_id' = $1") + .bind(fx.source.to_string()) + .execute(&fx.pool) + .await?; + 1 + } else { + 0 + }; + fx.reject_sync_jobs(true).await?; + assert!(matches!( + fx.choose_folder("saved-folder").await, + Err(ApiErr(AppError::Db(_))) + )); + assert_eq!(fx.config().await?["folders"][0]["id"], "saved-folder"); + assert_eq!(fx.sync_state().await?, (initial.into(), existing_jobs)); + fx.reject_sync_jobs(false).await?; + let _ = fx + .choose_folder("saved-folder") + .await + .map_err(|ApiErr(error)| error)?; + let expected_status = if initial == "running" { + "running" + } else { + "queued" + }; + assert_eq!( + fx.sync_state().await?, + (expected_status.into(), existing_jobs + 1) + ); + } + Ok(()) } fn denied_code(result: ApiResult) -> Option<&'static str> { diff --git a/crates/utopia-server/src/api/sources_routes.rs b/crates/utopia-server/src/api/sources_routes.rs index ddd859dd8..f698e78b4 100644 --- a/crates/utopia-server/src/api/sources_routes.rs +++ b/crates/utopia-server/src/api/sources_routes.rs @@ -132,13 +132,7 @@ pub async fn create( .await?; // 拉取型来源建好立即同步一次（有配置即产出，无需等下一个调度周期） if matches!(source.kind.as_str(), "url" | "rss" | "custom") { - let _ = utopia_store::sources::mark_queued(&state.pool, source.id).await; - utopia_store::jobs::enqueue( - &state.pool, - "sync_source", - json!({ "source_id": source.id }), - ) - .await?; + utopia_store::sources::enqueue_sync(&state.pool, source.id).await?; } // api 来源：生成专属推送密钥（此后可随时经 get_token 查看） let mut ingest_token: Option = None; @@ -396,14 +390,8 @@ pub async fn sync_now( ) -> ApiResult> { require_kb(&state, &user, kb_id, Role::Editor).await?; source_in_kb(&state, kb_id, source_id).await?; - let queued = utopia_store::sources::mark_queued(&state.pool, source_id).await?; + let queued = utopia_store::sources::enqueue_sync(&state.pool, source_id).await?; if queued { - utopia_store::jobs::enqueue( - &state.pool, - "sync_source", - json!({ "source_id": source_id }), - ) - .await?; state.emit_source(kb_id); } Ok(Json(json!({ "queued": queued }))) diff --git a/crates/utopia-server/src/main.rs b/crates/utopia-server/src/main.rs index df3141134..c0486cbd7 100644 --- a/crates/utopia-server/src/main.rs +++ b/crates/utopia-server/src/main.rs @@ -232,20 +232,10 @@ async fn main() -> anyhow::Result<()> { match utopia_store::sources::due_sources(&sched_state.pool).await { Ok(due) => { for s in due { - match utopia_store::sources::mark_queued(&sched_state.pool, s.id).await { - Ok(true) => { - if let Err(e) = utopia_store::jobs::enqueue( - &sched_state.pool, - "sync_source", - serde_json::json!({ "source_id": s.id }), - ) - .await - { - tracing::warn!(source_id = %s.id, error = %e, "同步任务入队失败"); - } - } - Ok(false) => {} - Err(e) => tracing::warn!(source_id = %s.id, error = %e, "标记入队失败"), + if let Err(e) = + utopia_store::sources::enqueue_sync(&sched_state.pool, s.id).await + { + tracing::warn!(source_id = %s.id, error = %e, "同步任务入队失败"); } } } diff --git a/crates/utopia-store/src/sources.rs b/crates/utopia-store/src/sources.rs index 6eb3d01c0..be013dfea 100644 --- a/crates/utopia-store/src/sources.rs +++ b/crates/utopia-store/src/sources.rs @@ -1,7 +1,7 @@ //! 摄入来源仓储："来源即文件夹"——source 是容器，挂着它摄入的文档，可定时同步。 use chrono::{DateTime, Utc}; -use sqlx::PgPool; +use sqlx::{PgPool, Postgres, Transaction}; use utopia_core::models::{Role, Source, SourceKind, SourceView, SyncRun, SOURCE_SECRET_KEYS}; use utopia_core::{secrets, AppError, AppResult}; use uuid::Uuid; @@ -378,14 +378,46 @@ pub async fn due_sources(pool: &PgPool) -> AppResult> { .collect()) } -/// 标记入队（幂等：已在队列/运行中则返回 false，避免重复入队）。 -pub async fn mark_queued(pool: &PgPool, id: Uuid) -> AppResult { +/// 状态、任务和 worker 通知一起提交；已在队列/运行中则返回 false。 +pub async fn enqueue_sync(pool: &PgPool, id: Uuid) -> AppResult { + let mut tx = pool.begin().await?; + let queued = mark_queued_tx(&mut tx, id).await?; + if queued { + crate::jobs::enqueue_with_max_attempts_tx( + &mut tx, + "sync_source", + serde_json::json!({ "source_id": id }), + 3, + ) + .await?; + } + tx.commit().await?; + Ok(queued) +} + +/// 换了 Drive 文件夹，即使旧同步正在运行，也要再排一次读取新配置。 +/// 配置由调用方先保存；这里只保证状态和后续任务一起提交。 +pub async fn enqueue_sync_after_config_change(pool: &PgPool, id: Uuid) -> AppResult<()> { + let mut tx = pool.begin().await?; + mark_queued_tx(&mut tx, id).await?; + crate::jobs::enqueue_with_max_attempts_tx( + &mut tx, + "sync_source", + serde_json::json!({ "source_id": id }), + 3, + ) + .await?; + tx.commit().await?; + Ok(()) +} + +async fn mark_queued_tx(tx: &mut Transaction<'_, Postgres>, id: Uuid) -> AppResult { let res = sqlx::query( "UPDATE sources SET last_sync_status = 'queued' WHERE id = $1 AND last_sync_status NOT IN ('queued', 'running')", ) .bind(id) - .execute(pool) + .execute(&mut **tx) .await?; Ok(res.rows_affected() > 0) }