runtime: restore persisted worker execution

This commit is contained in:
Keisuke Hirata 2026-07-12 23:38:56 +09:00
parent c56b8c04ea
commit c948b669c1
No known key found for this signature in database
6 changed files with 727 additions and 56 deletions

View File

@ -1,8 +1,8 @@
--- ---
title: 'Runtime restart-crossing Worker restore path' title: 'Runtime restart-crossing Worker restore path'
state: 'inprogress' state: 'done'
created_at: '2026-07-12T13:21:38Z' created_at: '2026-07-12T13:21:38Z'
updated_at: '2026-07-12T14:16:14Z' updated_at: '2026-07-12T14:38:56Z'
assignee: null assignee: null
queued_by: 'yoi ticket' queued_by: 'yoi ticket'
queued_at: '2026-07-12T14:16:14Z' queued_at: '2026-07-12T14:16:14Z'

View File

@ -39,4 +39,39 @@ Ticket を `yoi ticket` が queued にしました。
State changed to `inprogress`. State changed to `inprogress`.
---
<!-- event: implementation_report author: hare at: 2026-07-12T14:38:35Z -->
## Implementation report
Implemented Runtime restart-crossing Worker restore path.
- Added explicit `Restore` execution operation and `WorkerExecutionRestoreRequest` / backend restore API.
- Added Runtime restore pass for fs-backed Runtime startup after execution backend install.
- Persisted connected mappings remain stale on load until a real backend handle is restored.
- Restore success commits a live execution handle, connected status, Workdir status, and `WorkerExecutionRestored` event.
- Restore failure keeps the Worker stale and records `worker_execution_restore_failed` diagnostics without failing Runtime startup.
- Worker runtime backend can restore controllers through `RuntimeWorkerFactory::restore_controller`, using `Worker::restore_from_worker_metadata_with_context` and a conservative pending/no-history fallback.
- Spawn and restore share event bridge/handle registration behavior.
Validation:
- `cargo fmt`
- `cargo check -q -p worker-runtime`
- `cargo check -q -p worker-runtime --features fs-store`
- `cargo check -q -p worker-runtime --features fs-store,ws-server`
- `cargo test -q -p worker-runtime --features fs-store,ws-server`
- `cargo check -q`
- `nix build .#yoi --no-link`
---
<!-- event: state_changed author: "yoi ticket" at: 2026-07-12T14:38:56Z from: inprogress to: done reason: cli_state field: state -->
## State changed
State changed to `done`.
--- ---

View File

@ -61,6 +61,7 @@ pub enum WorkerExecutionRunState {
#[serde(rename_all = "snake_case")] #[serde(rename_all = "snake_case")]
pub enum WorkerExecutionOperation { pub enum WorkerExecutionOperation {
Spawn, Spawn,
Restore,
Input, Input,
Stop, Stop,
Cancel, Cancel,
@ -291,6 +292,20 @@ pub struct WorkerExecutionSpawnRequest {
pub config_bundle: Option<ConfigBundle>, pub config_bundle: Option<ConfigBundle>,
} }
/// Restore request passed to an execution backend for a persisted Runtime Worker.
///
/// The persisted execution status is a restore hint, not a live handle. Backends
/// must create a fresh controller/handle before returning `Connected`.
#[derive(Clone, Debug)]
pub struct WorkerExecutionRestoreRequest {
pub worker_ref: WorkerRef,
pub request: CreateWorkerRequest,
pub context: WorkerExecutionContext,
pub previous_execution: WorkerExecutionStatus,
pub working_directory: Option<WorkingDirectoryBinding>,
pub config_bundle: Option<ConfigBundle>,
}
/// Result of backend Worker spawn/initialization. /// Result of backend Worker spawn/initialization.
#[derive(Clone, Debug, PartialEq, Eq)] #[derive(Clone, Debug, PartialEq, Eq)]
pub enum WorkerExecutionSpawnResult { pub enum WorkerExecutionSpawnResult {
@ -313,6 +328,16 @@ pub trait WorkerExecutionBackend: Send + Sync + 'static {
fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult; fn spawn_worker(&self, request: WorkerExecutionSpawnRequest) -> WorkerExecutionSpawnResult;
fn restore_worker(
&self,
_request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::unsupported(
WorkerExecutionOperation::Restore,
"execution backend does not support restoring workers",
))
}
fn create_working_directory( fn create_working_directory(
&self, &self,
_request: &WorkingDirectoryRequest, _request: &WorkingDirectoryRequest,
@ -394,6 +419,19 @@ impl WorkerExecutionBackendRef {
self.backend.spawn_worker(request) self.backend.spawn_worker(request)
} }
#[cfg(feature = "fs-store")]
pub(crate) fn backend_id(&self) -> &str {
&self.id
}
#[cfg(feature = "fs-store")]
pub(crate) fn restore_worker(
&self,
request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
self.backend.restore_worker(request)
}
pub(crate) fn create_working_directory( pub(crate) fn create_working_directory(
&self, &self,
request: &WorkingDirectoryRequest, request: &WorkingDirectoryRequest,

View File

@ -39,6 +39,7 @@ pub enum RuntimeEventKind {
RuntimeStarted, RuntimeStarted,
RuntimeStopped, RuntimeStopped,
WorkerCreated, WorkerCreated,
WorkerExecutionRestored,
WorkerInputAccepted, WorkerInputAccepted,
WorkerStopped, WorkerStopped,
WorkerCancelled, WorkerCancelled,

View File

@ -11,6 +11,8 @@ use crate::config_bundle::{
use crate::diagnostics::DiagnosticSeverity; use crate::diagnostics::DiagnosticSeverity;
use crate::diagnostics::RuntimeDiagnostic; use crate::diagnostics::RuntimeDiagnostic;
use crate::error::RuntimeError; use crate::error::RuntimeError;
#[cfg(feature = "fs-store")]
use crate::execution::WorkerExecutionRestoreRequest;
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionBackendKind, WorkerExecutionBackendRef, WorkerExecutionBackend, WorkerExecutionBackendKind, WorkerExecutionBackendRef,
WorkerExecutionBindingIdentity, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionBindingIdentity, WorkerExecutionHandle, WorkerExecutionOperation,
@ -128,9 +130,11 @@ impl Runtime {
state state
}; };
state.execution_backend = execution_backend; state.execution_backend = execution_backend;
Ok(Self { let runtime = Self {
inner: Arc::new(Mutex::new(state)), inner: Arc::new(Mutex::new(state)),
}) };
runtime.restore_persisted_worker_executions()?;
Ok(runtime)
} }
/// Runtime id half of public Worker authority. /// Runtime id half of public Worker authority.
@ -639,7 +643,9 @@ impl Runtime {
let result = match operation { let result = match operation {
WorkerExecutionOperation::Stop => backend.stop_worker(&handle), WorkerExecutionOperation::Stop => backend.stop_worker(&handle),
WorkerExecutionOperation::Cancel => backend.cancel_worker(&handle), WorkerExecutionOperation::Cancel => backend.cancel_worker(&handle),
WorkerExecutionOperation::Spawn | WorkerExecutionOperation::Input => return Ok(()), WorkerExecutionOperation::Spawn
| WorkerExecutionOperation::Restore
| WorkerExecutionOperation::Input => return Ok(()),
}; };
if result.is_accepted() { if result.is_accepted() {
self.record_execution_result(worker_ref, result)?; self.record_execution_result(worker_ref, result)?;
@ -981,6 +987,136 @@ impl Runtime {
crate::execution::WorkerExecutionContext::new(worker_ref) crate::execution::WorkerExecutionContext::new(worker_ref)
} }
#[cfg(feature = "fs-store")]
fn restore_persisted_worker_executions(&self) -> Result<(), RuntimeError> {
#[derive(Clone)]
struct RestoreCandidate {
worker_ref: WorkerRef,
request: CreateWorkerRequest,
previous_execution: WorkerExecutionStatus,
}
let candidates = {
let mut state = self.lock()?;
let Some(backend) = state.execution_backend.clone() else {
return Ok(());
};
let backend_id = backend.backend_id().to_string();
let mut candidates = Vec::new();
let worker_ids: Vec<_> = state.workers.keys().cloned().collect();
for worker_id in worker_ids {
let Some(worker) = state.workers.get(&worker_id) else {
continue;
};
if !worker.status.is_active()
|| worker.execution_handle.is_some()
|| worker.execution.backend != WorkerExecutionBackendKind::Stale
|| worker
.execution
.binding
.as_ref()
.is_none_or(|binding| binding.backend_id != backend_id)
{
continue;
}
if let Some(working_directory) = worker.execution.working_directory.as_ref() {
if let Some(owner) = state.active_primary_worker_id_for_workdir_excluding(
&working_directory.summary.working_directory_id,
&worker.worker_id,
) {
let worker_ref = worker.worker_ref.clone();
let message = format!(
"worker {} cannot restore working directory {} because active worker {} is already the primary assignment",
worker.worker_id, working_directory.summary.working_directory_id, owner
);
state.record_restore_failure(
&worker_ref,
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
message,
),
)?;
continue;
}
}
candidates.push(RestoreCandidate {
worker_ref: worker.worker_ref.clone(),
request: worker.request.clone(),
previous_execution: worker.execution.clone(),
});
}
candidates
};
for candidate in candidates {
let backend = {
let state = self.lock()?;
state.execution_backend.clone()
};
let Some(backend) = backend else {
return Ok(());
};
let request = WorkerExecutionRestoreRequest {
worker_ref: candidate.worker_ref.clone(),
request: candidate.request,
context: self.execution_context(candidate.worker_ref.clone()),
previous_execution: candidate.previous_execution,
working_directory: None,
config_bundle: None,
};
match backend.restore_worker(request) {
WorkerExecutionSpawnResult::Connected {
handle,
run_state,
working_directory,
} => self.commit_restored_worker_execution(
&candidate.worker_ref,
handle,
run_state,
working_directory,
)?,
WorkerExecutionSpawnResult::Rejected(result)
| WorkerExecutionSpawnResult::Errored(result) => {
let mut state = self.lock()?;
state.record_restore_failure(&candidate.worker_ref, result)?;
}
}
}
Ok(())
}
#[cfg(feature = "fs-store")]
fn commit_restored_worker_execution(
&self,
worker_ref: &WorkerRef,
handle: WorkerExecutionHandle,
run_state: WorkerExecutionRunState,
working_directory: Option<CatalogWorkingDirectoryStatus>,
) -> Result<(), RuntimeError> {
let mut state = self.lock()?;
state.ensure_worker_ref(worker_ref)?;
let event_id = state.push_event(
Some(worker_ref.clone()),
RuntimeEventKind::WorkerExecutionRestored,
format!("worker {} execution restored", worker_ref.worker_id),
);
{
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = Some(handle.clone());
let mut execution = WorkerExecutionStatus::connected(run_state)
.with_binding(WorkerExecutionBindingIdentity::from_handle(&handle));
if let Some(status) = working_directory {
execution = execution.with_working_directory(status);
}
worker.execution = execution;
worker.last_event_id = event_id;
}
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
state.persist_event_by_id(event_id)?;
Ok(())
}
fn lock(&self) -> Result<MutexGuard<'_, RuntimeState>, RuntimeError> { fn lock(&self) -> Result<MutexGuard<'_, RuntimeState>, RuntimeError> {
self.inner.lock().map_err(|_| RuntimeError::StatePoisoned) self.inner.lock().map_err(|_| RuntimeError::StatePoisoned)
} }
@ -1394,6 +1530,62 @@ impl RuntimeState {
}) })
} }
#[cfg(feature = "fs-store")]
fn active_primary_worker_id_for_workdir_excluding(
&self,
working_directory_id: &str,
excluded_worker_id: &WorkerId,
) -> Option<WorkerId> {
self.workers.values().find_map(|worker| {
if worker.worker_id == *excluded_worker_id || !worker.status.is_active() {
return None;
}
if worker
.execution
.working_directory
.as_ref()
.is_some_and(|binding| binding.summary.working_directory_id == working_directory_id)
|| requested_primary_workdir_id(&worker.request) == Some(working_directory_id)
{
Some(worker.worker_id)
} else {
None
}
})
}
#[cfg(feature = "fs-store")]
fn record_restore_failure(
&mut self,
worker_ref: &WorkerRef,
result: WorkerExecutionResult,
) -> Result<(), RuntimeError> {
let message = result
.message
.clone()
.unwrap_or_else(|| "worker execution restore failed".to_string());
let diagnostic_id = self.next_diagnostic_id;
self.next_diagnostic_id += 1;
self.diagnostics.push(RuntimeDiagnostic {
id: diagnostic_id,
severity: DiagnosticSeverity::Warning,
code: "worker_execution_restore_failed".to_string(),
message: format!(
"worker {} execution restore failed: {message}",
worker_ref.worker_id
),
worker_ref: Some(worker_ref.clone()),
});
let worker = self.worker_mut(worker_ref)?;
worker.execution_handle = None;
let mut execution = WorkerExecutionStatus::stale(worker.execution.clone());
execution.last_result = Some(result);
worker.execution = execution;
self.persist_runtime_snapshot()?;
self.persist_worker(&worker_ref.worker_id)?;
Ok(())
}
fn push_event( fn push_event(
&mut self, &mut self,
worker_ref: Option<WorkerRef>, worker_ref: Option<WorkerRef>,
@ -1646,7 +1838,7 @@ mod tests {
}; };
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle, WorkerExecutionBackend, WorkerExecutionContext, WorkerExecutionHandle,
WorkerExecutionRunState, WorkerExecutionRestoreRequest, WorkerExecutionRunState,
}; };
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
@ -1714,6 +1906,8 @@ mod tests {
#[derive(Default)] #[derive(Default)]
struct TestExecutionBackend { struct TestExecutionBackend {
dispatch_result: Mutex<Option<WorkerExecutionResult>>, dispatch_result: Mutex<Option<WorkerExecutionResult>>,
restore_result: Mutex<Option<WorkerExecutionSpawnResult>>,
restore_count: Mutex<u64>,
contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>, contexts: Mutex<BTreeMap<WorkerId, WorkerExecutionContext>>,
} }
@ -1754,6 +1948,28 @@ mod tests {
} }
} }
fn restore_worker(
&self,
request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
*self.restore_count.lock().unwrap() += 1;
if let Some(result) = self.restore_result.lock().unwrap().clone() {
return result;
}
self.contexts
.lock()
.unwrap()
.insert(request.worker_ref.worker_id.clone(), request.context);
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(request.worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: request
.working_directory
.as_ref()
.map(|binding| binding.status()),
}
}
fn dispatch_input( fn dispatch_input(
&self, &self,
_handle: &WorkerExecutionHandle, _handle: &WorkerExecutionHandle,
@ -2386,6 +2602,134 @@ mod tests {
let _ = std::fs::remove_dir_all(root); let _ = std::fs::remove_dir_all(root);
} }
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_restores_active_worker_execution_handles() {
let root = fs_store_root("execution-restore");
let runtime_id = RuntimeId::new("runtime-execution-restore").unwrap();
let runtime = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: Some(runtime_id.clone()),
display_name: None,
limits: RuntimeLimits::default(),
},
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("restore active worker"))
.unwrap();
drop(runtime);
let restoring_backend = Arc::new(TestExecutionBackend::default());
let restored = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: Some(runtime_id),
display_name: None,
limits: RuntimeLimits::default(),
},
restoring_backend.clone(),
)
.unwrap();
assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 1);
let restored_worker = restored.worker_detail(&worker.worker_ref).unwrap();
assert_eq!(restored_worker.status, WorkerStatus::Running);
assert_eq!(
restored_worker.execution.backend,
WorkerExecutionBackendKind::Connected
);
assert!(restored_worker.execution.binding.is_some());
restored
.send_input(&worker.worker_ref, WorkerInput::user("after restart"))
.unwrap();
let cursor = restored.event_cursor_from_start().unwrap();
assert!(
restored
.read_events(&cursor, 16)
.unwrap()
.events
.iter()
.any(
|event| event.kind == RuntimeEventKind::WorkerExecutionRestored
&& event.worker_ref.as_ref() == Some(&worker.worker_ref)
)
);
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")]
#[test]
fn fs_store_keeps_worker_stale_when_execution_restore_fails() {
let root = fs_store_root("execution-restore-failed");
let runtime_id = RuntimeId::new("runtime-execution-restore-failed").unwrap();
let runtime = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: Some(runtime_id.clone()),
display_name: None,
limits: RuntimeLimits::default(),
},
Arc::new(TestExecutionBackend::default()),
)
.unwrap();
runtime.store_config_bundle(test_bundle()).unwrap();
let worker = runtime
.create_worker(task_request("restore failure"))
.unwrap();
drop(runtime);
let restoring_backend = Arc::new(TestExecutionBackend::default());
*restoring_backend.restore_result.lock().unwrap() =
Some(WorkerExecutionSpawnResult::Errored(
WorkerExecutionResult::errored(WorkerExecutionOperation::Restore, "restore boom"),
));
let restored = Runtime::with_fs_store_and_execution_backend(
crate::fs_store::FsRuntimeStoreOptions {
root: root.clone(),
runtime_id: Some(runtime_id),
display_name: None,
limits: RuntimeLimits::default(),
},
restoring_backend.clone(),
)
.unwrap();
assert_eq!(*restoring_backend.restore_count.lock().unwrap(), 1);
let restored_worker = restored.worker_detail(&worker.worker_ref).unwrap();
assert_eq!(restored_worker.status, WorkerStatus::Running);
assert_eq!(
restored_worker.execution.backend,
WorkerExecutionBackendKind::Stale
);
assert!(
restored
.diagnostics()
.unwrap()
.iter()
.any(
|diagnostic| diagnostic.code == "worker_execution_restore_failed"
&& diagnostic.worker_ref.as_ref() == Some(&worker.worker_ref)
)
);
let err = restored
.send_input(
&worker.worker_ref,
WorkerInput::user("after failed restore"),
)
.unwrap_err();
assert!(matches!(
err,
RuntimeError::WorkerExecutionUnavailable { .. }
));
let _ = std::fs::remove_dir_all(root);
}
#[cfg(feature = "fs-store")] #[cfg(feature = "fs-store")]
#[test] #[test]
fn fs_store_reports_corrupt_and_missing_data() { fn fs_store_reports_corrupt_and_missing_data() {

View File

@ -15,12 +15,13 @@ use std::sync::{Arc, Mutex, mpsc};
use std::time::Duration; use std::time::Duration;
use crate::catalog::{ use crate::catalog::{
ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource, WorkingDirectoryRequest, CreateWorkerRequest, ProfileSourceArchiveHttpRef, ProfileSourceArchiveSource,
WorkingDirectoryStatus, WorkingDirectoryRequest, WorkingDirectoryStatus,
}; };
use crate::execution::{ use crate::execution::{
WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation, WorkerExecutionResult, WorkerExecutionBackend, WorkerExecutionHandle, WorkerExecutionOperation,
WorkerExecutionRunState, WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult, WorkerExecutionRestoreRequest, WorkerExecutionResult, WorkerExecutionRunState,
WorkerExecutionSpawnRequest, WorkerExecutionSpawnResult,
}; };
use crate::interaction::{WorkerInput, WorkerInputKind}; use crate::interaction::{WorkerInput, WorkerInputKind};
use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache}; use crate::resource::{BackendResourceClient, ProfileSourceArchiveCache};
@ -33,11 +34,12 @@ use protocol::{Method, Segment, WorkerStatus};
use session_store::FsStore; use session_store::FsStore;
use session_store::{CombinedStore, FsWorkerStore}; use session_store::{CombinedStore, FsWorkerStore};
use tokio::runtime::Runtime; use tokio::runtime::Runtime;
#[cfg(feature = "ws-server")]
use tokio::sync::broadcast; use tokio::sync::broadcast;
use worker::{ use worker::{
Worker, WorkerController, WorkerFilesystemAuthority, WorkerHandle, WorkerWorkspaceContext, Worker, WorkerController, WorkerError, WorkerFilesystemAuthority, WorkerHandle,
WorkspaceId, WorkerWorkspaceContext, WorkspaceId,
}; };
const DEFAULT_BACKEND_ID: &str = "worker-crate"; const DEFAULT_BACKEND_ID: &str = "worker-crate";
@ -51,6 +53,11 @@ pub trait RuntimeWorkerFactory: Send + Sync + 'static {
&self, &self,
request: WorkerExecutionSpawnRequest, request: WorkerExecutionSpawnRequest,
) -> Result<WorkerHandle, String>; ) -> Result<WorkerHandle, String>;
async fn restore_controller(
&self,
request: WorkerExecutionRestoreRequest,
) -> Result<WorkerHandle, String>;
} }
/// Production factory that resolves a normal Worker profile and spawns it under /// Production factory that resolves a normal Worker profile and spawns it under
@ -139,14 +146,18 @@ impl ProfileRuntimeWorkerFactory {
.ok_or_else(|| "could not resolve worker runtime directory".to_string()) .ok_or_else(|| "could not resolve worker runtime directory".to_string())
} }
fn runtime_worker_name(request: &WorkerExecutionSpawnRequest) -> String { fn runtime_worker_name_for_ref(worker_ref: &crate::identity::WorkerRef) -> String {
format!( format!(
"runtime-{}-{}", "runtime-{}-{}",
sanitize_worker_name_component(request.worker_ref.runtime_id.as_str()), sanitize_worker_name_component(worker_ref.runtime_id.as_str()),
request.worker_ref.worker_id worker_ref.worker_id
) )
} }
fn runtime_worker_name(request: &WorkerExecutionSpawnRequest) -> String {
Self::runtime_worker_name_for_ref(&request.worker_ref)
}
fn runtime_profile_value( fn runtime_profile_value(
profile: &crate::catalog::ProfileSelector, profile: &crate::catalog::ProfileSelector,
) -> Option<std::borrow::Cow<'_, str>> { ) -> Option<std::borrow::Cow<'_, str>> {
@ -165,14 +176,21 @@ impl ProfileRuntimeWorkerFactory {
} }
} }
fn runtime_profile<'a>( fn runtime_profile_for_request<'a>(
&'a self, &'a self,
request: &'a WorkerExecutionSpawnRequest, request: &'a CreateWorkerRequest,
) -> Option<std::borrow::Cow<'a, str>> { ) -> Option<std::borrow::Cow<'a, str>> {
if let Some(profile) = self.profile.as_deref() { if let Some(profile) = self.profile.as_deref() {
return Some(std::borrow::Cow::Borrowed(profile)); return Some(std::borrow::Cow::Borrowed(profile));
} }
Self::runtime_profile_value(&request.request.profile) Self::runtime_profile_value(&request.profile)
}
fn runtime_profile<'a>(
&'a self,
request: &'a WorkerExecutionSpawnRequest,
) -> Option<std::borrow::Cow<'a, str>> {
self.runtime_profile_for_request(&request.request)
} }
async fn resolve_profile_source_archive( async fn resolve_profile_source_archive(
&self, &self,
@ -397,6 +415,110 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
.map_err(|err| format!("failed to spawn Worker controller: {err}"))?; .map_err(|err| format!("failed to spawn Worker controller: {err}"))?;
Ok(handle) Ok(handle)
} }
async fn restore_controller(
&self,
request: WorkerExecutionRestoreRequest,
) -> Result<WorkerHandle, String> {
let worker_name = Self::runtime_worker_name_for_ref(&request.worker_ref);
let profile = self.runtime_profile_for_request(&request.request);
let worker_root = request
.working_directory
.as_ref()
.map(|binding| binding.root().to_path_buf())
.unwrap_or_else(|| self.profile_base_dir.clone());
let filesystem_authority = request
.working_directory
.as_ref()
.map(|binding| {
WorkerFilesystemAuthority::local(
binding.root().to_path_buf(),
binding.cwd().to_path_buf(),
)
})
.unwrap_or(WorkerFilesystemAuthority::None);
let workspace_backend_ref =
RuntimeWorkspaceBackendRef::from_working_directory(request.working_directory.as_ref());
let workspace_context = workspace_backend_ref.worker_context();
let selector = profile.as_deref().unwrap_or("builtin:default");
let archive = self
.resolve_profile_source_archive(&request.request.profile_source)
.await?;
let (mut manifest, loader) = {
let manifest = archive
.resolve_profile(selector, &worker_root, &worker_name)
.map_err(|err| format!("failed to resolve profile source archive: {err}"))?;
worker::entrypoint::resolve_runtime_profile_manifest_from_manifest(
manifest,
&worker_root,
&worker_name,
)?
};
manifest.worker.name = worker_name.clone();
let store_dir = self.store_dir()?;
let session_store = FsStore::new(&store_dir).map_err(|err| {
format!(
"failed to initialize session store at {}: {err}",
store_dir.display()
)
})?;
let worker_metadata_dir = self.worker_metadata_dir(&store_dir);
let worker_metadata_store = FsWorkerStore::new(&worker_metadata_dir).map_err(|err| {
format!(
"failed to initialize worker metadata store at {}: {err}",
worker_metadata_dir.display()
)
})?;
let store = CombinedStore::new(session_store, worker_metadata_store);
let worker = match Worker::restore_from_worker_metadata_with_context(
&worker_name,
manifest.clone(),
store,
loader.clone(),
workspace_context.clone(),
filesystem_authority.clone(),
)
.await
{
Ok(worker) => worker,
Err(WorkerError::WorkerMetadataPending { .. })
if request.request.initial_input.is_none() =>
{
let session_store = FsStore::new(&store_dir).map_err(|err| {
format!(
"failed to initialize session store at {}: {err}",
store_dir.display()
)
})?;
let worker_metadata_store =
FsWorkerStore::new(&worker_metadata_dir).map_err(|err| {
format!(
"failed to initialize worker metadata store at {}: {err}",
worker_metadata_dir.display()
)
})?;
let store = CombinedStore::new(session_store, worker_metadata_store);
Worker::from_manifest_with_context(
manifest,
store,
loader,
workspace_context,
filesystem_authority,
)
.await
.map_err(|err| format!("failed to recreate pending Worker from profile: {err}"))?
}
Err(err) => return Err(format!("failed to restore Worker from metadata: {err}")),
};
let runtime_base = self.runtime_base_dir()?;
let (handle, _shutdown_rx) = WorkerController::spawn(worker, &runtime_base)
.await
.map_err(|err| format!("failed to spawn restored Worker controller: {err}"))?;
Ok(handle)
}
} }
struct RuntimeWorkerExecution { struct RuntimeWorkerExecution {
@ -532,6 +654,69 @@ where
.map(|_| WorkerExecutionResult::accepted(operation, accepted_run_state)) .map(|_| WorkerExecutionResult::accepted(operation, accepted_run_state))
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message)) .unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
} }
fn connect_handle(
&self,
operation: WorkerExecutionOperation,
worker_ref: crate::identity::WorkerRef,
bridge_context: crate::execution::WorkerExecutionContext,
handle: WorkerHandle,
working_directory: Option<WorkingDirectoryBinding>,
) -> WorkerExecutionSpawnResult {
let busy = Arc::new(AtomicBool::new(false));
#[cfg(feature = "ws-server")]
{
let mut events = handle.subscribe();
let bridge_handle = handle.clone();
let bridge_busy = busy.clone();
if let Err(message) = self.spawn_on_adapter_runtime(async move {
loop {
match events.recv().await {
Ok(event) => {
let _ = bridge_context.publish_protocol_event(event);
if bridge_handle.shared_state.get_status() == WorkerStatus::Idle {
bridge_busy.store(false, Ordering::SeqCst);
}
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => break,
}
}
}) {
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
operation, message,
));
}
}
#[cfg(not(feature = "ws-server"))]
{
let _ = bridge_context;
}
let mut workers = match self.workers.lock() {
Ok(workers) => workers,
Err(_) => {
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
operation,
"worker adapter registry lock is poisoned",
));
}
};
workers.insert(
worker_ref.clone(),
RuntimeWorkerExecution {
handle,
busy,
working_directory: working_directory.clone(),
},
);
WorkerExecutionSpawnResult::Connected {
handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()),
run_state: WorkerExecutionRunState::Idle,
working_directory: working_directory.map(|binding| binding.status()),
}
}
} }
impl<F> Drop for WorkerRuntimeExecutionBackend<F> { impl<F> Drop for WorkerRuntimeExecutionBackend<F> {
@ -694,53 +879,108 @@ where
} }
}; };
let mut events = handle.subscribe(); self.connect_handle(
let bridge_handle = handle.clone();
let busy = Arc::new(AtomicBool::new(false));
let bridge_busy = busy.clone();
if let Err(message) = self.spawn_on_adapter_runtime(async move {
loop {
match events.recv().await {
Ok(event) => {
let _ = bridge_context.publish_protocol_event(event);
if bridge_handle.shared_state.get_status() == WorkerStatus::Idle {
bridge_busy.store(false, Ordering::SeqCst);
}
}
Err(broadcast::error::RecvError::Lagged(_)) => continue,
Err(broadcast::error::RecvError::Closed) => break,
}
}
}) {
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
WorkerExecutionOperation::Spawn, WorkerExecutionOperation::Spawn,
worker_ref,
bridge_context,
handle,
working_directory,
)
}
fn restore_worker(
&self,
mut request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
let working_directory = match request.previous_execution.working_directory.clone() {
Some(status) => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted worker has a working directory binding, but no materializer is configured for this runtime backend",
));
};
let relative_cwd = request
.request
.working_directory
.as_ref()
.and_then(|working_directory| working_directory.relative_cwd.as_deref());
match materializer
.bind_working_directory(&status.summary.working_directory_id, relative_cwd)
{
Ok(binding) => {
request.working_directory = Some(binding.clone());
Some(binding)
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
error.to_string(),
),
);
}
}
}
None if request.request.working_directory_request.is_some() => {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted worker requested a working directory, but no persisted working directory binding is available to restore",
));
}
None if request.request.working_directory.is_some() => {
let Some(materializer) = self.working_directory_materializer.as_ref() else {
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
"persisted worker has a working directory claim, but no materializer is configured for this runtime backend",
));
};
let working_directory =
request.request.working_directory.as_ref().expect("checked");
match materializer.bind_working_directory(
&working_directory.working_directory_id,
working_directory.relative_cwd.as_deref(),
) {
Ok(binding) => {
request.working_directory = Some(binding.clone());
Some(binding)
}
Err(error) => {
return WorkerExecutionSpawnResult::Rejected(
WorkerExecutionResult::rejected(
WorkerExecutionOperation::Restore,
error.to_string(),
),
);
}
}
}
None => None,
};
let factory = self.factory.clone();
let bridge_context = request.context.clone();
let worker_ref = request.worker_ref.clone();
let restore_result =
self.run_on_adapter_runtime(async move { factory.restore_controller(request).await });
let handle = match restore_result {
Ok(handle) => handle,
Err(message) => {
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
WorkerExecutionOperation::Restore,
message, message,
)); ));
} }
let mut workers = match self.workers.lock() {
Ok(workers) => workers,
Err(_) => {
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
WorkerExecutionOperation::Spawn,
"worker adapter registry lock is poisoned",
));
}
}; };
workers.insert(
worker_ref.clone(),
RuntimeWorkerExecution {
handle,
busy,
working_directory: working_directory.clone(),
},
);
WorkerExecutionSpawnResult::Connected { self.connect_handle(
handle: WorkerExecutionHandle::new(worker_ref, self.backend_id()), WorkerExecutionOperation::Restore,
run_state: WorkerExecutionRunState::Idle, worker_ref,
working_directory: working_directory.map(|binding| binding.status()), bridge_context,
} handle,
working_directory,
)
} }
fn dispatch_input( fn dispatch_input(
@ -988,6 +1228,19 @@ mod tests {
.map_err(|err| err.to_string())?; .map_err(|err| err.to_string())?;
Ok(handle) Ok(handle)
} }
async fn restore_controller(
&self,
request: WorkerExecutionRestoreRequest,
) -> Result<WorkerHandle, String> {
let request = WorkerExecutionSpawnRequest {
worker_ref: request.worker_ref,
request: request.request,
context: request.context,
working_directory: request.working_directory,
config_bundle: request.config_bundle,
};
self.spawn_controller(request).await
}
} }
fn core_filesystem_tool_names() -> BTreeSet<&'static str> { fn core_filesystem_tool_names() -> BTreeSet<&'static str> {