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

fix: reconcile worker restores by operation

Browse Source
This commit is contained in:
Hare
2026-09-26 22:35:43 +09:00
parent a9c7f79012
commit 343af9ea3c
3 changed files with 274 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/worker-runtime/src/execution.rs
+14
View File
@@ -388,6 +388,13 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
Ok(())
}
fn reconcile_restore(
&self,
request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
self.restore_worker(request)
}
fn restore_worker(
&self,
_request: WorkerExecutionRestoreRequest,
@@ -608,6 +615,13 @@ impl WorkerExecutionBackendRef {
self.backend.preflight_restore(request)
}
pub(crate) fn reconcile_restore(
&self,
request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
self.backend.reconcile_restore(request)
}
pub(crate) fn restore_worker(
&self,
request: WorkerExecutionRestoreRequest,
crates/worker-runtime/src/runtime.rs
+100 -50
View File
@@ -1668,19 +1668,25 @@ impl Runtime {
worker_ref: &WorkerRef,
mode: WorkerRestoreMode,
) -> Result<RuntimeWorkerRestoreResult, RuntimeError> {
let operation_id = WorkerLifecycleOperationId::new();
let (backend, request) = {
let fresh_operation_id = WorkerLifecycleOperationId::new();
let (backend, request, operation_id, operation_mode, reconcile) = {
let state = self.lock()?;
state.ensure_running()?;
let worker = state.worker(worker_ref)?;
if worker.has_pending_lifecycle_operation() {
if worker.pending_stop.is_some() {
return Ok(RuntimeWorkerRestoreResult::failed(
WorkerRestoreState::ReconciliationRequired,
"worker_restore_reconciliation_pending",
"Worker restore reconciliation is still pending",
"worker_stop_reconciliation_pending",
"Worker stop reconciliation is still pending",
));
}
if worker.execution_handle.is_some() {
let pending_restore = worker.pending_restore;
let reconcile = pending_restore.is_some();
let operation_id = pending_restore
.map(|pending| pending.operation_id)
.unwrap_or(fresh_operation_id);
let operation_mode = pending_restore.map(|pending| pending.mode).unwrap_or(mode);
if !reconcile && worker.execution_handle.is_some() {
if worker.status.is_active() {
return Ok(RuntimeWorkerRestoreResult::accepted(worker.detail()));
}
@@ -1695,7 +1701,8 @@ impl Runtime {
// before Workspace Backend renewed operation-scoped Repository
// access). Preserve its automatic intent so an explicit restore can
// serve as the authorized retry for that state.
if mode == WorkerRestoreMode::Explicit
if !reconcile
&& mode == WorkerRestoreMode::Explicit
&& worker.execution_metadata_available
&& worker.status != WorkerStatus::Stopped
&& worker.restore_intent != WorkerRestoreIntent::Automatic
@@ -1706,7 +1713,8 @@ impl Runtime {
"Worker restore requires a stopped Worker",
));
}
if mode == WorkerRestoreMode::Automatic
if !reconcile
&& mode == WorkerRestoreMode::Automatic
&& (!worker.status.is_active()
|| worker.restore_intent != WorkerRestoreIntent::Automatic)
{
@@ -1743,10 +1751,10 @@ impl Runtime {
workdir_attachments: BTreeMap::new(),
config_bundle: None,
};
(backend, request)
(backend, request, operation_id, operation_mode, reconcile)
};
if let Err(_result) = backend.preflight_restore(&request) {
if !reconcile && let Err(_result) = backend.preflight_restore(&request) {
return Ok(RuntimeWorkerRestoreResult::failed(
WorkerRestoreState::Rejected,
"worker_restore_preflight_rejected",
@@ -1754,7 +1762,10 @@ impl Runtime {
));
}
if let Err(error) = self.begin_restore_operation(worker_ref, operation_id, mode) {
if !reconcile
&& let Err(error) =
self.begin_restore_operation(worker_ref, operation_id, operation_mode)
{
if matches!(error, RuntimeError::StoreCommitOutcomeUnknown { .. }) {
return Ok(RuntimeWorkerRestoreResult::failed(
WorkerRestoreState::ReconciliationRequired,
@@ -1765,7 +1776,12 @@ impl Runtime {
return Err(error);
}
match backend.restore_worker(request) {
let restore_result = if reconcile {
backend.reconcile_restore(request)
} else {
backend.restore_worker(request)
};
match restore_result {
WorkerExecutionSpawnResult::Connected {
handle,
worker_state,
@@ -1860,6 +1876,13 @@ impl Runtime {
}
}
WorkerExecutionSpawnResult::Rejected(_result) => {
if reconcile {
return Ok(RuntimeWorkerRestoreResult::failed(
WorkerRestoreState::ReconciliationRequired,
"worker_restore_reconciliation_rejected",
"Worker restore reconciliation was not completed",
));
}
if let Err(error) = self.settle_restore_rollback(worker_ref) {
tracing::error!(
worker_id = %worker_ref.worker_id,
@@ -2922,6 +2945,32 @@ impl Runtime {
}
}
let pending_restores = {
let state = self.lock()?;
state
.workers
.values()
.filter_map(|worker| {
worker
.pending_restore
.map(|pending| (worker.worker_ref.clone(), pending.mode))
})
.collect::<Vec<_>>()
};
for (worker_ref, mode) in pending_restores {
let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| RuntimeError::StatePoisoned)?;
if let Err(error) = self.restore_worker_under_lock(&worker_ref, mode) {
tracing::warn!(
worker_id = %worker_ref.worker_id,
error = %error,
"Worker restore reconciliation remains pending after Runtime startup"
);
}
}
let candidates = {
let state = self.lock()?;
if state.execution_backend.is_none() {
@@ -8806,18 +8855,20 @@ mod tests {
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);
assert_eq!(
backend.restore_operation_ids.lock().unwrap().as_slice(),
&[operation_id, operation_id]
);
assert!(
runtime
.lock()
.unwrap()
.worker(&created.worker_ref)
.unwrap()
.pending_restore
.unwrap()
.operation_id,
operation_id
.is_none()
);
}
@@ -9226,24 +9277,28 @@ mod tests {
let retry = restarted
.restore_worker_operation(&worker.worker_ref)
.unwrap();
assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 0);
assert_eq!(retry.state, WorkerRestoreState::Accepted);
assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 1);
assert_eq!(
restarted_backend
.restore_operation_ids
.lock()
.unwrap()
.as_slice(),
&[operation_id]
);
assert!(
restarted
.lock()
.unwrap()
.worker(&worker.worker_ref)
.unwrap()
.pending_restore
.unwrap()
.operation_id,
operation_id
.is_none()
);
let delete_error = restarted.delete_worker(&worker.worker_ref).unwrap_err();
assert!(
delete_error
.to_string()
.contains("pending restore reconciliation")
assert_eq!(
restarted.worker_detail(&worker.worker_ref).unwrap().status,
WorkerStatus::Idle
);
let _ = std::fs::remove_dir_all(root);
@@ -9298,37 +9353,32 @@ mod tests {
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);
drop(restored);
let restarted_backend = Arc::new(TestExecutionBackend::default());
let restarted = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: "test-runtime".to_string(),
display_name: None,
},
restarted_backend.clone(),
)
.unwrap();
assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 0);
assert_eq!(retry.state, WorkerRestoreState::Accepted);
assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 2);
assert_eq!(
restarted
restoring_backend
.restore_operation_ids
.lock()
.unwrap()
.as_slice(),
&[operation_id, operation_id]
);
assert!(
restored
.lock()
.unwrap()
.worker(&worker.worker_ref)
.unwrap()
.pending_restore
.unwrap()
.operation_id,
operation_id
.is_none()
);
let retry = restarted
.restore_worker_operation(&worker.worker_ref)
restored
.send_input(
&worker.worker_ref,
WorkerInput::user("after restore reconciliation"),
)
.unwrap();
assert_eq!(retry.state, WorkerRestoreState::ReconciliationRequired);
assert_eq!(*restarted_backend.restore_count.lock().unwrap(), 0);
restored.stop_worker(&worker.worker_ref, None).unwrap();
let _ = std::fs::remove_dir_all(root);
}
crates/worker-runtime/src/worker_backend.rs
+160 -20
View File
@@ -472,6 +472,17 @@ impl ProfileRuntimeWorkerFactory {
.join(uuid::Uuid::now_v7().to_string()))
}
fn worker_restore_run_dir(
&self,
worker_ref: &WorkerRef,
operation_id: crate::execution::WorkerLifecycleOperationId,
) -> Result<PathBuf, String> {
Ok(self
.worker_aggregate_dir(worker_ref)?
.join("runs")
.join(format!("restore-{operation_id}")))
}
fn runtime_worker_name_for_ref(worker_ref: &crate::identity::WorkerRef) -> String {
format!("worker-runtime-{}", worker_ref.worker_id)
}
@@ -1423,7 +1434,8 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
}
let workspace_client = worker.workspace_client_handle();
let run_dir = self.worker_run_dir(&request.worker_ref)?;
let run_dir =
self.worker_restore_run_dir(&request.worker_ref, request.operation_id)?;
let bash_output_dir = bash_output_dir_for_worker_id(&request.worker_ref.worker_id);
let started = PreparedWorker::new(
worker,
@@ -1567,6 +1579,8 @@ struct RuntimeWorkerExecution {
tasks: RuntimeExecutionTaskScope,
worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
restore_operation_id: Option<crate::execution::WorkerLifecycleOperationId>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
}
/// `worker-runtime` execution backend backed by real `worker` crate Workers.
@@ -1836,6 +1850,8 @@ where
tasks: RuntimeExecutionTaskScope,
worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
restore_operation_id: Option<crate::execution::WorkerLifecycleOperationId>,
workdir_attachments: Vec<crate::catalog::WorkingDirectoryAttachmentStatus>,
) -> Result<WorkerExecutionHandle, String> {
let mut workers = self
.workers
@@ -1853,6 +1869,8 @@ where
tasks,
worker_state,
workspace_client,
restore_operation_id,
workdir_attachments,
},
);
Ok(WorkerExecutionHandle::new(
@@ -1907,6 +1925,7 @@ where
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
controller_task: tokio::task::JoinHandle<()>,
restore_operation_id: Option<crate::execution::WorkerLifecycleOperationId>,
workdir_attachments: BTreeMap<WorkdirAttachmentAlias, WorkingDirectoryBinding>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
) -> WorkerExecutionSpawnResult {
@@ -1921,6 +1940,15 @@ where
let worker_state_snapshot = handle.shared_state.snapshot();
let worker_state = Arc::new(RwLock::new(worker_state_snapshot));
let tasks = RuntimeExecutionTaskScope::new(controller_task);
let workdir_attachment_statuses = workdir_attachments
.iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
},
)
.collect::<Vec<_>>();
#[cfg(feature = "ws-server")]
{
let mut events = streams.events;
@@ -1987,13 +2015,8 @@ where
Err(_) => {
let retained_state =
worker_state.read().ok().map(|state| state.clone());
let retained_workdir_attachments = workdir_attachments
.iter()
.map(|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias: alias.clone(),
working_directory: binding.status(),
})
.collect();
let retained_workdir_attachments =
workdir_attachment_statuses.clone();
let retained_handle = self
.retain_uncertain_unconnected_controller(
&worker_ref,
@@ -2002,6 +2025,8 @@ where
tasks,
worker_state,
workspace_client,
restore_operation_id,
retained_workdir_attachments.clone(),
)
.ok();
WorkerExecutionSpawnResult::ReconciliationRequired {
@@ -2072,21 +2097,15 @@ where
tasks,
worker_state,
workspace_client,
restore_operation_id,
workdir_attachments: workdir_attachment_statuses.clone(),
},
);
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()),
worker_state: connected_worker_state,
workdir_attachments: workdir_attachments
.into_iter()
.map(
|(alias, binding)| crate::catalog::WorkingDirectoryAttachmentStatus {
alias,
working_directory: binding.status(),
},
)
.collect(),
workdir_attachments: workdir_attachment_statuses,
}
}
}
@@ -2338,6 +2357,7 @@ where
controller.handle,
controller.shutdown,
controller.controller_task,
None,
BTreeMap::new(),
Some(controller.workspace_client),
);
@@ -2465,6 +2485,7 @@ where
controller.handle,
controller.shutdown,
controller.controller_task,
None,
workdir_attachments,
Some(controller.workspace_client),
)
@@ -2577,10 +2598,25 @@ where
.workers
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
if workers.contains_key(&request.worker_ref) {
if let Some(existing) = workers.get(&request.worker_ref) {
if existing.restore_operation_id == Some(request.operation_id) {
let worker_state = existing
.worker_state
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.clone();
return WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(
request.worker_ref.clone(),
self.backend_id(),
),
worker_state,
workdir_attachments: existing.workdir_attachments.clone(),
};
}
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy(
WorkerExecutionOperation::Restore,
"Worker is already connected to execution backend",
"Worker is already connected to execution backend by another lifecycle operation",
));
}
drop(workers);
@@ -2671,6 +2707,7 @@ where
let factory = self.factory.clone();
let bridge_context = request.context.clone();
let worker_ref = request.worker_ref.clone();
let operation_id = request.operation_id;
let restore_result = self
.run_cancellable_on_adapter_runtime(self.spawn_restore_timeout, async move {
factory.restore_controller(request).await
@@ -2698,6 +2735,7 @@ where
controller.handle,
controller.shutdown,
controller.controller_task,
Some(operation_id),
workdir_attachments,
Some(controller.workspace_client),
)
@@ -3555,7 +3593,13 @@ mod tests {
fn test_execution_context(worker_ref: WorkerRef) -> WorkerExecutionContext {
WorkerExecutionContext::new(
worker_ref,
Arc::new(|_, _| panic!("unused test event sink")),
Arc::new(|worker_ref, payload| {
Ok(crate::observation::WorkerObservationEvent::new(
1,
worker_ref,
payload,
))
}),
)
}
@@ -4562,6 +4606,102 @@ mod tests {
drop(restored_backend);
}
#[test]
#[serial_test::serial(worker_allocation)]
fn restore_reconciliation_reuses_the_operation_controller() {
let client = MockClient::sequential(Vec::new());
let runtime_base = tempfile::tempdir().unwrap();
let cwd = tempfile::tempdir().unwrap();
let store = tempfile::tempdir().unwrap();
let observed_workspace_clients = Arc::new(Mutex::new(Vec::new()));
let factory = MockFactory {
client,
runtime_base: runtime_base.path().to_path_buf(),
cwd: cwd.path().to_path_buf(),
store_dir: store.path().join("sessions"),
worker_metadata_dir: store.path().join("workers"),
observed_cwds: Arc::new(Mutex::new(Vec::new())),
observed_workspace_clients: observed_workspace_clients.clone(),
};
let backend = Arc::new(WorkerRuntimeExecutionBackend::new(factory).unwrap());
let runtime = EmbeddedRuntime::with_execution_backend(
RuntimeOptions::default(),
backend.clone(),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let request = create_request("idempotent restore reconciliation");
let worker = runtime.create_worker(request.clone()).unwrap();
runtime.stop_worker(&worker.worker_ref, None).unwrap();
let operation_id = crate::execution::WorkerLifecycleOperationId::new();
let restore_request = WorkerExecutionRestoreRequest {
operation_id,
worker_ref: worker.worker_ref.clone(),
request,
workspace_scope: None,
context: test_execution_context(worker.worker_ref.clone()),
previous_workdir_attachments: Vec::new(),
logical_workdir_attachments: Vec::new(),
workdir_attachments: BTreeMap::new(),
config_bundle: None,
};
let first = backend.restore_worker(restore_request.clone());
let first_handle = match first {
WorkerExecutionSpawnResult::Connected { handle, .. } => handle,
other => panic!("initial restore was not connected: {other:?}"),
};
let factory_calls_after_restore = observed_workspace_clients.lock().unwrap().len();
let reconciled = backend.reconcile_restore(restore_request);
assert!(matches!(
reconciled,
WorkerExecutionSpawnResult::Connected { .. }
));
assert_eq!(
observed_workspace_clients.lock().unwrap().len(),
factory_calls_after_restore,
"reconciliation must not construct a second controller"
);
let workers = backend.workers.lock().unwrap();
let execution = workers.get(&worker.worker_ref).unwrap();
assert_eq!(execution.restore_operation_id, Some(operation_id));
assert_eq!(workers.len(), 1);
drop(workers);
assert!(backend.stop_worker(&first_handle).is_accepted());
}
#[test]
fn profile_restore_run_directory_is_operation_scoped() {
let root = tempfile::tempdir().unwrap();
let runtime_store_dir = root.path().join("runtime");
let factory = ProfileRuntimeWorkerFactory::new(root.path())
.with_runtime_store_dir(&runtime_store_dir);
let worker_ref = WorkerRef::new(WorkerId::now_v7());
let operation_id = crate::execution::WorkerLifecycleOperationId::new();
let first = factory
.worker_restore_run_dir(&worker_ref, operation_id)
.unwrap();
let retry = factory
.worker_restore_run_dir(&worker_ref, operation_id)
.unwrap();
let different = factory
.worker_restore_run_dir(
&worker_ref,
crate::execution::WorkerLifecycleOperationId::new(),
)
.unwrap();
assert_eq!(first, retry);
assert_ne!(first, different);
assert_eq!(
first.file_name().unwrap().to_string_lossy(),
format!("restore-{operation_id}")
);
}
#[test]
fn builtin_profile_selector_is_not_double_prefixed() {
assert_eq!(