From 4d9811944065b3a273f8b7582b2138ce5566a221 Mon Sep 17 00:00:00 2001 From: Hare Date: Mon, 12 Oct 2026 15:07:27 +0900 Subject: [PATCH] fix(server): bound embedded protocol and Feature shutdown ownership --- crates/workspace-server/src/server.rs | 138 ++++++++---- .../tests/embedded_owned_execution_tests.rs | 197 ++++++++++++++++++ .../tests/protocol_backpressure_tests.rs | 23 +- .../src/server_residual_execution_tests.rs | 61 ++++++ 4 files changed, 367 insertions(+), 52 deletions(-) create mode 100644 crates/workspace-server/src/server/tests/embedded_owned_execution_tests.rs diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 4d75d5985..1a5593d6d 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -2651,16 +2651,7 @@ impl WorkspaceServerApi { async fn shutdown_feature_storage(&self) -> Result<()> { let apis = self.apis.lock().await.values().cloned().collect::>(); for api in apis { - api.worker_projection.shutdown_async().await; - let storage = api.feature_storage.clone(); - api.store - .db_executor() - .run(move || { - storage.shutdown().map_err(|error| { - Error::Store(format!("failed to shut down Feature storage: {error}")) - }) - }) - .await?; + shutdown_workspace_feature_storage(&api).await?; } Ok(()) } @@ -12271,15 +12262,28 @@ pub async fn serve( let orchestrator_hook = tokio::spawn(run_orchestrator_turn_end_hook(api.clone())); let result = axum::serve(listener, build_router(api.clone())).await; orchestrator_hook.abort(); - let storage_result = api - .feature_storage - .shutdown() - .map_err(|error| Error::Store(format!("failed to shut down Feature storage: {error}"))); + let _ = orchestrator_hook.await; + let storage_result = shutdown_workspace_feature_storage(&api).await; result?; storage_result?; Ok(()) } +// Feature shutdown fences new operations and waits for synchronous connection +// holders. Accepted close retains ownership after its async waiter disappears. +async fn shutdown_workspace_feature_storage(api: &WorkspaceApi) -> Result<()> { + api.worker_projection.shutdown_async().await; + let storage = api.feature_storage.clone(); + api.store + .db_executor() + .run(move || { + storage.shutdown().map_err(|error| { + Error::Store(format!("failed to shut down Feature storage: {error}")) + }) + }) + .await +} + #[derive(Debug, Serialize, Deserialize)] pub struct ProjectListResponse { pub workspace_id: String, @@ -31749,7 +31753,7 @@ pub(crate) async fn connect_workspace_worker_protocol( RuntimeObservationSource::Embedded(source) => { let scope = worker_runtime::RuntimeWorkspaceScope::new(api.workspace_id(), "embedded-backend"); - connect_embedded_worker_protocol(source, scope).await + connect_embedded_worker_protocol(source, scope, api.host_executors.host()?).await } } } @@ -31845,18 +31849,25 @@ async fn connect_remote_worker_protocol( async fn connect_embedded_worker_protocol( source: crate::observation::EmbeddedRuntimeObservationSource, scope: worker_runtime::RuntimeWorkspaceScope, + executor: crate::db_executor::DbExecutor, ) -> Result { - // Capture one real execution endpoint, including its greeting/snapshot. - // The WorkerRef-wide Workspace observation bus is a separate catalogue - // channel and must not reconnect this operational session after Restore. - let mut transport = source - .runtime - .attach_worker_protocol_scoped(&scope, &source.worker_ref) - .map_err(|error| Error::RuntimeOperationFailed { - runtime_id: source.worker.runtime_id.clone(), - code: "embedded_worker_protocol_connect_failed".to_string(), - message: error.to_string(), - })?; + // Endpoint capture takes Runtime operation/state locks. Keep it, and each + // intact method (including fs-store commit/publication), off Tokio and the + // sole DB lane. A discarded reply retains the endpoint until cleanup. + let capture_source = source.clone(); + let capture_scope = scope.clone(); + let mut transport = executor + .run(move || { + capture_source + .runtime + .attach_worker_protocol_scoped(&capture_scope, &capture_source.worker_ref) + .map_err(|error| Error::RuntimeOperationFailed { + runtime_id: capture_source.worker.runtime_id.clone(), + code: "embedded_worker_protocol_connect_failed".to_string(), + message: error.to_string(), + }) + }) + .await?; Ok(spawn_worker_protocol_pump( move |mut methods, events| async move { if let Err(reason) = enqueue_worker_protocol_event(&events, transport.snapshot.clone()) @@ -31868,9 +31879,26 @@ async fn connect_embedded_worker_protocol( method = methods.recv() => { let Some(method) = method else { return WorkerProtocolEnd::Disconnected; }; let shutdown = matches!(&method, protocol::Method::Shutdown { .. }); - match source.runtime.send_connected_protocol_method_scoped( - &scope, &source.worker_ref, &transport, method, - ) { + let method_source = source.clone(); + let method_scope = scope.clone(); + let owned = executor.run(move || { + let result = method_source.runtime.send_connected_protocol_method_scoped( + &method_scope, &method_source.worker_ref, &transport, method, + ); + // Endpoint validation is synchronous too. + let terminal = result.as_ref().err().is_some_and(|error| + embedded_protocol_dispatch_is_terminal(error, &transport)); + Ok::<_, Error>((transport, result, terminal)) + }).await; + let (returned, result, terminal) = match owned { + Ok(value) => value, + Err(error) => { + let _ = enqueue_worker_protocol_event(&events, protocol_error_event(error.to_string())); + return WorkerProtocolEnd::Disconnected; + } + }; + transport = returned; + match result { Ok(direct_events) => { for event in direct_events { let terminal = matches!(event, protocol::Event::Shutdown); @@ -31879,7 +31907,6 @@ async fn connect_embedded_worker_protocol( } } Err(error) => { - let terminal = embedded_protocol_dispatch_is_terminal(&error, &transport); if let Err(reason) = enqueue_worker_protocol_event(&events, protocol_error_event(error.to_string())) { return reason; } if terminal { return WorkerProtocolEnd::Disconnected; } } @@ -31892,7 +31919,6 @@ async fn connect_embedded_worker_protocol( if let Err(reason) = enqueue_worker_protocol_event(&events, event) { return reason; } if terminal { return WorkerProtocolEnd::Disconnected; } } - // Stop, disconnect and upstream lag never cause an implicit reattach. None => return WorkerProtocolEnd::Disconnected, } } @@ -35184,6 +35210,7 @@ mod tests { mod drive_api_tests; mod drive_tests; mod drive_web_roundtrip_tests; + mod embedded_owned_execution_tests; mod external_checkout_tests; mod protocol_backpressure_tests; mod subject_spawn_tests; @@ -45338,6 +45365,7 @@ mod tests { // The receiver belongs to one mock execution, not to a Runtime WorkerRecord. // Bound protocol senders become terminal when Stop drops this execution. struct DeterministicWorkerExecution { + dispatch_hook: std::sync::Mutex>>, context: worker_runtime::execution::WorkerExecutionContext, state: std::sync::Mutex, method_tx: tokio::sync::mpsc::Sender, @@ -45357,6 +45385,7 @@ mod tests { let (method_tx, method_rx) = tokio::sync::mpsc::channel(32); let (events, _) = tokio::sync::broadcast::channel(32); Self { + dispatch_hook: std::sync::Mutex::new(None), context, state: std::sync::Mutex::new(protocol::WorkerStateSnapshot::initial()), method_tx, @@ -45380,6 +45409,9 @@ mod tests { Vec, worker_runtime::execution::WorkerExecutionResult, > { + if let Some(hook) = self.dispatch_hook.lock().unwrap().take() { + hook(); + } // Drive the actual mock receiver synchronously; no method is routed // through a WorkerRef lookup after a protocol connection is attached. let mut receiver = self.method_rx.lock().unwrap(); @@ -61492,10 +61524,13 @@ mod tests { for delayed in delayed_protocol_test_methods() { let (runtime, worker_ref, execution) = runtime_with_worker_backend(); let source = embedded_protocol_test_source(&runtime, &worker_ref); - let mut old = - connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope()) - .await - .unwrap(); + let mut old = connect_embedded_worker_protocol( + source.clone(), + embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), + ) + .await + .unwrap(); assert_embedded_protocol_snapshot( next_embedded_protocol_event(&mut old).await, &worker_ref, @@ -61559,10 +61594,13 @@ mod tests { &[worker_ref.clone()] ); - let mut current = - connect_embedded_worker_protocol(source, embedded_protocol_test_scope()) - .await - .unwrap(); + let mut current = connect_embedded_worker_protocol( + source, + embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), + ) + .await + .unwrap(); assert_embedded_protocol_snapshot( next_embedded_protocol_event(&mut current).await, &worker_ref, @@ -61710,10 +61748,13 @@ mod tests { for lagged in [false, true] { let (runtime, worker_ref, execution) = runtime_with_worker_backend(); let source = embedded_protocol_test_source(&runtime, &worker_ref); - let mut bridge = - connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope()) - .await - .unwrap(); + let mut bridge = connect_embedded_worker_protocol( + source.clone(), + embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), + ) + .await + .unwrap(); assert_embedded_protocol_snapshot( next_embedded_protocol_event(&mut bridge).await, &worker_ref, @@ -61755,6 +61796,7 @@ mod tests { let mut bridge = connect_embedded_worker_protocol( embedded_protocol_test_source(&runtime, &worker_ref), embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), ) .await .unwrap(); @@ -61815,9 +61857,13 @@ mod tests { let wrong_scope = worker_runtime::RuntimeWorkspaceScope::new("other-workspace", "local-token"); assert!( - connect_embedded_worker_protocol(source.clone(), wrong_scope.clone()) - .await - .is_err() + connect_embedded_worker_protocol( + source.clone(), + wrong_scope.clone(), + HostExecutors::default().host().unwrap() + ) + .await + .is_err() ); let (mut socket, server) = spawn_embedded_protocol_test_ws(source.clone(), wrong_scope, false).await; diff --git a/crates/workspace-server/src/server/tests/embedded_owned_execution_tests.rs b/crates/workspace-server/src/server/tests/embedded_owned_execution_tests.rs new file mode 100644 index 000000000..55444a031 --- /dev/null +++ b/crates/workspace-server/src/server/tests/embedded_owned_execution_tests.rs @@ -0,0 +1,197 @@ +use super::*; +use crate::db_executor::{ + DbExecutor, DbExecutorLimits, DbOperationLifetime, with_operation_lifetime, +}; +use std::{ + num::NonZeroUsize, + sync::{ + atomic::{AtomicUsize, Ordering}, + mpsc, + }, + time::Duration, +}; + +fn owner() -> DbExecutor { + DbExecutor::new(DbExecutorLimits { + workers: NonZeroUsize::new(1).unwrap(), + queued: NonZeroUsize::new(1).unwrap(), + }) + .unwrap() +} +async fn observed(future: impl std::future::Future) -> T { + tokio::time::timeout(Duration::from_secs(2), future) + .await + .expect("owned protocol stage did not progress") +} +struct Release(Option>); +impl Release { + fn release(&mut self) { + if let Some(tx) = self.0.take() { + let _ = tx.send(()); + } + } +} +impl Drop for Release { + fn drop(&mut self) { + self.release(); + } +} +struct Lifetime(Arc); +struct Retained(Arc); +impl Drop for Retained { + fn drop(&mut self) { + self.0.fetch_sub(1, Ordering::SeqCst); + } +} +impl DbOperationLifetime for Lifetime { + fn retain(&self) -> Box { + self.0.fetch_add(1, Ordering::SeqCst); + Box::new(Retained(self.0.clone())) + } +} + +#[tokio::test(flavor = "current_thread")] +async fn cancelled_embedded_endpoint_capture_keeps_actual_host_budget_off_tokio() { + let (runtime, worker, execution) = runtime_with_worker_backend(); + let (started, held) = tokio::sync::oneshot::channel(); + let (tx, rx) = mpsc::channel(); + let mut release = Release(Some(tx)); + let thread = std::thread::spawn(move || { + let _guard = execution.contexts.lock().unwrap(); + let _ = started.send(()); + let _ = rx.recv_timeout(Duration::from_secs(5)); + }); + observed(held).await.unwrap(); + let executor = owner(); + let leases = Arc::new(AtomicUsize::new(0)); + let mut capture = Box::pin(with_operation_lifetime( + Arc::new(Lifetime(leases.clone())), + connect_embedded_worker_protocol( + embedded_protocol_test_source(&runtime, &worker), + embedded_protocol_test_scope(), + executor.clone(), + ), + )); + assert!( + futures::poll!(&mut capture).is_pending(), + "endpoint Runtime lock must not run on Tokio" + ); + observed(tokio::time::sleep(Duration::from_millis(1))).await; + assert_eq!( + observed(tokio::spawn(async { "control" })).await.unwrap(), + "control" + ); + drop(capture); + assert_eq!(leases.load(Ordering::SeqCst), 1); + executor.dispatch(|| {}).unwrap(); // only the one queued slot is left + let error = executor.run(|| Ok::<_, Error>(())).await.unwrap_err(); + assert!(matches!(error, Error::DbCapacityExceeded)); + release.release(); + thread.join().unwrap(); + observed(executor.drain()).await.unwrap(); + assert_eq!(leases.load(Ordering::SeqCst), 0); +} + +#[tokio::test(flavor = "current_thread")] +async fn cancelled_embedded_shutdown_commits_fs_receipt_after_the_pump_disappears() { + let dir = tempfile::tempdir().unwrap(); + let execution = Arc::new(DeterministicExecutionBackend::default()); + let options = worker_runtime::fs_store::FsRuntimeStoreOptions { + root: dir.path().join("runtime"), + runtime_id: "owned-protocol".into(), + display_name: None, + }; + let runtime = worker_runtime::Runtime::with_fs_store_and_execution_backend( + options.clone(), + execution.clone(), + ) + .unwrap(); + runtime.store_config_bundle(runtime_test_bundle()).unwrap(); + let worker = runtime + .create_worker_scoped(&embedded_protocol_test_scope(), runtime_create_request()) + .unwrap() + .worker_ref; + let executor = owner(); + let mut connection = connect_embedded_worker_protocol( + embedded_protocol_test_source(&runtime, &worker), + embedded_protocol_test_scope(), + executor.clone(), + ) + .await + .unwrap(); + assert!(matches!( + observed(connection.events.recv()).await.unwrap(), + protocol::Event::Snapshot { .. } + )); + let (entered, held) = tokio::sync::oneshot::channel(); + let (tx, rx) = mpsc::channel(); + let mut release = Release(Some(tx)); + let endpoint = execution + .contexts + .lock() + .unwrap() + .get(&worker) + .unwrap() + .clone(); + let gate_finished = Arc::new(std::sync::atomic::AtomicBool::new(false)); + let finished = gate_finished.clone(); + *endpoint.dispatch_hook.lock().unwrap() = Some(Box::new(move || { + let _ = entered.send(()); + let _ = rx.recv_timeout(Duration::from_secs(5)); + finished.store(true, Ordering::SeqCst); + })); + connection + .methods + .send(protocol::Method::Shutdown { + command: protocol::WorkerCommandEnvelope::new(700), + }) + .await + .unwrap(); + observed(held).await.unwrap(); // actual Runtime method, after its durable stop admission + observed(tokio::time::sleep(Duration::from_millis(1))).await; + assert_eq!( + observed(tokio::spawn(async { "control" })).await.unwrap(), + "control" + ); + assert!( + !gate_finished.load(Ordering::SeqCst), + "watchdog bailout is not Tokio progress" + ); + connection + .termination + .terminate(WorkerProtocolEnd::Disconnected); + drop(connection); + executor.dispatch(|| {}).unwrap(); + assert!(matches!( + executor.run(|| Ok::<_, Error>(())).await.unwrap_err(), + Error::DbCapacityExceeded + )); + release.release(); + observed(executor.drain()).await.unwrap(); + assert_eq!( + runtime.worker_detail(&worker).unwrap().status, + worker_runtime::catalog::WorkerStatus::Stopped + ); + // Read the actual fs-store aggregate written by the accepted method. The + // deterministic execution fixture has no restorable Profile archive, so + // restarting it would test archive recovery rather than this commit owner. + let record: serde_json::Value = serde_json::from_slice( + &std::fs::read( + options + .root + .join("workers") + .join(worker.worker_id.to_string()) + .join("worker.json"), + ) + .unwrap(), + ) + .unwrap(); + assert_eq!( + record["status"], "stopped", + "cancelled pump lost its fs-store commit: {record}" + ); + assert_ne!( + record["execution_state"]["state"], + "reconciliation_required" + ); +} diff --git a/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs b/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs index 358b9a7bb..06759a67e 100644 --- a/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs +++ b/crates/workspace-server/src/server/tests/protocol_backpressure_tests.rs @@ -11,6 +11,7 @@ async fn embedded_event_overflow_terminates_without_losing_the_admitted_prefix() let mut connection = connect_embedded_worker_protocol( embedded_protocol_test_source(&runtime, &worker_ref), embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), ) .await .unwrap(); @@ -69,6 +70,7 @@ async fn embedded_protocol_delivers_snapshot_then_ordered_replies_and_live_event let mut connection = connect_embedded_worker_protocol( embedded_protocol_test_source(&runtime, &worker_ref), embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), ) .await .unwrap(); @@ -105,9 +107,13 @@ async fn embedded_protocol_delivers_snapshot_then_ordered_replies_and_live_event async fn embedded_overflow_requires_explicit_reconnect_and_stop_ends_the_new_connection() { let (runtime, worker_ref, execution) = runtime_with_worker_backend(); let source = embedded_protocol_test_source(&runtime, &worker_ref); - let mut old = connect_embedded_worker_protocol(source.clone(), embedded_protocol_test_scope()) - .await - .unwrap(); + let mut old = connect_embedded_worker_protocol( + source.clone(), + embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), + ) + .await + .unwrap(); assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut old).await, &worker_ref); tokio::time::timeout(PROTOCOL_TEST_TIMEOUT, async { for index in 0..=WORKER_PROTOCOL_EVENT_CAPACITY { @@ -127,9 +133,13 @@ async fn embedded_overflow_requires_explicit_reconnect_and_stop_ends_the_new_con // Overflow ends only the attached connection, not its Worker. Reconnection // is explicit and starts from a fresh execution snapshot, not old backlog. - let mut fresh = connect_embedded_worker_protocol(source, embedded_protocol_test_scope()) - .await - .unwrap(); + let mut fresh = connect_embedded_worker_protocol( + source, + embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), + ) + .await + .unwrap(); assert_embedded_protocol_snapshot(next_embedded_protocol_event(&mut fresh).await, &worker_ref); fresh .methods @@ -797,6 +807,7 @@ async fn embedded_disconnect_and_last_owner_drop_release_the_execution_event_sub let mut connection = connect_embedded_worker_protocol( embedded_protocol_test_source(&runtime, &worker_ref), embedded_protocol_test_scope(), + HostExecutors::default().host().unwrap(), ) .await .unwrap(); diff --git a/crates/workspace-server/src/server_residual_execution_tests.rs b/crates/workspace-server/src/server_residual_execution_tests.rs index 8b0e3620a..abbe46393 100644 --- a/crates/workspace-server/src/server_residual_execution_tests.rs +++ b/crates/workspace-server/src/server_residual_execution_tests.rs @@ -841,3 +841,64 @@ async fn ticket_tracking_and_case_authority_runtime_lookup_leave_db_and_tokio_fr ); } } + +#[tokio::test(flavor = "current_thread")] +async fn cancelled_feature_shutdown_keeps_sql_close_owned_without_blocking_control() { + let fixture = fixture().await; + let (database, _) = fixture.feature_database(); + let mutation = database.clone(); + let (entered, waiting) = oneshot::channel(); + let (release, permit) = mpsc::channel(); + let released = Arc::new(AtomicBool::new(false)); + let flag = released.clone(); + let thread = std::thread::spawn(move || { + mutation + .with_connection(|conn| { + conn.execute_batch("BEGIN IMMEDIATE; CREATE TABLE close_owner_probe(id INTEGER)")?; + let _ = entered.send(()); + let _ = permit.recv_timeout(WATCHDOG); + flag.store(true, Ordering::SeqCst); + conn.execute_batch("ROLLBACK")?; + Ok(()) + }) + .unwrap(); + }); + let mut lock = WriteLock { + release: Some(release), + thread: Some(thread), + released, + }; + observed(waiting).await.unwrap(); // actual SQL transaction owns the Feature connection + let leases = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut close = Box::pin(crate::db_executor::with_operation_lifetime( + Arc::new(CountJobLifetime(leases.clone())), + shutdown_workspace_feature_storage(&fixture.api), + )); + assert!( + futures::poll!(&mut close).is_pending(), + "Feature connection close must not wait on Tokio" + ); + runtime_progress_while_locked(&lock).await; + assert_eq!( + observed(fixture.host_executor.run(|| Ok::<_, Error>("control"))) + .await + .unwrap(), + "control" + ); + fixture.executor.dispatch(|| {}).unwrap(); + assert!(matches!( + shutdown_workspace_feature_storage(&fixture.api) + .await + .unwrap_err(), + Error::DbCapacityExceeded + )); + drop(close); + assert_eq!(leases.load(Ordering::SeqCst), 1); + lock.release(); + observed(fixture.executor.drain()).await.unwrap(); + assert_eq!(leases.load(Ordering::SeqCst), 0); + assert!( + database.with_connection(|_| Ok(())).is_err(), + "accepted shutdown must fence later use" + ); +}