Hare/yoi
1
1
Fork 0
Code Actions Packages Releases Activity

fix(workspace-server): make worker removal catalog commits atomic

Browse Source
This commit is contained in:
Hare
2026-09-19 01:34:19 +09:00
parent 308ae19f5e
commit b272bc37e8
5 changed files with 778 additions and 318 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/workspace-server/src/retention.rs
+132 -5
View File
@@ -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<WorkerRemovalPlan, WorkerRetentionError> {
) -> Result<WorkerRemovalCommit, WorkerRetentionError> {
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<WorkerRemovalPlanState> {
#[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()
crates/workspace-server/src/server.rs
+399 -153
View File
@@ -659,7 +659,7 @@ impl From<&WorkspaceApi> for ServerAuthApi {
}
#[derive(Clone)]
struct WorkspaceWorkerRemoveExecutor {
struct WorkerRemovalService {
workspace_id: String,
store: Arc<dyn ControlPlaneStore>,
runtime: Weak<RuntimeRegistry>,
@@ -667,10 +667,10 @@ struct WorkspaceWorkerRemoveExecutor {
workdir_session_locks: Arc<Mutex<HashMap<RuntimeWorkerRef, Arc<tokio::sync::Mutex<()>>>>>,
worker_remove_locks: Arc<Mutex<HashMap<RuntimeWorkerRef, Arc<tokio::sync::Mutex<()>>>>>,
worker_control_locks: Arc<Mutex<HashMap<String, Arc<tokio::sync::Mutex<()>>>>>,
worker_projection: std::sync::Weak<crate::worker_projection::WorkerProjectionService>,
worker_projection: Arc<crate::worker_projection::WorkerProjectionService>,
}
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<RuntimeDiagnostic>) {
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<RuntimeDiagnostic> {
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<RuntimeDiagnostic>) {
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,
))
crates/workspace-server/src/store.rs
+227 -105
View File
@@ -563,10 +563,27 @@ pub struct WorkerRegistryProjectionRecord {
pub observation: Option<WorkerRegistryObservationRecord>,
}
#[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<RuntimeWorkerRef>,
pub changes: Vec<WorkerCatalogChange>,
}
#[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<WorkerRegistryProjectionCommit> {
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<WorkerRemovalCommit, crate::retention::WorkerRetentionError> {
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<Option<DeviceLoginFlowRecord>>;
fn upsert_worker_registry(&self, record: &WorkerRegistryRecord) -> Result<()>;
fn upsert_worker_registry(
&self,
record: &WorkerRegistryRecord,
) -> Result<WorkerRegistryProjectionCommit>;
fn get_worker_registry(
&self,
workspace_id: &str,
@@ -1369,12 +1396,7 @@ pub trait ControlPlaneStore: Send + Sync + WorkspaceDeletionStore {
subject_revision: Option<u64>,
observed_at: &str,
) -> Result<WorkerRegistryProjectionCommit>;
fn publish_worker_registry_catalog_change(
&self,
workspace_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<WorkerRegistryProjectionCommit>;
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<bool>;
fn delete_worker_registry(&self, workspace_id: &str, worker: &RuntimeWorkerRef)
-> Result<bool>;
) -> Result<WorkerRegistryProjectionCommit>;
fn delete_worker_registry(
&self,
workspace_id: &str,
worker: &RuntimeWorkerRef,
) -> Result<WorkerRegistryProjectionCommit>;
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<WorkerRegistryProjectionCommit> {
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<WorkerRemovalCommit, crate::retention::WorkerRetentionError> {
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<WorkerRegistryProjectionCommit> {
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<WorkerRegistryProjectionCommit> {
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<bool> {
self.with_conn(|conn| {
let changed = conn.execute(
) -> Result<WorkerRegistryProjectionCommit> {
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<bool> {
self.with_conn(|conn| {
let tx = conn.unchecked_transaction()?;
) -> Result<WorkerRegistryProjectionCommit> {
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<Vec<WorkerCatalogChange>> {
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<WorkerRegistryProjectionCommit> {
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
crates/workspace-server/src/worker_projection.rs
+17 -54
View File
@@ -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::<Vec<_>>();
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(())
}
}
crates/workspace-server/src/workspace_deletion.rs
+3 -1
View File
@@ -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| {