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

Harden restore preflight and retry reconciliation

Browse Source
This commit is contained in:
Hare
2026-09-18 00:48:56 +09:00
parent 68c59afbe7
commit 0d86ae66ad
3 changed files with 142 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/runtime.rs
+9 -18
View File
@@ -1483,13 +1483,6 @@ impl Runtime {
let state = self.lock()?;
state.ensure_running()?;
let worker = state.worker(worker_ref)?;
if worker.execution_bound && !worker.execution_metadata_available {
return Ok(RuntimeWorkerRestoreResult::failed(
WorkerRestoreState::ReconciliationRequired,
"worker_restore_reconciliation_pending",
"Worker restore cleanup or reconciliation is still pending",
));
}
if worker.execution_handle.is_some() {
if worker.status.is_active() {
return Ok(RuntimeWorkerRestoreResult::accepted(worker.detail()));
@@ -7520,11 +7513,12 @@ mod tests {
runtime.worker_detail(&created.worker_ref).unwrap().status,
WorkerStatus::Idle
);
backend.restore_result.lock().unwrap().take();
let retry = runtime
.restore_worker_operation(&created.worker_ref)
.unwrap();
assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(retry.state, WorkerRestoreState::Accepted);
assert_eq!(*backend.restore_count.lock().unwrap(), 2);
}
#[cfg(feature = "fs-store")]
@@ -7898,21 +7892,18 @@ mod tests {
let restored_worker = restored.worker_detail(&worker.worker_ref).unwrap();
assert_eq!(restored_worker.status, WorkerStatus::Idle);
assert!(!restored_worker.execution_metadata_available);
restoring_backend.restore_result.lock().unwrap().take();
let retry = restored
.restore_worker_operation(&worker.worker_ref)
.unwrap();
assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 1);
let err = restored
assert_eq!(retry.state, WorkerRestoreState::Accepted);
assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 2);
restored
.send_input(
&worker.worker_ref,
WorkerInput::user("after failed restore"),
WorkerInput::user("after reconciled restore"),
)
.unwrap_err();
assert!(matches!(
err,
RuntimeError::WorkerExecutionUnavailable { .. }
));
.unwrap();
let _ = std::fs::remove_dir_all(root);
}
crates/worker-runtime/src/worker_backend.rs
+128 -36
View File
@@ -59,7 +59,9 @@ fn next_internal_command(
.max(floor);
Ok(WorkerCommandEnvelope::new(command_id))
}
use session_store::{CombinedStore, WorkerAggregateStore, WorkerSessionStore};
use session_store::{
CombinedStore, WorkerAggregateStore, WorkerMetadataStore, WorkerSessionStore,
};
#[cfg(test)]
use session_store::{FsStore, FsWorkerStore};
use tokio::runtime::Runtime;
@@ -122,6 +124,13 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static {
request: WorkerExecutionSpawnRequest,
) -> Result<RuntimeWorkerController, String>;
async fn preflight_restore(
&self,
_request: &WorkerExecutionRestoreRequest,
) -> Result<(), String> {
Ok(())
}
async fn restore_controller(
&self,
request: WorkerExecutionRestoreRequest,
@@ -1002,6 +1011,41 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
})
}
async fn preflight_restore(
&self,
request: &WorkerExecutionRestoreRequest,
) -> Result<(), String> {
let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref);
let (mut manifest, _) = Self::restore_fallback_manifest(&worker_name)?;
bind_workspace_memory_settings(&mut manifest, &request.request)?;
let worker_aggregate_dir = self.worker_aggregate_dir(&request.worker_ref)?;
if !worker_aggregate_dir.is_dir() {
return Err("Persisted Worker aggregate metadata is unavailable".to_string());
}
if !worker_aggregate_dir.join("session").is_dir() {
return Err("Persisted Worker Session metadata is unavailable".to_string());
}
let metadata_store = WorkerAggregateStore::new(worker_aggregate_dir, &worker_name)
.map_err(|error| format!("failed to read Worker aggregate metadata: {error}"))?;
let metadata = metadata_store
.read_by_name(&worker_name)
.map_err(|error| format!("failed to read Worker metadata: {error}"))?
.ok_or_else(|| "Persisted Worker metadata is unavailable".to_string())?;
if request.request.workspace_api.is_some()
&& metadata
.active
.as_ref()
.is_some_and(|active| active.segment_id.is_none())
&& request.config_bundle.is_none()
{
return Err(
"Pending Workspace Worker restore requires current profile launch authority"
.to_string(),
);
}
Ok(())
}
async fn restore_controller(
&self,
request: WorkerExecutionRestoreRequest,
@@ -1310,6 +1354,7 @@ pub struct WorkerRuntimeExecutionBackend<F = ProfileRuntimeWorkerFactory> {
working_directory_materializer: Option<Arc<dyn WorkingDirectoryMaterializer>>,
runtime: Mutex<Option<Runtime>>,
workers: Mutex<HashMap<crate::identity::WorkerRef, RuntimeWorkerExecution>>,
restore_lock: Mutex<()>,
spawn_restore_timeout: Duration,
}
@@ -1338,6 +1383,7 @@ where
working_directory_materializer: None,
runtime: Mutex::new(Some(runtime)),
workers: Mutex::new(HashMap::new()),
restore_lock: Mutex::new(()),
spawn_restore_timeout: SPAWN_RESTORE_TASK_TIMEOUT,
})
}
@@ -1760,40 +1806,19 @@ where
let _ = bridge_context;
}
let mut workers = match self.workers.lock() {
Ok(workers) => workers,
Err(_) => {
let cleanup = self.cleanup_unconnected_controller(
&handle,
&shutdown,
&tasks,
&worker_state,
);
let result = WorkerExecutionResult::errored(
operation,
match &cleanup {
Ok(()) => "worker adapter registry lock is poisoned".to_string(),
Err(cleanup) => format!(
"worker adapter registry lock is poisoned; controller cleanup failed: {cleanup}"
),
},
);
return if operation == WorkerExecutionOperation::Restore {
match cleanup {
Ok(()) => WorkerExecutionSpawnResult::RolledBack(result),
Err(_) => WorkerExecutionSpawnResult::ReconciliationRequired {
result,
handle: None,
worker_state: None,
working_directory: None,
},
}
} else {
WorkerExecutionSpawnResult::Errored(result)
};
}
};
let mut workers = self
.workers
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if workers.contains_key(&worker_ref) {
let existing_worker_state = workers.get(&worker_ref).and_then(|worker| {
worker
.worker_state
.read()
.ok()
.map(|state| state.clone())
});
let existing_handle = WorkerExecutionHandle::new(worker_ref.clone(), self.backend_id());
drop(workers);
let cleanup = self.cleanup_unconnected_controller(
&handle,
@@ -1815,8 +1840,8 @@ where
Ok(()) => WorkerExecutionSpawnResult::RolledBack(result),
Err(_) => WorkerExecutionSpawnResult::ReconciliationRequired {
result,
handle: None,
worker_state: None,
handle: Some(existing_handle),
worker_state: existing_worker_state,
working_directory: None,
},
}
@@ -2125,6 +2150,54 @@ where
"Worker is already connected to execution backend",
));
}
drop(workers);
if request.previous_working_directory.is_some()
&& self.working_directory_materializer.is_none()
{
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir binding cannot be restored by this Runtime",
));
}
if request.previous_working_directory.is_none()
&& request.request.working_directory_request.is_some()
{
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir allocation is unavailable",
));
}
if request.previous_working_directory.is_none()
&& request.request.working_directory.is_some()
&& self.working_directory_materializer.is_none()
{
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Worker Workdir claim cannot be restored by this Runtime",
));
}
if let Some(workspace_api) = request.request.workspace_api.as_ref() {
let Some(scope) = request.workspace_scope.as_ref() else {
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Workspace restore authorization is unavailable",
));
};
if scope.workspace_id != workspace_api.workspace_id {
return Err(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"Persisted Workspace restore authorization does not match Worker authority",
));
}
}
let factory = Arc::clone(&self.factory);
let request = request.clone();
self.run_on_adapter_runtime(async move { factory.preflight_restore(&request).await })
.map_err(|message| {
WorkerExecutionResult::rejected(WorkerExecutionOperation::Restore, message)
})?;
Ok(())
}
@@ -2132,6 +2205,25 @@ where
&self,
mut request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
// Serialize backend restore starts so duplicate detection is completed
// before any controller or bridge side effect. Runtime also serializes
// per Worker; this guard protects direct/concurrent backend callers.
let _restore_guard = self
.restore_lock
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let workers = self
.workers
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if workers.contains_key(&request.worker_ref) {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy(
WorkerExecutionOperation::Restore,
"Worker is already connected to execution backend",
));
}
drop(workers);
let working_directory = match request.previous_working_directory.clone() {
Some(status) => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
crates/workspace-server/src/server.rs
+5 -3
View File
@@ -20722,7 +20722,7 @@ mod tests {
}
#[tokio::test]
async fn production_profile_backend_reports_uncertain_pending_orchestrator_restore() {
async fn production_profile_backend_preflight_rejects_pending_orchestrator_restore() {
let workspace = tempfile::tempdir().unwrap();
init_clean_git_workspace(workspace.path());
let config = test_server_config(workspace.path());
@@ -20772,9 +20772,11 @@ mod tests {
}),
)
.await
.expect_err("pending Workspace Orchestrator restore without durable Prompt must require reconciliation");
.expect_err(
"pending Workspace Orchestrator restore without durable Prompt must be rejected",
);
assert!(
format!("{error:?}").contains("worker_restore_reconciliation_required"),
format!("{error:?}").contains("worker_restore_preflight_rejected"),
"unexpected restore error: {error:?}"
);
}