From ac08ccbe397be7e5b5168e5ceb53a851fae48d7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E5=BE=B7=E5=AE=87?= Date: Mon, 21 Sep 2026 14:33:51 +0800 Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=AD=A3spacetime-client=20checkpoint?= =?UTF-8?q?=20facade=E7=9A=84=E8=90=BD=E7=82=B9=20-=20facade=20=E6=96=B9?= =?UTF-8?q?=E6=B3=95=E8=A1=A5=E8=BF=9B=E7=9C=9F=E6=AD=A3=E5=8F=82=E4=B8=8E?= =?UTF-8?q?=E7=BC=96=E8=AF=91=E7=9A=84=20active/external=5Fgeneration.rs?= =?UTF-8?q?=EF=BC=8C=E6=AD=A4=E5=89=8D=E8=AF=AF=E5=86=99=E5=9C=A8=E6=9C=AA?= =?UTF-8?q?=E8=A2=AB=E7=BC=96=E8=AF=91=E7=9A=84=E9=87=8D=E5=A4=8D=E6=96=87?= =?UTF-8?q?=E4=BB=B6=E9=87=8C=20-=20mapper=20=E9=87=8D=E6=96=B0=E5=AF=BC?= =?UTF-8?q?=E5=87=BA=20ExternalGenerationJobProviderCheckpointRecordInput?= =?UTF-8?q?=EF=BC=8C=E8=B0=83=E7=94=A8=E6=96=B9=E6=97=A0=E9=9C=80=E4=BB=8E?= =?UTF-8?q?=E5=AD=90=E6=A8=A1=E5=9D=97=E8=B7=AF=E5=BE=84=E5=BC=95=E7=94=A8?= =?UTF-8?q?=20-=20=E5=88=A0=E9=99=A4=E6=9C=AA=E8=A2=AB=E7=BC=96=E8=AF=91?= =?UTF-8?q?=E7=9A=84=20src/external=5Fgeneration.rs=20=E9=87=8D=E5=A4=8D?= =?UTF-8?q?=E5=89=AF=E6=9C=AC=EF=BC=8C=E9=81=BF=E5=85=8D=E5=86=8D=E6=AC=A1?= =?UTF-8?q?=E6=94=B9=E5=88=B0=E6=AD=BB=E6=96=87=E4=BB=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/active/external_generation.rs | 30 + .../spacetime-client/src/active/mapper.rs | 6 +- .../src/external_generation.rs | 641 ------------------ 3 files changed, 33 insertions(+), 644 deletions(-) delete mode 100644 server-rs/crates/spacetime-client/src/external_generation.rs diff --git a/server-rs/crates/spacetime-client/src/active/external_generation.rs b/server-rs/crates/spacetime-client/src/active/external_generation.rs index 12f74f3d2..c693a3776 100644 --- a/server-rs/crates/spacetime-client/src/active/external_generation.rs +++ b/server-rs/crates/spacetime-client/src/active/external_generation.rs @@ -359,6 +359,36 @@ impl SpacetimeClient { .await } + /// 写入 provider 任务 checkpoint。 + /// + /// 只在 submit 成功、拿到 provider task id 之后调用一次；procedure 侧对已有 + /// checkpoint 的 job 直接拒绝，调用方不得用重试来覆盖，这是 at-most-once submit + /// 的数据层保证。写入失败意味着本次 attempt 无法安全续跑，按失败处理而不是重新 submit。 + pub async fn set_external_generation_job_provider_checkpoint( + &self, + input: ExternalGenerationJobProviderCheckpointRecordInput, + ) -> Result { + let procedure_input = input.into(); + + self.call_after_connect( + "set_external_generation_job_provider_checkpoint_and_return", + move |connection, sender| { + connection + .procedures() + .set_external_generation_job_provider_checkpoint_and_return_then( + procedure_input, + move |_, result| { + let mapped = result + .map_err(SpacetimeClientError::from_sdk_error) + .and_then(map_external_generation_job_procedure_result); + send_once(&sender, mapped); + }, + ); + }, + ) + .await + } + pub async fn get_external_generation_job( &self, input: ExternalGenerationJobGetRecordInput, diff --git a/server-rs/crates/spacetime-client/src/active/mapper.rs b/server-rs/crates/spacetime-client/src/active/mapper.rs index 1f725e294..8e2480e14 100644 --- a/server-rs/crates/spacetime-client/src/active/mapper.rs +++ b/server-rs/crates/spacetime-client/src/active/mapper.rs @@ -85,9 +85,9 @@ pub use self::external_generation::{ ExternalGenerationJobEnqueueRecordInput, ExternalGenerationJobFailRecordInput, ExternalGenerationJobGetRecordInput, ExternalGenerationJobListRecord, ExternalGenerationJobListRecordInput, ExternalGenerationJobPhaseUpdateRecordInput, - ExternalGenerationJobRecord, ExternalGenerationJobRenewLeaseRecordInput, - ExternalGenerationJobSummaryListRecord, ExternalGenerationJobSummaryRecord, - ExternalGenerationQueueStatsRecord, + ExternalGenerationJobProviderCheckpointRecordInput, ExternalGenerationJobRecord, + ExternalGenerationJobRenewLeaseRecordInput, ExternalGenerationJobSummaryListRecord, + ExternalGenerationJobSummaryRecord, ExternalGenerationQueueStatsRecord, }; pub use self::llm_router_account::{LlmRouterAccountRecord, LlmRouterAccountUpsertRecordInput}; pub use self::runtime::FeatureGateConfigRecord; diff --git a/server-rs/crates/spacetime-client/src/external_generation.rs b/server-rs/crates/spacetime-client/src/external_generation.rs deleted file mode 100644 index 4f96e7817..000000000 --- a/server-rs/crates/spacetime-client/src/external_generation.rs +++ /dev/null @@ -1,641 +0,0 @@ -use super::*; -use crate::mapper::*; -use spacetimedb_sdk::TableWithPrimaryKey; -use tokio::sync::watch; - -const EXTERNAL_GENERATION_QUEUE_WAKE_SUBSCRIPTION_QUERIES: [&str; 2] = [ - "SELECT * FROM external_generation_job WHERE status = 'pending'", - "SELECT * FROM external_generation_job WHERE status = 'running'", -]; - -#[derive(Debug)] -pub enum ExternalGenerationJobPhaseUpdateError { - LeaseFencingRejected(String), - Rejected(String), - Rpc(SpacetimeClientError), -} - -impl ExternalGenerationJobPhaseUpdateError { - pub fn is_lease_fencing_rejected(&self) -> bool { - matches!(self, Self::LeaseFencingRejected(_)) - } - - pub fn is_retryable_transport(&self) -> bool { - matches!( - self, - Self::Rpc( - SpacetimeClientError::Build(_) - | SpacetimeClientError::ConnectDropped - | SpacetimeClientError::Timeout(_) - ) - ) - } -} - -impl std::fmt::Display for ExternalGenerationJobPhaseUpdateError { - fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::LeaseFencingRejected(message) | Self::Rejected(message) => { - formatter.write_str(message) - } - Self::Rpc(error) => std::fmt::Display::fmt(error, formatter), - } - } -} - -impl std::error::Error for ExternalGenerationJobPhaseUpdateError {} - -pub struct ExternalGenerationQueueWakeSubscription { - connection: DbConnection, - _subscriptions: Vec, - _insert_callback: ExternalGenerationJobInsertCallbackId, - _update_callback: ExternalGenerationJobUpdateCallbackId, - _delete_callback: ExternalGenerationJobDeleteCallbackId, - runner: Option>, - receiver: watch::Receiver, -} - -impl ExternalGenerationQueueWakeSubscription { - pub async fn changed(&mut self) -> Result<(), SpacetimeClientError> { - self.receiver - .changed() - .await - .map_err(|_| SpacetimeClientError::ConnectDropped) - } -} - -impl Drop for ExternalGenerationQueueWakeSubscription { - fn drop(&mut self) { - let _ = self.connection.disconnect(); - if let Some(runner) = self.runner.take() { - drop(runner); - } - } -} - -impl SpacetimeClient { - pub async fn subscribe_external_generation_queue_wake( - &self, - ) -> Result { - let config = self.config.clone(); - let operation_timeout = config.procedure_timeout; - let (connected_sender, connected_receiver) = - oneshot::channel::>(); - let connected_sender = Arc::new(Mutex::new(Some(connected_sender))); - let connect_sender = connected_sender.clone(); - let disconnect_sender = connected_sender.clone(); - let connection = timeout( - operation_timeout, - tokio::task::spawn_blocking(move || { - DbConnection::builder() - .with_uri(config.server_url) - .with_database_name(config.database) - .with_token(config.token) - .on_connect(move |_, _, _| { - send_connect_once(&connect_sender, Ok(())); - }) - .on_disconnect(move |_, error| { - let message = error - .map(|error| error.to_string()) - .unwrap_or_else(|| "SpacetimeDB 队列订阅连接已断开".to_string()); - send_connect_once( - &disconnect_sender, - Err(SpacetimeClientError::Procedure(message)), - ); - }) - .build() - .map_err(|error| SpacetimeClientError::Build(error.to_string())) - }), - ) - .await - .map_err(|_| SpacetimeClientError::Timeout(SpacetimeClientStage::ConnectBuild))? - .map_err(|error| SpacetimeClientError::Runtime(error.to_string()))??; - - let runner = connection.run_threaded(); - timeout(operation_timeout, connected_receiver) - .await - .map_err(|_| SpacetimeClientError::Timeout(SpacetimeClientStage::ConnectHandshake))? - .map_err(|_| SpacetimeClientError::ConnectDropped)??; - - let (wake_sender, wake_receiver) = watch::channel(0u64); - let wake_counter = Arc::new(AtomicU64::new(0)); - let insert_sender = wake_sender.clone(); - let insert_counter = wake_counter.clone(); - let update_sender = wake_sender.clone(); - let update_counter = wake_counter.clone(); - let delete_sender = wake_sender.clone(); - let delete_counter = wake_counter.clone(); - let insert_callback = connection - .db() - .external_generation_job() - .on_insert(move |_, row| { - if external_generation_queue_row_should_wake(row) { - send_external_generation_queue_wake(&insert_sender, &insert_counter); - } - }); - let update_callback = - connection - .db() - .external_generation_job() - .on_update(move |_, old, new| { - if external_generation_queue_row_should_wake(old) - || external_generation_queue_row_should_wake(new) - { - send_external_generation_queue_wake(&update_sender, &update_counter); - } - }); - let delete_callback = connection - .db() - .external_generation_job() - .on_delete(move |_, row| { - if external_generation_queue_row_should_wake(row) { - send_external_generation_queue_wake(&delete_sender, &delete_counter); - } - }); - - let mut subscriptions = Vec::new(); - for query in EXTERNAL_GENERATION_QUEUE_WAKE_SUBSCRIPTION_QUERIES { - let (applied_sender, applied_receiver) = - oneshot::channel::>(); - let applied_sender = Arc::new(Mutex::new(Some(applied_sender))); - let on_applied_sender = applied_sender.clone(); - let on_error_sender = applied_sender.clone(); - let subscription = connection - .subscription_builder() - .on_applied(move |_| { - send_connect_once(&on_applied_sender, Ok(())); - }) - .on_error(move |_, error| { - send_connect_once( - &on_error_sender, - Err(SpacetimeClientError::Procedure(error.to_string())), - ); - }) - .subscribe(query); - - timeout(operation_timeout, applied_receiver) - .await - .map_err(|_| { - SpacetimeClientError::Timeout(SpacetimeClientStage::ReadModelSubscribe) - })? - .map_err(|_| SpacetimeClientError::ConnectDropped)??; - subscriptions.push(subscription); - } - send_external_generation_queue_wake(&wake_sender, &wake_counter); - - Ok(ExternalGenerationQueueWakeSubscription { - connection, - _subscriptions: subscriptions, - _insert_callback: insert_callback, - _update_callback: update_callback, - _delete_callback: delete_callback, - runner: Some(runner), - receiver: wake_receiver, - }) - } - - pub async fn enqueue_external_generation_job( - &self, - input: ExternalGenerationJobEnqueueRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "enqueue_external_generation_job_and_return", - move |connection, sender| { - connection - .procedures() - .enqueue_external_generation_job_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn claim_external_generation_jobs( - &self, - input: ExternalGenerationJobClaimRecordInput, - ) -> Result, SpacetimeClientError> { - let procedure_input = input.into(); - - self.call_after_connect( - "claim_external_generation_jobs_and_return", - move |connection, sender| { - connection - .procedures() - .claim_external_generation_jobs_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_claim_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn complete_external_generation_job( - &self, - input: ExternalGenerationJobCompleteRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "complete_external_generation_job_and_return", - move |connection, sender| { - connection - .procedures() - .complete_external_generation_job_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn renew_external_generation_job_lease( - &self, - input: ExternalGenerationJobRenewLeaseRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "renew_external_generation_job_lease_and_return", - move |connection, sender| { - connection - .procedures() - .renew_external_generation_job_lease_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn update_external_generation_job_phase( - &self, - input: ExternalGenerationJobPhaseUpdateRecordInput, - ) -> Result { - let procedure_input = input.into(); - - let outcome = self - .call_after_connect( - "update_external_generation_job_phase_and_return", - move |connection, sender| { - connection - .procedures() - .update_external_generation_job_phase_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .map(map_external_generation_job_phase_update_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - .map_err(ExternalGenerationJobPhaseUpdateError::Rpc)?; - - match outcome { - ExternalGenerationJobPhaseUpdateProcedureOutcome::Updated(job) => Ok(job), - ExternalGenerationJobPhaseUpdateProcedureOutcome::Rejected { kind, message } => { - match kind { - ExternalGenerationJobPhaseUpdateFailureKind::LeaseFencingRejected => Err( - ExternalGenerationJobPhaseUpdateError::LeaseFencingRejected(message), - ), - ExternalGenerationJobPhaseUpdateFailureKind::OtherRejected => { - Err(ExternalGenerationJobPhaseUpdateError::Rejected(message)) - } - } - } - } - } - - pub async fn fail_external_generation_job( - &self, - input: ExternalGenerationJobFailRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "fail_external_generation_job_and_return", - move |connection, sender| { - connection - .procedures() - .fail_external_generation_job_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - /// 写入 provider 任务 checkpoint。 - /// - /// 只在 submit 成功、拿到 provider task id 之后调用一次；procedure 侧对已有 - /// checkpoint 的 job 直接拒绝，调用方不得用重试来覆盖，这是 at-most-once submit - /// 的数据层保证。写入失败意味着本次 attempt 无法安全续跑，按失败处理而不是重新 submit。 - pub async fn set_external_generation_job_provider_checkpoint( - &self, - input: ExternalGenerationJobProviderCheckpointRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "set_external_generation_job_provider_checkpoint_and_return", - move |connection, sender| { - connection - .procedures() - .set_external_generation_job_provider_checkpoint_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn get_external_generation_job( - &self, - input: ExternalGenerationJobGetRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "get_external_generation_job_and_return", - move |connection, sender| { - connection - .procedures() - .get_external_generation_job_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - // TODO icons gen's result query is not prepared for now. - pub async fn get_external_generation_job_generated_artifacts( - &self, - input: ExternalGenerationJobGetRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "get_external_generation_job_result_and_return", - move |connection, sender| { - connection - .procedures() - .get_external_generation_job_result_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then( - map_external_generation_job_result_procedure_generated_artifacts, - ); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn get_external_generation_job_summary( - &self, - input: ExternalGenerationJobGetRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "get_external_generation_job_summary_and_return", - move |connection, sender| { - connection - .procedures() - .get_external_generation_job_summary_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_summary_procedure_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn list_external_generation_jobs( - &self, - input: ExternalGenerationJobListRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "list_external_generation_jobs_and_return", - move |connection, sender| { - connection - .procedures() - .list_external_generation_jobs_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_list_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn list_external_generation_job_summaries( - &self, - input: ExternalGenerationJobListRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "list_external_generation_job_summaries_and_return", - move |connection, sender| { - connection - .procedures() - .list_external_generation_job_summaries_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_summary_list_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn acknowledge_external_generation_jobs( - &self, - input: ExternalGenerationJobAcknowledgeRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "acknowledge_external_generation_jobs_and_return", - move |connection, sender| { - connection - .procedures() - .acknowledge_external_generation_jobs_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_list_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn acknowledge_external_generation_job_summaries( - &self, - input: ExternalGenerationJobAcknowledgeRecordInput, - ) -> Result { - let procedure_input = input.into(); - - self.call_after_connect( - "acknowledge_external_generation_job_summaries_and_return", - move |connection, sender| { - connection - .procedures() - .acknowledge_external_generation_job_summaries_and_return_then( - procedure_input, - move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_job_summary_list_result); - send_once(&sender, mapped); - }, - ); - }, - ) - .await - } - - pub async fn get_external_generation_queue_stats( - &self, - ) -> Result { - self.call_after_connect( - "get_external_generation_queue_stats_and_return", - move |connection, sender| { - connection - .procedures() - .get_external_generation_queue_stats_and_return_then(move |_, result| { - let mapped = result - .map_err(SpacetimeClientError::from_sdk_error) - .and_then(map_external_generation_queue_stats_result); - send_once(&sender, mapped); - }); - }, - ) - .await - } -} - -fn external_generation_queue_row_should_wake(row: &ExternalGenerationJob) -> bool { - matches!(row.status.as_str(), "pending" | "running") -} - -fn send_external_generation_queue_wake(sender: &watch::Sender, counter: &AtomicU64) { - let next = counter.fetch_add(1, Ordering::Relaxed).saturating_add(1); - let _ = sender.send(next); -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn phase_update_only_retries_transport_errors() { - let build = ExternalGenerationJobPhaseUpdateError::Rpc(SpacetimeClientError::Build( - "initial websocket connect failed".to_string(), - )); - assert!(build.is_retryable_transport()); - - let timeout = ExternalGenerationJobPhaseUpdateError::Rpc(SpacetimeClientError::Timeout( - SpacetimeClientStage::ProcedureResult, - )); - assert!(timeout.is_retryable_transport()); - assert!(!timeout.is_lease_fencing_rejected()); - - let disconnected = - ExternalGenerationJobPhaseUpdateError::Rpc(SpacetimeClientError::ConnectDropped); - assert!(disconnected.is_retryable_transport()); - - let lease = ExternalGenerationJobPhaseUpdateError::LeaseFencingRejected( - "lease token 不匹配".to_string(), - ); - assert!(!lease.is_retryable_transport()); - assert!(lease.is_lease_fencing_rejected()); - - let rejected = ExternalGenerationJobPhaseUpdateError::Rejected("phase 非法".to_string()); - assert!(!rejected.is_retryable_transport()); - assert!(!rejected.is_lease_fencing_rejected()); - - let procedure = ExternalGenerationJobPhaseUpdateError::Rpc( - SpacetimeClientError::Procedure("procedure rejected".to_string()), - ); - assert!(!procedure.is_retryable_transport()); - - let runtime = ExternalGenerationJobPhaseUpdateError::Rpc(SpacetimeClientError::Runtime( - "runtime invariant failed".to_string(), - )); - assert!(!runtime.is_retryable_transport()); - } -}