From b272bc37e83c81c9f02a89cb1a2af2a0a61a9c15 Mon Sep 17 00:00:00 2001 From: Hare Date: Sat, 19 Sep 2026 01:34:19 +0900 Subject: [PATCH] fix(workspace-server): make worker removal catalog commits atomic --- crates/workspace-server/src/retention.rs | 137 ++++- crates/workspace-server/src/server.rs | 552 +++++++++++++----- crates/workspace-server/src/store.rs | 332 +++++++---- .../workspace-server/src/worker_projection.rs | 71 +-- .../src/workspace_deletion.rs | 4 +- 5 files changed, 778 insertions(+), 318 deletions(-) diff --git a/crates/workspace-server/src/retention.rs b/crates/workspace-server/src/retention.rs index 6aa4925c6..7b4cc6869 100644 --- a/crates/workspace-server/src/retention.rs +++ b/crates/workspace-server/src/retention.rs @@ -2,7 +2,10 @@ //! Runtime receives only resolved dispositions and stable ids; policy authority //! never comes from prompts, profiles, or model input. -use crate::{Error as StoreError, store::SqliteWorkspaceStore}; +use crate::{ + Error as StoreError, + store::{SqliteWorkspaceStore, WorkerRemovalCommit, commit_worker_catalog_removal}, +}; use chrono::Utc; use rusqlite::{Connection, OptionalExtension, TransactionBehavior, params}; use serde::{Deserialize, Serialize}; @@ -417,13 +420,19 @@ impl SqliteWorkspaceStore { operation_id: &str, fp: &str, result: &WorkerRetentionExecutionResult, - ) -> Result { + ) -> Result { self.with_conn_mut(|conn|{ let tx=conn.transaction_with_behavior(TransactionBehavior::Immediate)?; let mut plan=load_plan_op(&tx,operation_id)?.ok_or_else(||StoreError::InvalidInput("operation missing".into()))?; if plan.workspace_id!=workspace_id{return Err(StoreError::InvalidInput("cross-workspace".into()));} if plan.input_fingerprint!=fp||result.input_fingerprint!=fp||result.operation_id!=operation_id{return Err(StoreError::InvalidInput(format!("fingerprint:{operation_id}")));} - if plan.state==WorkerRemovalPlanState::Succeeded{tx.commit()?;return Ok(plan);} + if plan.state==WorkerRemovalPlanState::Succeeded{ + // Repair projection fences left by the pre-atomic implementation and + // return a replayable change for a broadcast missed after commit. + let catalog=commit_worker_catalog_removal(&tx,workspace_id,&plan.worker)?; + tx.commit()?; + return Ok(WorkerRemovalCommit { plan, catalog }); + } if plan.state != WorkerRemovalPlanState::Executing { return Err(StoreError::InvalidInput(format!("stale:{}:plan state {} is not committable", plan.plan_id, state_s(plan.state)))); } @@ -460,12 +469,14 @@ impl SqliteWorkspaceStore { if plan.metadata_disposition==MetadataDisposition::Tombstone{ tx.execute("INSERT OR IGNORE INTO worker_tombstones(workspace_id,runtime_id,worker_id,display_name,profile,worker_created_at,removed_at,archive_id,policy_id,policy_revision,operation_id) VALUES(?1,?2,?3,?4,?5,?6,?7,?8,?9,?10,?11)",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,worker.display_name,worker.profile,worker.created_at,now,plan.archive_id,plan.policy_id,plan.policy_revision,operation_id])?; } + tx.execute("UPDATE worker_workdir_links SET unlinked_at=?4 WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND unlinked_at IS NULL",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,now])?; let deleted=tx.execute("DELETE FROM worker_registry WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND updated_at=?4",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,plan.worker_revision])?; tx.execute("UPDATE worker_create_reservations SET state='removed',updated_at=?4 WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND state='created'",params![workspace_id,plan.worker.runtime_id,plan.worker.worker_id,now])?; if deleted!=1{return Err(StoreError::InvalidInput(format!("stale:{}:removal fence changed",plan.plan_id)));} + let catalog=commit_worker_catalog_removal(&tx,workspace_id,&plan.worker)?; tx.execute("UPDATE worker_removal_operations SET state='succeeded',failure_category=NULL,updated_at=?1 WHERE operation_id=?2",params![now,operation_id])?; tx.execute("INSERT OR IGNORE INTO worker_retention_audit_events(event_id,operation_id,workspace_id,event_kind,detail,created_at) VALUES(?1,?2,?3,'worker_removed',?4,?5)",params![stable("wre",operation_id),operation_id,workspace_id,format!("runtime_id={} worker_id={} session={} metadata={} diagnostics={}",plan.worker.runtime_id,plan.worker.worker_id,sess(plan.session_disposition),meta(plan.metadata_disposition),diag(plan.diagnostics_disposition)),now])?; - tx.commit()?;plan.state=WorkerRemovalPlanState::Succeeded;plan.updated_at=now;Ok(plan) + tx.commit()?;plan.state=WorkerRemovalPlanState::Succeeded;plan.updated_at=now;Ok(WorkerRemovalCommit { plan, catalog }) }).map_err(map_error) } @@ -948,7 +959,9 @@ fn parse_state(v: &str) -> rusqlite::Result { #[cfg(test)] mod tests { use super::*; - use crate::store::{ControlPlaneStore, TicketCoderAssignmentRecord, WorkerRegistryRecord}; + use crate::store::{ + ControlPlaneStore, TicketCoderAssignmentRecord, WorkerCatalogChange, WorkerRegistryRecord, + }; use worker_runtime::identity::WorkerId; fn worker_id() -> WorkerId { WorkerId::from_legacy_u64(1) @@ -1203,6 +1216,7 @@ mod tests { assert_eq!( s.commit_worker_removal("w", &p.operation_id, &p.input_fingerprint, &result) .unwrap() + .plan .state, WorkerRemovalPlanState::Succeeded ); @@ -1219,6 +1233,7 @@ mod tests { assert_eq!( s.commit_worker_removal("w", &p.operation_id, &p.input_fingerprint, &result) .unwrap() + .plan .state, WorkerRemovalPlanState::Succeeded ); @@ -1351,6 +1366,98 @@ mod tests { .unwrap(); assert!(s.worker_tombstone("w", &p.worker).unwrap().is_none()); } + #[test] + fn terminal_catalog_failure_rolls_back_removal_projection_and_operation() { + let store = setup(); + let purge = WorkerRetentionPolicyUpdate { + policy_id: "purge".into(), + session_disposition: SessionDisposition::Purge, + metadata_disposition: MetadataDisposition::Tombstone, + archive_retention: ArchiveRetention::Forever, + diagnostics_disposition: DiagnosticsDisposition::Purge, + diagnostics_retention_seconds: None, + }; + store + .update_worker_retention_policy("w", 1, &purge) + .unwrap(); + let plan = store.plan_worker_removal(&req(), &inv()).unwrap(); + store + .prepare_worker_removal_execution("w", &plan.plan_id, &plan.input_fingerprint) + .unwrap(); + let result = WorkerRetentionExecutionResult { + operation_id: plan.operation_id.clone(), + input_fingerprint: plan.input_fingerprint.clone(), + expected_worker_revision: plan.worker_revision.clone(), + worker_id: worker_id(), + session_disposition: SessionDisposition::Purge, + diagnostics_disposition: DiagnosticsDisposition::Purge, + archive: None, + source_removed: true, + diagnostics_retained: false, + }; + store + .with_conn(|conn| { + conn.execute_batch( + "CREATE TRIGGER fail_worker_projection_removal + BEFORE INSERT ON worker_registry_projection_removals + BEGIN SELECT RAISE(ABORT, 'injected projection failure'); END;", + )?; + Ok(()) + }) + .unwrap(); + + assert!( + store + .commit_worker_removal("w", &plan.operation_id, &plan.input_fingerprint, &result,) + .is_err() + ); + let (registry, removal, state, tombstone): (i64, i64, String, i64) = store + .with_conn(|conn| { + Ok(( + conn.query_row( + "SELECT COUNT(*) FROM worker_registry WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + |row| row.get(0), + )?, + conn.query_row( + "SELECT COUNT(*) FROM worker_registry_projection_removals WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + |row| row.get(0), + )?, + conn.query_row( + "SELECT state FROM worker_removal_operations WHERE operation_id=?1", + [&plan.operation_id], + |row| row.get(0), + )?, + conn.query_row( + "SELECT COUNT(*) FROM worker_tombstones WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + |row| row.get(0), + )?, + )) + }) + .unwrap(); + assert_eq!(registry, 1, "catalog deletion must roll back"); + assert_eq!(removal, 0, "removal fence must not partially commit"); + assert_eq!(state, "executing", "operation must remain recoverable"); + assert_eq!(tombstone, 0, "retention metadata must roll back"); + + store + .with_conn(|conn| { + conn.execute_batch("DROP TRIGGER fail_worker_projection_removal;")?; + Ok(()) + }) + .unwrap(); + let committed = store + .commit_worker_removal("w", &plan.operation_id, &plan.input_fingerprint, &result) + .unwrap(); + assert!(matches!( + committed.catalog.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &plan.worker + )); + assert_eq!(committed.plan.state, WorkerRemovalPlanState::Succeeded); + } + #[test] fn commit_requires_executing_state_and_exact_runtime_result() { let store = setup(); @@ -1597,6 +1704,26 @@ mod tests { .unwrap() .is_none() ); + s.with_conn(|conn| { + conn.execute( + "DELETE FROM worker_registry_projection_removals WHERE workspace_id='w' AND runtime_id='r' AND worker_id=?1", + [worker_id().to_string()], + )?; + Ok(()) + }) + .unwrap(); + let repaired = s + .recover_succeeded_worker_removal_catalog("w", &request.worker) + .unwrap(); + assert!(matches!( + repaired.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &request.worker + )); + assert!( + s.get_worker_registry("w", &request.worker) + .unwrap() + .is_none() + ); let recovered = s .recover_worker_removal_execution("w", &request.worker) .unwrap() diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 47fff2f3f..16d3a899c 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -659,7 +659,7 @@ impl From<&WorkspaceApi> for ServerAuthApi { } #[derive(Clone)] -struct WorkspaceWorkerRemoveExecutor { +struct WorkerRemovalService { workspace_id: String, store: Arc, runtime: Weak, @@ -667,10 +667,10 @@ struct WorkspaceWorkerRemoveExecutor { workdir_session_locks: Arc>>>>, worker_remove_locks: Arc>>>>, worker_control_locks: Arc>>>>, - worker_projection: std::sync::Weak, + worker_projection: Arc, } -impl WorkspaceWorkerRemoveExecutor { +impl WorkerRemovalService { fn new(api: &WorkspaceApi) -> Self { Self { workspace_id: api.config.workspace_id.clone(), @@ -680,10 +680,168 @@ impl WorkspaceWorkerRemoveExecutor { workdir_session_locks: api.workdir_session_locks.clone(), worker_remove_locks: api.worker_remove_locks.clone(), worker_control_locks: api.worker_control_locks.clone(), - worker_projection: Arc::downgrade(&api.worker_projection), + worker_projection: api.worker_projection.clone(), } } + /// Terminal catalog deletion shared by cleanup and spawn compensation. + /// The Store atomically commits membership, reservation, link and projection + /// state; this service owns the only production-side publication step. + fn commit_catalog_removal(&self, target: &RuntimeWorkerRef) -> crate::Result<()> { + let commit = self + .store + .delete_worker_registry(&self.workspace_id, target)?; + self.worker_projection.publish_commit(commit) + } + + async fn execute_cleanup_removal(&self, candidate: &CleanupWorkerCandidate) -> ApiResult<()> { + let target = RuntimeWorkerRef::new( + candidate.runtime_id.clone(), + candidate.runtime_worker_id.clone(), + ); + let remove_lock = { + let mut locks = self + .worker_remove_locks + .lock() + .map_err(|_| Error::Store("Worker removal lock registry was poisoned".into()))?; + locks + .entry(target.clone()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + }; + let _remove_guard = remove_lock.lock().await; + let session_lock = { + let mut locks = self + .workdir_session_locks + .lock() + .map_err(|_| Error::Store("Workdir session lock registry was poisoned".into()))?; + locks + .entry(target.clone()) + .or_insert_with(|| Arc::new(tokio::sync::Mutex::new(()))) + .clone() + }; + let _session_guard = session_lock.lock().await; + close_worker_workdir_sessions(&self.workdir_sessions, &target) + .await + .map_err(|message| Error::RuntimeOperationFailed { + runtime_id: target.runtime_id.clone(), + code: "workdir_session_close_failed".to_string(), + message, + })?; + let runtime = self.runtime.upgrade().ok_or_else(|| { + Error::Store("Workspace Runtime registry is unavailable during cleanup".to_string()) + })?; + match runtime.stop_worker( + &target, + WorkerLifecycleRequest { + reason: Some("cleanup worker before deletion".to_string()), + ticket_assignment: None, + }, + ) { + Ok(result) if result.state == WorkerOperationState::Accepted => {} + Ok(result) => { + return Err(ApiError::with_diagnostics( + Error::RuntimeOperationFailed { + runtime_id: target.runtime_id.clone(), + code: "workspace_cleanup_worker_runtime_stop_rejected".to_string(), + message: "Runtime did not stop selected Worker before deletion".to_string(), + }, + result.diagnostics, + )); + } + Err(RuntimeRegistryError::UnknownWorker { .. }) => {} + Err(error) => return Err(error.into_error().into()), + } + match runtime.delete_worker(&target) { + Ok(result) if result.deleted && result.state == WorkerOperationState::Accepted => {} + Ok(result) => { + return Err(ApiError::with_diagnostics( + Error::RuntimeOperationFailed { + runtime_id: target.runtime_id.clone(), + code: "workspace_cleanup_worker_runtime_delete_rejected".to_string(), + message: "Runtime did not delete selected Worker after stopping it" + .to_string(), + }, + result.diagnostics, + )); + } + Err(RuntimeRegistryError::UnknownWorker { .. }) => {} + Err(error) => return Err(error.into_error().into()), + } + self.commit_catalog_removal(&target)?; + Ok(()) + } + + fn delete_spawn_compensation_runtime_worker( + &self, + worker_ref: &RuntimeWorkerRef, + ) -> (bool, Vec) { + let mut diagnostics = Vec::new(); + let Some(runtime) = self.runtime.upgrade() else { + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_runtime_delete_failed", + "Workspace Runtime registry is unavailable during spawn compensation".to_string(), + )); + return (false, diagnostics); + }; + let lifecycle_request = WorkerLifecycleRequest { + reason: Some("Backend spawn finalize failed; compensating Runtime Worker".to_string()), + ticket_assignment: None, + }; + let cancellation = runtime.cancel_worker(worker_ref, lifecycle_request.clone()); + let stop = runtime.stop_worker(worker_ref, lifecycle_request); + let stop_accepted = stop + .as_ref() + .is_ok_and(|result| result.state == WorkerOperationState::Accepted); + let termination_detail = (!stop_accepted).then(|| { + let cancellation = lifecycle_failure_detail("cancel", &cancellation); + let stop = lifecycle_failure_detail("stop", &stop); + format!("{cancellation}; {stop}") + }); + + let runtime_deleted = match runtime.delete_worker(worker_ref) { + Ok(result) if result.state == WorkerOperationState::Accepted && result.deleted => true, + Ok(result) => { + let mut message = format!( + "Runtime did not delete Worker {}:{}: state={:?}, deleted={}", + worker_ref.runtime_id, worker_ref.worker_id, result.state, result.deleted + ); + if let Some(detail) = termination_detail.as_deref() { + message.push_str(&format!("; cancellation: {detail}")); + } + if !result.diagnostics.is_empty() { + message.push_str(&format!( + "; delete diagnostics: {}", + runtime_diagnostics_message(&result.diagnostics) + )); + } + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_runtime_delete_failed", + message, + )); + false + } + Err(RuntimeRegistryError::UnknownWorker { .. }) => true, + Err(error) => { + let mut message = format!( + "Failed to delete Runtime Worker {}:{}: {}", + worker_ref.runtime_id, + worker_ref.worker_id, + error.message() + ); + if let Some(detail) = termination_detail.as_deref() { + message.push_str(&format!("; cancellation: {detail}")); + } + diagnostics.push(spawn_compensation_diagnostic( + "worker_spawn_compensation_runtime_delete_failed", + message, + )); + false + } + }; + (runtime_deleted, diagnostics) + } + async fn resume_worker_retention( &self, runtime: &RuntimeRegistry, @@ -707,10 +865,10 @@ impl WorkspaceWorkerRemoveExecutor { &prepared.plan.input_fingerprint, &result, ) { - Ok(_) => { - if let Some(worker_projection) = self.worker_projection.upgrade() { - let _ = worker_projection.publish_removed(target.clone()); - } + Ok(commit) => { + self.worker_projection + .publish_commit(commit.catalog) + .map_err(|error| error.to_string())?; Ok(worker_remove_success_response(target)) } Err(error) => Ok(worker_retention_error_response(error)), @@ -834,7 +992,14 @@ impl WorkspaceWorkerRemoveExecutor { .map_err(|_| "Worker removal recovery authority is unavailable".to_string())?; if let Some(prepared) = prepared { if prepared.plan.state == crate::retention::WorkerRemovalPlanState::Succeeded { - return Ok(worker_remove_success_response(&target)); + let commit = self + .store + .recover_succeeded_worker_removal_catalog(&self.workspace_id, target) + .map_err(|error| error.to_string())?; + self.worker_projection + .publish_commit(commit) + .map_err(|error| error.to_string())?; + return Ok(worker_remove_success_response(target)); } let prepared = if matches!( prepared.plan.state, @@ -1023,10 +1188,10 @@ impl WorkspaceWorkerRemoveExecutor { &plan.input_fingerprint, &retention_result, ) { - Ok(_) => { - if let Some(worker_projection) = self.worker_projection.upgrade() { - let _ = worker_projection.publish_removed(plan.worker.clone()); - } + Ok(commit) => { + self.worker_projection + .publish_commit(commit.catalog) + .map_err(|error| error.to_string())?; } Err(error) => { let _ = self.store.fail_worker_removal( @@ -1042,7 +1207,7 @@ impl WorkspaceWorkerRemoveExecutor { } } -impl crate::worker_source::VerifiedWorkerRemoveExecutor for WorkspaceWorkerRemoveExecutor { +impl crate::worker_source::VerifiedWorkerRemoveExecutor for WorkerRemovalService { fn execute( &self, source: crate::worker_source::VerifiedWorkerMutationSource, @@ -1252,7 +1417,7 @@ impl WorkspaceServerApi { { let worker_key = worker.display_name.clone(); let target = worker.worker; - let response = WorkspaceWorkerRemoveExecutor::new(&api) + let response = WorkerRemovalService::new(&api) .execute_target_removal( api.runtime.as_ref(), &target, @@ -2473,7 +2638,7 @@ impl WorkspaceApi { }; if let Some(dispatcher) = worker_remove_dispatcher { dispatcher - .install_executor(Arc::new(WorkspaceWorkerRemoveExecutor::new(&api))) + .install_executor(Arc::new(WorkerRemovalService::new(&api))) .map_err(|message| Error::Config(message.to_string()))?; } recover_workdir_removals(&api)?; @@ -8530,7 +8695,7 @@ async fn scoped_attach_current_worker_workdir( return Err(error.into()); } api.worker_projection - .publish_catalog_change(&worker) + .refresh(&worker) .map_err(ApiError::from)?; Ok(Json(CurrentWorkerWorkdirAttachmentResponse { workspace_id: api.config.workspace_id.clone(), @@ -8557,7 +8722,7 @@ async fn scoped_detach_current_worker_workdir( &Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), )?; api.worker_projection - .publish_catalog_change(&worker) + .refresh(&worker) .map_err(ApiError::from)?; Ok(Json(CurrentWorkerWorkdirAttachmentResponse { workspace_id: api.config.workspace_id.clone(), @@ -10008,7 +10173,7 @@ async fn scoped_worker_remove_source_boundary( .await { Ok(source) => { - let executor = WorkspaceWorkerRemoveExecutor::new(&api); + let executor = WorkerRemovalService::new(&api); match executor .execute_async( source, @@ -11903,19 +12068,22 @@ async fn set_worker_retention( } } let retention_state = if pinned { "pinned" } else { "normal" }; - let changed = api.store.update_worker_retention( + let commit = api.store.update_worker_retention( &api.config.workspace_id, &worker_ref, retention_state, now_registry_timestamp().as_str(), )?; - if !changed { + if commit.changes.is_empty() { return Err(cleanup_api_error( runtime_id.as_str(), "workspace_worker_retention_unknown_worker", "Worker is not known to the Backend registry", )); } + api.worker_projection + .publish_commit(commit) + .map_err(ApiError::from)?; Ok(Json(WorkerRetentionResponse { workspace_id: api.config.workspace_id, worker_ref, @@ -12196,13 +12364,9 @@ async fn execute_runtime_cleanup( )); } parse_runtime_worker_id_for_registry(&candidate.runtime_worker_id)?; - let session_lock = current_worker_session_lock(api, &worker); - let _session_guard = session_lock.lock().await; - close_current_worker_session_locked(api, &worker).await?; - cleanup_runtime_worker_for_execution(api, runtime_id, candidate)?; - api.store - .delete_worker_registry(&api.config.workspace_id, &worker)?; - drop(_session_guard); + WorkerRemovalService::new(api) + .execute_cleanup_removal(candidate) + .await?; api.workdir_session_locks .lock() .expect("Workdir session lock registry poisoned") @@ -12297,49 +12461,6 @@ async fn execute_runtime_cleanup( }) } -fn cleanup_runtime_worker_for_execution( - api: &WorkspaceApi, - runtime_id: &str, - candidate: &CleanupWorkerCandidate, -) -> ApiResult<()> { - let worker = RuntimeWorkerRef::new(runtime_id, &candidate.runtime_worker_id); - match api.runtime.stop_worker( - &worker, - WorkerLifecycleRequest { - reason: Some("cleanup worker before deletion".to_string()), - ticket_assignment: None, - }, - ) { - Ok(result) if result.state == WorkerOperationState::Accepted => {} - Ok(result) => { - return Err(ApiError::with_diagnostics( - Error::RuntimeOperationFailed { - runtime_id: runtime_id.to_string(), - code: "workspace_cleanup_worker_runtime_stop_rejected".to_string(), - message: "Runtime did not stop selected Worker before deletion".to_string(), - }, - result.diagnostics, - )); - } - Err(RuntimeRegistryError::UnknownWorker { .. }) => return Ok(()), - Err(error) => return Err(error.into_error().into()), - } - - match api.runtime.delete_worker(&worker) { - Ok(result) if result.deleted && result.state == WorkerOperationState::Accepted => Ok(()), - Ok(result) => Err(ApiError::with_diagnostics( - Error::RuntimeOperationFailed { - runtime_id: runtime_id.to_string(), - code: "workspace_cleanup_worker_runtime_delete_rejected".to_string(), - message: "Runtime did not delete selected Worker after stopping it".to_string(), - }, - result.diagnostics, - )), - Err(RuntimeRegistryError::UnknownWorker { .. }) => Ok(()), - Err(error) => Err(error.into_error().into()), - } -} - fn cleanup_api_error(runtime_id: &str, code: &str, message: &str) -> ApiError { Error::RuntimeOperationFailed { runtime_id: runtime_id.to_string(), @@ -15690,7 +15811,7 @@ fn compensate_failed_workspace_worker_create( let mut reserved_worker_absent = false; if let Some(worker) = worker { let (returned_worker_absent, delete_diagnostics) = - delete_runtime_worker_for_spawn_compensation(api, &worker.worker); + WorkerRemovalService::new(api).delete_spawn_compensation_runtime_worker(&worker.worker); diagnostics.extend(delete_diagnostics); if returned_worker_absent { diagnostics.extend(finalize_spawn_compensation_after_worker_delete( @@ -15702,8 +15823,8 @@ fn compensate_failed_workspace_worker_create( } } if !reserved_worker_absent { - let (worker_absent, delete_diagnostics) = - delete_runtime_worker_for_spawn_compensation(api, &reserved_worker_ref); + let (worker_absent, delete_diagnostics) = WorkerRemovalService::new(api) + .delete_spawn_compensation_runtime_worker(&reserved_worker_ref); reserved_worker_absent = worker_absent; diagnostics.extend(delete_diagnostics); } @@ -15762,7 +15883,7 @@ fn compensate_failed_worker_spawn( context: &WorkerSpawnCompensationContext<'_>, ) -> Vec { let (runtime_deleted, mut diagnostics) = - delete_runtime_worker_for_spawn_compensation(api, &worker.worker); + WorkerRemovalService::new(api).delete_spawn_compensation_runtime_worker(&worker.worker); if !runtime_deleted { return diagnostics; } @@ -15772,71 +15893,6 @@ fn compensate_failed_worker_spawn( diagnostics } -fn delete_runtime_worker_for_spawn_compensation( - api: &WorkspaceApi, - worker_ref: &RuntimeWorkerRef, -) -> (bool, Vec) { - let mut diagnostics = Vec::new(); - let lifecycle_request = WorkerLifecycleRequest { - reason: Some("Backend spawn finalize failed; compensating Runtime Worker".to_string()), - ticket_assignment: None, - }; - let cancellation = api - .runtime - .cancel_worker(worker_ref, lifecycle_request.clone()); - let stop = api.runtime.stop_worker(worker_ref, lifecycle_request); - let stop_accepted = stop - .as_ref() - .is_ok_and(|result| result.state == WorkerOperationState::Accepted); - let termination_detail = (!stop_accepted).then(|| { - let cancellation = lifecycle_failure_detail("cancel", &cancellation); - let stop = lifecycle_failure_detail("stop", &stop); - format!("{cancellation}; {stop}") - }); - - let runtime_deleted = match api.runtime.delete_worker(worker_ref) { - Ok(result) if result.state == WorkerOperationState::Accepted && result.deleted => true, - Ok(result) => { - let mut message = format!( - "Runtime did not delete Worker {}:{}: state={:?}, deleted={}", - worker_ref.runtime_id, worker_ref.worker_id, result.state, result.deleted - ); - if let Some(detail) = termination_detail.as_deref() { - message.push_str(&format!("; cancellation: {detail}")); - } - if !result.diagnostics.is_empty() { - message.push_str(&format!( - "; delete diagnostics: {}", - runtime_diagnostics_message(&result.diagnostics) - )); - } - diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_runtime_delete_failed", - message, - )); - false - } - Err(RuntimeRegistryError::UnknownWorker { .. }) => true, - Err(error) => { - let mut message = format!( - "Failed to delete Runtime Worker {}:{}: {}", - worker_ref.runtime_id, - worker_ref.worker_id, - error.message() - ); - if let Some(detail) = termination_detail.as_deref() { - message.push_str(&format!("; cancellation: {detail}")); - } - diagnostics.push(spawn_compensation_diagnostic( - "worker_spawn_compensation_runtime_delete_failed", - message, - )); - false - } - }; - (runtime_deleted, diagnostics) -} - fn finalize_spawn_compensation_after_worker_delete( api: &WorkspaceApi, worker: &WorkerSummary, @@ -15860,10 +15916,7 @@ fn finalize_spawn_compensation_after_worker_delete( )); } } - if let Err(error) = api - .store - .delete_worker_registry(&api.config.workspace_id, &worker.worker) - { + if let Err(error) = WorkerRemovalService::new(api).commit_catalog_removal(&worker.worker) { diagnostics.push(spawn_compensation_diagnostic( "worker_spawn_compensation_registry_delete_failed", format!( @@ -17554,7 +17607,10 @@ fn record_worker_summary( created_at: timestamp.clone(), updated_at: timestamp, }; - api.store.upsert_worker_registry(&record)?; + let catalog_commit = api.store.upsert_worker_registry(&record)?; + api.worker_projection + .publish_commit(catalog_commit) + .map_err(ApiError::from)?; if let Ok(worker_id) = protocol::subscription::SubscriptionWorkerId::new(worker.worker.worker_id.clone()) { @@ -17590,10 +17646,6 @@ fn record_worker_summary( .unwrap_or(record.updated_at.as_str()), ) .map_err(ApiError::from)?; - } else { - api.worker_projection - .publish_catalog_change(&worker_ref) - .map_err(ApiError::from)?; } Ok(api .store @@ -18133,7 +18185,7 @@ fn link_worker_to_workdir( api.store.attach_worker_workdir(&record)?; } api.worker_projection - .publish_catalog_change(&worker_record.worker) + .refresh(&worker_record.worker) .map_err(ApiError::from)?; Ok(()) } @@ -26994,7 +27046,7 @@ mod tests { permission: worker_runtime::auth::WORKER_REMOVE_PERMISSION.to_string(), jti: "caller-guard-proof".to_string(), }; - let executor = WorkspaceWorkerRemoveExecutor::new(&api); + let executor = WorkerRemovalService::new(&api); let self_response = executor .execute_async( verified_source(), @@ -27125,8 +27177,9 @@ mod tests { let summary = api.runtime.worker(&target).unwrap(); sync_worker_observation(&api, &summary).unwrap(); seed_worker_control_grant(&api, &source, &target, "embedded-valid-proof"); + let mut subscriber = api.worker_projection.subscribe(); - let response = WorkspaceWorkerRemoveExecutor::new(&api) + let response = WorkerRemovalService::new(&api) .execute_async( crate::worker_source::VerifiedWorkerMutationSource { runtime_id: source.runtime_id, @@ -27157,6 +27210,84 @@ mod tests { .unwrap() .is_none() ); + let removal = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let event = subscriber.recv().await.expect("projection event"); + if event.changes.iter().any(|change| { + matches!(change, crate::store::WorkerCatalogChange::Removed(worker) if worker == &target) + }) { + return event; + } + } + }) + .await + .expect("WorkerRemove catalog removal broadcast"); + assert!(removal.revision > 0); + } + + #[tokio::test] + async fn published_spawn_compensation_uses_catalog_removal_commit() { + let temp = tempfile::tempdir().unwrap(); + let api = test_api(temp.path()).await; + let spawned = api + .spawn_workspace_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("compensation-target".to_string()), + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: worker_runtime::catalog::ProfileSelector::Builtin( + "builtin:companion".to_string(), + ), + ticket_assignment: None, + initial_submit: Vec::new(), + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: Some(runtime_test_bundle()), + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }, + ) + .unwrap(); + let worker = spawned.worker.unwrap().worker; + let summary = api.runtime.worker(&worker).unwrap(); + sync_worker_observation(&api, &summary).unwrap(); + let mut subscriber = api.worker_projection.subscribe(); + let diagnostics = compensate_failed_worker_spawn( + &api, + &summary, + &WorkerSpawnCompensationContext { + assignment: None, + prepared_workdir_id: None, + cleanup_spawned_workdir: false, + }, + ); + assert!(diagnostics.is_empty(), "{diagnostics:?}"); + assert!( + api.store + .get_worker_registry(TEST_WORKSPACE_ID, &worker) + .unwrap() + .is_none() + ); + let removal = tokio::time::timeout(std::time::Duration::from_secs(1), async { + loop { + let event = subscriber.recv().await.expect("projection event"); + if event.changes.iter().any(|change| { + matches!(change, crate::store::WorkerCatalogChange::Removed(removed) if removed == &worker) + }) { + return event; + } + } + }) + .await + .expect("spawn compensation catalog removal broadcast"); + assert!(removal.revision > 0); } #[tokio::test] @@ -28534,6 +28665,121 @@ mod tests { ); } + #[tokio::test] + async fn cleanup_http_entry_publishes_atomic_catalog_removal_without_runtime_event() { + let workspace = tempfile::tempdir().unwrap(); + init_clean_git_workspace(workspace.path()); + let api = test_api(workspace.path()).await; + let spawned = api + .spawn_workspace_worker( + EMBEDDED_WORKER_RUNTIME_ID, + WorkerSpawnRequest { + intent: WorkerSpawnIntent::WorkspaceCompanion, + requested_worker_name: Some("cleanup-projection-target".to_string()), + acceptance: WorkerSpawnAcceptanceRequirement::RunAccepted { + expected_segments: 0, + }, + profile: worker_runtime::catalog::ProfileSelector::Builtin( + "builtin:companion".to_string(), + ), + ticket_assignment: None, + initial_submit: Vec::new(), + working_directory_request: None, + resolved_working_directory_request: None, + resolved_working_directory: None, + resolved_config_bundle: Some(runtime_test_bundle()), + resolved_worker_observation_enabled: false, + resolved_worker_observation_grants: Vec::new(), + resolved_workspace_api: None, + resolved_memory_settings: None, + resolved_control_operation: None, + }, + ) + .unwrap(); + let worker = spawned.worker.unwrap().worker; + let stopped = api + .runtime + .stop_worker( + &worker, + WorkerLifecycleRequest { + reason: Some("prepare cleanup projection regression".to_string()), + ticket_assignment: None, + }, + ) + .unwrap(); + assert_eq!(stopped.state, WorkerOperationState::Accepted); + let summary = api.runtime.worker(&worker).unwrap(); + sync_worker_observation(&api, &summary).unwrap(); + let worker_id = worker.worker_id.clone(); + let mut subscriber = api.worker_projection.subscribe(); + let plan = build_runtime_cleanup_plan(&api, EMBEDDED_WORKER_RUNTIME_ID) + .unwrap_or_else(|error| panic!("cleanup plan: {}", error.error)); + let candidate = plan + .workers + .iter() + .find(|candidate| candidate.worker_id == worker_id) + .unwrap(); + let request = ExecuteRuntimeCleanupRequest { + expected_plan_revision: plan.revision, + expected_plan_digest: plan.digest, + worker_target_ids: vec![candidate.target_id.clone()], + workdir_target_ids: Vec::new(), + confirm_dirty_discard_target_ids: Vec::new(), + }; + + let Json(response) = scoped_execute_runtime_cleanup( + State(api.clone()), + AxumPath(ScopedRuntimePath { + workspace_id: TEST_WORKSPACE_ID.to_string(), + runtime_id: EMBEDDED_WORKER_RUNTIME_ID.to_string(), + }), + Json(request), + ) + .await + .unwrap(); + assert_eq!(response.results[0].status, "deleted"); + + let event = tokio::time::timeout(std::time::Duration::from_secs(1), subscriber.recv()) + .await + .expect("catalog removal broadcast") + .expect("projection event"); + assert!(matches!( + event.changes.as_slice(), + [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker + )); + assert!( + api.store + .get_worker_registry(&api.config.workspace_id, &worker) + .unwrap() + .is_none() + ); + let (snapshot_revision, snapshot) = api + .store + .worker_registry_projection_snapshot(&api.config.workspace_id, 100) + .unwrap(); + assert_eq!(snapshot_revision, event.revision); + assert!( + snapshot + .iter() + .all(|record| record.registry.worker != worker) + ); + + // A retry replays the durable fence at the same revision, healing an + // in-process broadcast omitted after the first DB commit. + WorkerRemovalService::new(&api) + .commit_catalog_removal(&worker) + .unwrap(); + let replay = tokio::time::timeout(std::time::Duration::from_secs(1), subscriber.recv()) + .await + .expect("catalog removal replay") + .expect("projection replay"); + assert_eq!(replay.revision, event.revision); + assert!(matches!( + replay.changes.as_slice(), + [crate::store::WorkerCatalogChange::Removed(removed)] if removed == &worker + )); + } + #[tokio::test] async fn cleanup_blocks_assigned_worker_before_runtime_deletion() { let workspace = tempfile::tempdir().unwrap(); @@ -32102,7 +32348,7 @@ mod tests { .unwrap(); projection_api .worker_projection - .publish_catalog_change(&RuntimeWorkerRef::new( + .refresh(&RuntimeWorkerRef::new( EMBEDDED_WORKER_RUNTIME_ID, &worker_id, )) diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 1e0f0cf66..edf027ec4 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -563,10 +563,27 @@ pub struct WorkerRegistryProjectionRecord { pub observation: Option, } +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum WorkerCatalogChange { + Upsert(WorkerRegistryProjectionRecord), + Removed(RuntimeWorkerRef), +} + +/// A catalog mutation and the exact projection payload committed at its revision. +/// +/// Callers may broadcast this value, but must never re-read catalog rows to infer +/// what happened: a later catalog mutation could otherwise be paired with an +/// older revision. #[derive(Debug, Clone, PartialEq, Eq)] pub struct WorkerRegistryProjectionCommit { pub revision: u64, - pub changed_workers: Vec, + pub changes: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct WorkerRemovalCommit { + pub plan: crate::retention::WorkerRemovalPlan, + pub catalog: WorkerRegistryProjectionCommit, } /// Durable authority describing which Runtime Worker another Runtime Worker may @@ -1130,16 +1147,23 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { "Worker retention authority is unavailable".to_string(), )) } + fn recover_succeeded_worker_removal_catalog( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result { + let _ = (workspace_id, worker); + Err(Error::Store( + "Worker removal catalog recovery authority is unavailable".to_string(), + )) + } fn commit_worker_removal( &self, workspace_id: &str, operation_id: &str, input_fingerprint: &str, result: &worker_runtime::retention::WorkerRetentionExecutionResult, - ) -> std::result::Result< - crate::retention::WorkerRemovalPlan, - crate::retention::WorkerRetentionError, - > { + ) -> std::result::Result { let _ = (workspace_id, operation_id, input_fingerprint, result); Err(crate::retention::WorkerRetentionError::Invalid( "Worker retention authority is unavailable".to_string(), @@ -1317,7 +1341,10 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { consumed_at: &str, ) -> Result>; - fn upsert_worker_registry(&self, record: &WorkerRegistryRecord) -> Result<()>; + fn upsert_worker_registry( + &self, + record: &WorkerRegistryRecord, + ) -> Result; fn get_worker_registry( &self, workspace_id: &str, @@ -1369,12 +1396,7 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { subject_revision: Option, observed_at: &str, ) -> Result; - fn publish_worker_registry_catalog_change( - &self, - workspace_id: &str, - worker: &RuntimeWorkerRef, - ) -> Result; - fn publish_worker_registry_removal( + fn commit_worker_registry_projection_refresh( &self, workspace_id: &str, worker: &RuntimeWorkerRef, @@ -1385,9 +1407,12 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore { worker: &RuntimeWorkerRef, retention_state: &str, updated_at: &str, - ) -> Result; - fn delete_worker_registry(&self, workspace_id: &str, worker: &RuntimeWorkerRef) - -> Result; + ) -> Result; + fn delete_worker_registry( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result; fn create_worker_control_grant( &self, @@ -4330,16 +4355,36 @@ impl ControlPlaneStore for SqliteWorkspaceStore { ) } + fn recover_succeeded_worker_removal_catalog( + &self, + workspace_id: &str, + worker: &RuntimeWorkerRef, + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let succeeded: bool = tx.query_row( + "SELECT EXISTS(SELECT 1 FROM worker_removal_operations WHERE workspace_id=?1 AND runtime_id=?2 AND worker_id=?3 AND state='succeeded')", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get(0), + )?; + if !succeeded { + return Err(Error::InvalidInput( + "Worker removal operation has not succeeded".to_string(), + )); + } + let commit = commit_worker_catalog_removal(&tx, workspace_id, worker)?; + tx.commit()?; + Ok(commit) + }) + } + fn commit_worker_removal( &self, workspace_id: &str, operation_id: &str, input_fingerprint: &str, result: &worker_runtime::retention::WorkerRetentionExecutionResult, - ) -> std::result::Result< - crate::retention::WorkerRemovalPlan, - crate::retention::WorkerRetentionError, - > { + ) -> std::result::Result { SqliteWorkspaceStore::commit_worker_removal( self, workspace_id, @@ -5527,7 +5572,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } - fn upsert_worker_registry(&self, record: &WorkerRegistryRecord) -> Result<()> { + fn upsert_worker_registry( + &self, + record: &WorkerRegistryRecord, + ) -> Result { self.with_conn(|conn| { let tx = conn.unchecked_transaction()?; let removal_blocks_upsert: bool = tx.query_row( @@ -5545,9 +5593,12 @@ impl ControlPlaneStore for SqliteWorkspaceStore { |row| row.get(0), )?; if removal_blocks_upsert { - return Ok(()); + return Ok(WorkerRegistryProjectionCommit { + revision: current_worker_projection_revision(&tx, &record.workspace_id)?, + changes: Vec::new(), + }); } - tx.execute( + let changed = tx.execute( r#"INSERT INTO worker_registry ( workspace_id, runtime_id, worker_id, display_name, profile, retention_state, transcript_ref, session_ref, summary_ref, @@ -5600,8 +5651,20 @@ impl ControlPlaneStore for SqliteWorkspaceStore { "DELETE FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", params![record.workspace_id, record.worker.runtime_id, record.worker.worker_id], )?; + let changed_workers = if changed == 0 { + Vec::new() + } else { + vec![record.worker.clone()] + }; + let revision = + commit_worker_projection_changes(&tx, &record.workspace_id, &changed_workers)?; + let changes = worker_catalog_upsert_changes( + &tx, + &record.workspace_id, + changed_workers.as_slice(), + )?; tx.commit()?; - Ok(()) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5715,7 +5778,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { tx.commit()?; return Ok(WorkerRegistryProjectionCommit { revision, - changed_workers: Vec::new(), + changes: Vec::new(), }); } upsert_worker_projection_cursor( @@ -5775,11 +5838,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { } } let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5810,7 +5872,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { tx.commit()?; return Ok(WorkerRegistryProjectionCommit { revision, - changed_workers: Vec::new(), + changes: Vec::new(), }); } let changed = upsert_worker_observation( @@ -5825,11 +5887,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { )?; let changed_workers = if changed { vec![identity] } else { Vec::new() }; let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5878,15 +5939,14 @@ impl ControlPlaneStore for SqliteWorkspaceStore { } } let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } - fn publish_worker_registry_catalog_change( + fn commit_worker_registry_projection_refresh( &self, workspace_id: &str, worker: &RuntimeWorkerRef, @@ -5899,46 +5959,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { Vec::new() }; let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) - }) - } - - fn publish_worker_registry_removal( - &self, - workspace_id: &str, - worker: &RuntimeWorkerRef, - ) -> Result { - self.with_conn_mut(|conn| { - let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; - let published_revision = tx - .query_row( - "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", - params![workspace_id, worker.runtime_id, worker.worker_id], - |row| row.get::<_, i64>(0), - ) - .optional()?; - if let Some(revision) = published_revision { - tx.commit()?; - return Ok(WorkerRegistryProjectionCommit { - revision: revision.max(0) as u64, - changed_workers: Vec::new(), - }); - } - let changed_workers = vec![worker.clone()]; - let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; - tx.execute( - "INSERT INTO worker_registry_projection_removals (workspace_id, runtime_id, worker_id, projection_revision) VALUES (?1, ?2, ?3, ?4)", - params![workspace_id, worker.runtime_id, worker.worker_id, revision as i64], - )?; - tx.commit()?; - Ok(WorkerRegistryProjectionCommit { - revision, - changed_workers, - }) + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5948,9 +5972,10 @@ impl ControlPlaneStore for SqliteWorkspaceStore { worker: &RuntimeWorkerRef, retention_state: &str, updated_at: &str, - ) -> Result { - self.with_conn(|conn| { - let changed = conn.execute( + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; + let changed = tx.execute( r#"UPDATE worker_registry SET retention_state = ?4, updated_at = ?5 WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3 @@ -5968,7 +5993,16 @@ impl ControlPlaneStore for SqliteWorkspaceStore { updated_at ], )?; - Ok(changed > 0) + let changed_workers = if changed == 0 { + Vec::new() + } else { + vec![worker.clone()] + }; + let revision = commit_worker_projection_changes(&tx, workspace_id, &changed_workers)?; + let changes = + worker_catalog_upsert_changes(&tx, workspace_id, changed_workers.as_slice())?; + tx.commit()?; + Ok(WorkerRegistryProjectionCommit { revision, changes }) }) } @@ -5976,9 +6010,9 @@ impl ControlPlaneStore for SqliteWorkspaceStore { &self, workspace_id: &str, worker: &RuntimeWorkerRef, - ) -> Result { - self.with_conn(|conn| { - let tx = conn.unchecked_transaction()?; + ) -> Result { + self.with_conn_mut(|conn| { + let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?; tx.execute( r#"UPDATE worker_workdir_links SET unlinked_at = strftime('%Y-%m-%dT%H:%M:%fZ', 'now') @@ -5989,7 +6023,7 @@ impl ControlPlaneStore for SqliteWorkspaceStore { "DELETE FROM worker_registry WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", params![workspace_id, worker.runtime_id, worker.worker_id], )?; - if changed != 0 { + let commit = if changed != 0 { tx.execute( "UPDATE worker_create_reservations SET state = 'removed', updated_at = ?4 @@ -6002,9 +6036,27 @@ impl ControlPlaneStore for SqliteWorkspaceStore { chrono::Utc::now().to_rfc3339(), ], )?; - } + commit_worker_catalog_removal(&tx, workspace_id, worker)? + } else if let Some(revision) = tx + .query_row( + "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + WorkerRegistryProjectionCommit { + revision: revision.max(0) as u64, + changes: vec![WorkerCatalogChange::Removed(worker.clone())], + } + } else { + WorkerRegistryProjectionCommit { + revision: current_worker_projection_revision(&tx, workspace_id)?, + changes: Vec::new(), + } + }; tx.commit()?; - Ok(changed > 0) + Ok(commit) }) } @@ -8820,6 +8872,60 @@ fn set_worker_observation_unavailable( Ok(true) } +fn worker_catalog_upsert_changes( + conn: &Connection, + workspace_id: &str, + changed_workers: &[RuntimeWorkerRef], +) -> Result> { + changed_workers + .iter() + .map(|worker| { + read_worker_registry_projection(conn, workspace_id, worker)?.map_or_else( + || { + Err(Error::Store(format!( + "committed Worker catalog row disappeared before projection payload was captured: {}:{}", + worker.runtime_id, worker.worker_id + ))) + }, + |record| Ok(WorkerCatalogChange::Upsert(record)), + ) + }) + .collect() +} + +/// Commits the durable removal fence in the same transaction as catalog deletion. +/// Existing fences are returned as replayable changes so a retry can heal a +/// process-local broadcast missed after the original DB commit. +pub(crate) fn commit_worker_catalog_removal( + conn: &Connection, + workspace_id: &str, + worker: &RuntimeWorkerRef, +) -> Result { + if let Some(revision) = conn + .query_row( + "SELECT projection_revision FROM worker_registry_projection_removals WHERE workspace_id = ?1 AND runtime_id = ?2 AND worker_id = ?3", + params![workspace_id, worker.runtime_id, worker.worker_id], + |row| row.get::<_, i64>(0), + ) + .optional()? + { + return Ok(WorkerRegistryProjectionCommit { + revision: revision.max(0) as u64, + changes: vec![WorkerCatalogChange::Removed(worker.clone())], + }); + } + let revision = + commit_worker_projection_changes(conn, workspace_id, std::slice::from_ref(worker))?; + conn.execute( + "INSERT INTO worker_registry_projection_removals (workspace_id, runtime_id, worker_id, projection_revision) VALUES (?1, ?2, ?3, ?4)", + params![workspace_id, worker.runtime_id, worker.worker_id, revision as i64], + )?; + Ok(WorkerRegistryProjectionCommit { + revision, + changes: vec![WorkerCatalogChange::Removed(worker.clone())], + }) +} + fn commit_worker_projection_changes( conn: &Connection, workspace_id: &str, @@ -13499,9 +13605,11 @@ mod tests { .is_err() ); assert!( - store + !store .delete_worker_registry("workspace-a", &reserved_worker) .unwrap() + .changes + .is_empty() ); let removed_state: String = store .with_conn(|conn| { @@ -14202,9 +14310,11 @@ INSERT INTO worker_registry ( .is_none() ); assert!( - store + !store .delete_worker_registry("workspace-role", &worker.worker) .unwrap() + .changes + .is_empty() ); let removed_worker_assignment = TicketRoleAssignmentRecord { assignment_id: "contributor-removed-worker".to_string(), @@ -14937,7 +15047,10 @@ INSERT INTO worker_registry ( "2", ) .unwrap(); - assert_eq!(first.changed_workers, vec![catalog.worker.clone()]); + assert!(matches!( + first.changes.as_slice(), + [WorkerCatalogChange::Upsert(record)] if record.registry.worker == catalog.worker + )); let (_, records) = store .worker_registry_projection_snapshot("local-dev", 10) .unwrap(); @@ -14959,7 +15072,10 @@ INSERT INTO worker_registry ( let unavailable = store .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[], "3") .unwrap(); - assert_eq!(unavailable.changed_workers, vec![catalog.worker.clone()]); + assert!(matches!( + unavailable.changes.as_slice(), + [WorkerCatalogChange::Upsert(record)] if record.registry.worker == catalog.worker + )); let projection = store .worker_registry_projection("local-dev", &catalog.worker) .unwrap() @@ -14978,7 +15094,7 @@ INSERT INTO worker_registry ( let duplicate = store .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[], "4") .unwrap(); - assert!(duplicate.changed_workers.is_empty()); + assert!(duplicate.changes.is_empty()); let reconnected = store .reconcile_worker_registry_snapshot( @@ -14990,7 +15106,10 @@ INSERT INTO worker_registry ( "5", ) .unwrap(); - assert_eq!(reconnected.changed_workers, vec![catalog.worker.clone()]); + assert!(matches!( + reconnected.changes.as_slice(), + [WorkerCatalogChange::Upsert(record)] if record.registry.worker == catalog.worker + )); let projection = store .worker_registry_projection("local-dev", &catalog.worker) .unwrap() @@ -15003,7 +15122,7 @@ INSERT INTO worker_registry ( let duplicate_reconnect = store .reconcile_worker_registry_snapshot("local-dev", "embedded", 1, 11, &[observed], "6") .unwrap(); - assert!(duplicate_reconnect.changed_workers.is_empty()); + assert!(duplicate_reconnect.changes.is_empty()); } #[tokio::test] @@ -15037,19 +15156,20 @@ INSERT INTO worker_registry ( updated_at: "1".to_string(), }) .unwrap(); - assert!( - store - .delete_worker_registry("local-dev", &worker_ref) - .unwrap() - ); let first_removal = store - .publish_worker_registry_removal("local-dev", &worker_ref) + .delete_worker_registry("local-dev", &worker_ref) .unwrap(); - assert_eq!(first_removal.changed_workers, vec![worker_ref.clone()]); + assert!(matches!( + first_removal.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &worker_ref + )); let duplicate_removal = store - .publish_worker_registry_removal("local-dev", &worker_ref) + .delete_worker_registry("local-dev", &worker_ref) .unwrap(); - assert!(duplicate_removal.changed_workers.is_empty()); + assert!(matches!( + duplicate_removal.changes.as_slice(), + [WorkerCatalogChange::Removed(worker)] if worker == &worker_ref + )); assert_eq!(duplicate_removal.revision, first_removal.revision); let event = SubscriptionWorker { worker_id: SubscriptionWorkerId::new("removed").unwrap(), @@ -15070,7 +15190,7 @@ INSERT INTO worker_registry ( let commit = store .apply_worker_registry_observation("local-dev", "embedded", 2, &event, "2") .unwrap(); - assert!(commit.changed_workers.is_empty()); + assert!(commit.changes.is_empty()); assert!( store .worker_registry_projection("local-dev", &worker_ref) @@ -15194,9 +15314,11 @@ INSERT INTO worker_registry ( .is_empty() ); assert!( - store + !store .delete_worker_registry("workspace-control", &subject_record.worker) .unwrap() + .changes + .is_empty() ); assert!( store diff --git a/crates/workspace-server/src/worker_projection.rs b/crates/workspace-server/src/worker_projection.rs index 36b01a001..22095f8c0 100644 --- a/crates/workspace-server/src/worker_projection.rs +++ b/crates/workspace-server/src/worker_projection.rs @@ -11,18 +11,12 @@ use tokio::task::JoinHandle; use crate::runtime_subscription::{ BrokerSubscriptionEvent, RuntimeSubscriptionBroker, RuntimeSubscriptionBrokerError, }; -use crate::store::{ - ControlPlaneStore, WorkerRegistryProjectionCommit, WorkerRegistryProjectionRecord, -}; +use crate::store::{ControlPlaneStore, WorkerCatalogChange, WorkerRegistryProjectionCommit}; use worker_runtime::identity::RuntimeWorkerRef; const PROJECTION_EVENT_CAPACITY: usize = 256; -#[derive(Debug, Clone)] -pub enum WorkerProjectionChange { - Upsert(WorkerRegistryProjectionRecord), - Removed(RuntimeWorkerRef), -} +pub type WorkerProjectionChange = WorkerCatalogChange; #[derive(Debug, Clone)] pub struct WorkerProjectionEvent { @@ -144,27 +138,13 @@ impl WorkerProjectionService { self.publish_commit(commit) } - pub fn publish_catalog_change(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> { + pub fn refresh(&self, worker: &RuntimeWorkerRef) -> crate::Result<()> { let commit = self .store - .publish_worker_registry_catalog_change(self.workspace_id.as_str(), worker)?; + .commit_worker_registry_projection_refresh(self.workspace_id.as_str(), worker)?; self.publish_commit(commit) } - pub fn publish_removed(&self, worker: RuntimeWorkerRef) -> crate::Result<()> { - let commit = self - .store - .publish_worker_registry_removal(self.workspace_id.as_str(), &worker)?; - if commit.changed_workers.is_empty() { - return Ok(()); - } - let _ = self.events.send(WorkerProjectionEvent { - revision: commit.revision, - changes: vec![WorkerProjectionChange::Removed(worker)], - }); - Ok(()) - } - async fn apply_broker_event( &self, runtime_id: &str, @@ -236,38 +216,21 @@ impl WorkerProjectionService { self.publish_commit(commit) } - fn publish_commit(&self, commit: WorkerRegistryProjectionCommit) -> crate::Result<()> { - if commit.changed_workers.is_empty() { + /// Broadcast an already committed catalog payload. Persistence always happens + /// before this call; a process crash therefore closes subscribers, whose next + /// subscription starts from the durable snapshot. Retrying a removal returns + /// the same durable change and can heal an in-process missed broadcast. + pub(crate) fn publish_commit( + &self, + commit: WorkerRegistryProjectionCommit, + ) -> crate::Result<()> { + if commit.changes.is_empty() { return Ok(()); } - let changes = commit - .changed_workers - .iter() - .filter_map(|worker| { - match self - .store - .worker_registry_projection(self.workspace_id.as_str(), worker) - { - Ok(Some(record)) => Some(WorkerProjectionChange::Upsert(record)), - Ok(None) => None, - Err(error) => { - tracing::warn!( - runtime_id = %worker.runtime_id, - worker_id = %worker.worker_id, - error = %error, - "failed to load committed Worker projection" - ); - None - } - } - }) - .collect::>(); - if !changes.is_empty() { - let _ = self.events.send(WorkerProjectionEvent { - revision: commit.revision, - changes, - }); - } + let _ = self.events.send(WorkerProjectionEvent { + revision: commit.revision, + changes: commit.changes, + }); Ok(()) } } diff --git a/crates/workspace-server/src/workspace_deletion.rs b/crates/workspace-server/src/workspace_deletion.rs index 3dbfa46f7..e4c5f5b54 100644 --- a/crates/workspace-server/src/workspace_deletion.rs +++ b/crates/workspace-server/src/workspace_deletion.rs @@ -1198,7 +1198,7 @@ mod tests { ); assert!( - store + !store .delete_worker_registry( &workspace_id, &worker_runtime::identity::RuntimeWorkerRef { @@ -1207,6 +1207,8 @@ mod tests { }, ) .expect("remove worker registry") + .changes + .is_empty() ); let state: String = store .with_conn(|conn| {