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

fix: serialize runtime worker lifecycle

Browse Source
This commit is contained in:
Hare
2026-09-17 00:14:33 +09:00
parent 93070a3a93
commit 9d8e0114f5
12 changed files with 942 additions and 224 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -6070,6 +6070,7 @@ dependencies = [
"base64 0.22.1",
"decodal",
"flow",
"fs4",
"futures",
"manifest",
"protocol",
crates/worker-runtime/Cargo.toml
+2 -1
View File
@@ -18,7 +18,7 @@ required-features = ["ws-server", "fs-store"]
[features]
default = ["ws-server", "fs-store"]
fs-store = []
fs-store = ["dep:fs4"]
http-server = ["dep:axum", "dep:tower", "dep:reqwest"]
ws-server = ["http-server", "axum/ws", "dep:futures", "tokio/sync"]
@@ -29,6 +29,7 @@ axum = { workspace = true, optional = true }
futures = { workspace = true, optional = true }
decodal.workspace = true
flow = { path = "../flow" }
fs4 = { workspace = true, optional = true }
manifest.workspace = true
protocol.workspace = true
serde = { workspace = true, features = ["derive"] }
crates/worker-runtime/src/error.rs
+3
View File
@@ -79,6 +79,9 @@ pub enum RuntimeError {
name: String,
},
#[error("Runtime store is already owned by another process")]
RuntimeStoreAlreadyOpen { path: PathBuf },
#[error("runtime store {operation} failed at {}: {source}", path.display())]
StoreIo {
operation: &'static str,
crates/worker-runtime/src/fs_store.rs
+150 -2
View File
@@ -8,11 +8,13 @@ use crate::identity::{
LegacyWorkerIdentityMapping, WorkerId, WorkerRef, legacy_worker_identity_mapping_digest,
};
use crate::management::{RuntimeBackendKind, RuntimeStatus};
use fs4::fs_std::FileExt;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;
use std::fs::{self, File, OpenOptions};
use std::io::{BufReader, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
const SCHEMA_VERSION: u32 = 6;
@@ -50,25 +52,47 @@ impl FsRuntimeStoreOptions {
}
}
#[derive(Debug)]
struct RuntimeStoreOwnerLock {
file: File,
}
impl Drop for RuntimeStoreOwnerLock {
fn drop(&mut self) {
let _ = FileExt::unlock(&self.file);
}
}
/// Filesystem persistence boundary for one Worker Runtime state.
///
/// Authority is the Workspace-owned typed Worker identity. Legacy pod paths, socket
/// paths, and session paths are deliberately not part of the layout or lookup API.
#[derive(Clone, Debug, PartialEq, Eq)]
#[derive(Clone, Debug)]
pub struct FsRuntimeStore {
root: PathBuf,
_owner_lock: Option<Arc<RuntimeStoreOwnerLock>>,
}
impl PartialEq for FsRuntimeStore {
fn eq(&self, other: &Self) -> bool {
self.root == other.root
}
}
impl Eq for FsRuntimeStore {}
impl FsRuntimeStore {
pub fn migration_plan(
options: &FsRuntimeStoreOptions,
) -> Result<FsRuntimeStoreMigrationPlan, RuntimeError> {
let _owner_lock = acquire_runtime_store_owner_lock(&options.root)?;
plan_runtime_store_migration(&options.root, &options.runtime_id).map(|(plan, _)| plan)
}
pub fn migrate(
options: &FsRuntimeStoreOptions,
) -> Result<FsRuntimeStoreMigrationPlan, RuntimeError> {
let _owner_lock = acquire_runtime_store_owner_lock(&options.root)?;
migrate_runtime_store(&options.root, &options.runtime_id)
}
@@ -93,6 +117,15 @@ impl FsRuntimeStore {
});
}
if !existed {
fs::create_dir_all(&root).map_err(|source| RuntimeError::StoreIo {
operation: "create runtime store root",
path: root.clone(),
source,
})?;
}
let owner_lock = acquire_runtime_store_owner_lock(&root)?;
fs::create_dir_all(root.join(WORKERS_DIR)).map_err(|source| RuntimeError::StoreIo {
operation: "create runtime store",
path: root.join(WORKERS_DIR),
@@ -114,7 +147,10 @@ impl FsRuntimeStore {
if existed {
migrate_runtime_store(&root, runtime_id)?;
}
let store = Self { root };
let store = Self {
root,
_owner_lock: Some(owner_lock),
};
let state = if existed {
Some(store.load_runtime_state()?)
} else {
@@ -256,6 +292,61 @@ impl FsRuntimeStore {
}
}
fn acquire_runtime_store_owner_lock(
root: &Path,
) -> Result<Arc<RuntimeStoreOwnerLock>, RuntimeError> {
let canonical_root = fs::canonicalize(root).map_err(|source| RuntimeError::StoreIo {
operation: "resolve runtime store owner lock",
path: root.to_path_buf(),
source,
})?;
let parent = canonical_root
.parent()
.ok_or_else(|| RuntimeError::StoreCorrupt {
operation: "resolve runtime store owner lock",
path: root.to_path_buf(),
message: "runtime store root has no parent".to_string(),
})?;
let name = canonical_root
.file_name()
.ok_or_else(|| RuntimeError::StoreCorrupt {
operation: "resolve runtime store owner lock",
path: root.to_path_buf(),
message: "runtime store root has no file name".to_string(),
})?;
let mut lock_name = std::ffi::OsString::from(".");
lock_name.push(name);
lock_name.push(".runtime-owner.lock");
let lock_path = parent.join(lock_name);
let file = OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&lock_path)
.map_err(|source| RuntimeError::StoreIo {
operation: "open runtime store owner lock",
path: root.to_path_buf(),
source,
})?;
match file.try_lock_exclusive() {
Ok(true) => Ok(Arc::new(RuntimeStoreOwnerLock { file })),
Ok(false) => Err(RuntimeError::RuntimeStoreAlreadyOpen {
path: root.to_path_buf(),
}),
Err(source) if source.kind() == std::io::ErrorKind::WouldBlock => {
Err(RuntimeError::RuntimeStoreAlreadyOpen {
path: root.to_path_buf(),
})
}
Err(source) => Err(RuntimeError::StoreIo {
operation: "acquire runtime store owner lock",
path: root.to_path_buf(),
source,
}),
}
}
fn remove_legacy_observations(worker_dir: &Path) {
let path = worker_dir.join(LEGACY_OBSERVATIONS_FILE);
if path.is_file() {
@@ -1193,6 +1284,7 @@ fn migrate_runtime_store(
};
let staged_store = FsRuntimeStore {
root: staging.clone(),
_owner_lock: None,
};
if let Err(error) = staged_store.load_runtime_state() {
let _ = fs::remove_dir_all(&staging);
@@ -1628,6 +1720,62 @@ fn sync_directory(path: &Path, operation: &'static str) -> Result<(), RuntimeErr
mod tests {
use super::*;
#[test]
fn runtime_store_owner_lock_child_probe() {
let Some(root) = std::env::var_os("YOI_TEST_RUNTIME_STORE_LOCK_ROOT") else {
return;
};
assert!(matches!(
FsRuntimeStore::open_or_create(PathBuf::from(root), "runtime-test").unwrap_err(),
RuntimeError::RuntimeStoreAlreadyOpen { .. }
));
}
#[test]
fn second_runtime_store_open_conflicts_before_store_mutation() {
let parent = tempfile::tempdir().unwrap();
let root = parent.path().join("runtime-store");
let first = FsRuntimeStore::open_or_create(root.clone(), "runtime-test").unwrap();
let legacy_events = root.join("events.jsonl");
fs::write(&legacy_events, b"must remain").unwrap();
let error = FsRuntimeStore::open_or_create(root.clone(), "runtime-test").unwrap_err();
assert!(matches!(
error,
RuntimeError::RuntimeStoreAlreadyOpen { path } if path == root
));
assert_eq!(fs::read(&legacy_events).unwrap(), b"must remain");
let child = std::process::Command::new(std::env::current_exe().unwrap())
.arg("fs_store::tests::runtime_store_owner_lock_child_probe")
.arg("--exact")
.env("YOI_TEST_RUNTIME_STORE_LOCK_ROOT", &root)
.output()
.unwrap();
assert!(
child.status.success(),
"cross-process lock probe failed: {}",
String::from_utf8_lossy(&child.stderr)
);
assert!(fs::read_dir(&root).unwrap().all(|entry| {
!entry
.unwrap()
.file_name()
.to_string_lossy()
.contains("owner.lock")
}));
let retained_clone = first.store.clone();
drop(first);
assert!(matches!(
FsRuntimeStore::open_or_create(root.clone(), "runtime-test").unwrap_err(),
RuntimeError::RuntimeStoreAlreadyOpen { .. }
));
drop(retained_clone);
fs::remove_dir_all(&root).unwrap();
FsRuntimeStore::open_or_create(root, "runtime-test").unwrap();
}
#[test]
fn schema_v4_migration_plan_ignores_orphan_worker_directories() {
let root = tempfile::tempdir().unwrap();
crates/worker-runtime/src/http_server.rs
+2
View File
@@ -2467,6 +2467,7 @@ fn status_for_runtime_error(error: &RuntimeError) -> StatusCode {
StatusCode::NOT_FOUND
}
RuntimeError::RuntimeStopped
| RuntimeError::RuntimeStoreAlreadyOpen { .. }
| RuntimeError::WorkerExecutionUnavailable { .. }
| RuntimeError::ExecutionBackendUnavailable { .. }
| RuntimeError::WorkerExecutionRejected { .. } => StatusCode::CONFLICT,
@@ -2488,6 +2489,7 @@ fn status_for_runtime_error(error: &RuntimeError) -> StatusCode {
fn code_for_runtime_error(error: &RuntimeError) -> String {
match error {
RuntimeError::RuntimeStopped => "runtime_stopped".to_string(),
RuntimeError::RuntimeStoreAlreadyOpen { .. } => "runtime_store_already_open".to_string(),
RuntimeError::WorkerNotFound { .. } => "worker_not_found".to_string(),
RuntimeError::WorkerExecutionUnavailable { .. } => {
"worker_execution_unavailable".to_string()
crates/worker-runtime/src/runtime.rs
+353 -97
View File
@@ -72,6 +72,12 @@ impl RuntimeWorkspaceScope {
const SUBSCRIPTION_QUEUE_CAPACITY: usize = 256;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum WorkerRestoreMode {
Explicit,
Automatic,
}
#[derive(Clone, Debug)]
pub struct RuntimeSubscriptionUpdate {
pub subject_revision: u64,
@@ -1295,16 +1301,51 @@ impl Runtime {
/// Attach a live execution to a persisted Worker definition.
///
/// Current liveness is never read from disk. If a handle is already
/// present this is idempotent; otherwise the configured backend is tried.
/// Every lifecycle mutation for a Worker is serialized by the same operation
/// lock. A concurrent exact restore waits for the first caller and then
/// converges on its already-installed execution rather than spawning another
/// controller.
pub fn restore_worker(&self, worker_ref: &WorkerRef) -> Result<WorkerDetail, RuntimeError> {
let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| RuntimeError::StatePoisoned)?;
self.restore_worker_under_lock(worker_ref, WorkerRestoreMode::Explicit)
}
fn restore_worker_under_lock(
&self,
worker_ref: &WorkerRef,
mode: WorkerRestoreMode,
) -> Result<WorkerDetail, RuntimeError> {
let (backend, request) = {
let mut state = self.lock()?;
state.ensure_running()?;
let (worker_request, previous_working_directory, run_generation) = {
let worker = state.worker(worker_ref)?;
if worker.execution_handle.is_some() {
return Ok(worker.detail());
if worker.status.is_active() {
return Ok(worker.detail());
}
return Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id,
message: "previous execution cleanup has not completed".to_string(),
});
}
match mode {
WorkerRestoreMode::Explicit if worker.status != WorkerStatus::Stopped => {
return Err(RuntimeError::InvalidRequest(format!(
"worker {} is not stopped",
worker_ref.worker_id
)));
}
WorkerRestoreMode::Automatic
if !worker.status.is_active()
|| worker.restore_intent != WorkerRestoreIntent::Automatic =>
{
return Ok(worker.detail());
}
_ => {}
}
(
worker.request.clone(),
@@ -1314,7 +1355,7 @@ impl Runtime {
};
let backend = state.execution_backend.clone().ok_or_else(|| {
RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id.clone(),
worker_id: worker_ref.worker_id,
message: "runtime has no execution backend".to_string(),
}
})?;
@@ -1322,7 +1363,6 @@ impl Runtime {
let worker = state.worker_mut(worker_ref)?;
worker.run_generation = run_generation;
worker.execution_bound = true;
worker.restore_intent = WorkerRestoreIntent::Automatic;
}
state.persist_worker(&worker_ref.worker_id)?;
let workspace_scope = worker_request.workspace_api.as_ref().and_then(|api| {
@@ -1350,13 +1390,17 @@ impl Runtime {
worker_state,
working_directory,
} => {
self.commit_restored_worker_execution(
let commit = self.commit_restored_worker_execution(
worker_ref,
handle,
handle.clone(),
worker_state,
WorkerStatus::Idle,
working_directory,
)?;
);
if let Err(error) = commit {
self.cleanup_failed_restore(&backend, worker_ref, &handle)?;
return Err(error);
}
self.worker_detail(worker_ref)
}
WorkerExecutionSpawnResult::Rejected(result)
@@ -1367,7 +1411,7 @@ impl Runtime {
.record_restore_failure(worker_ref, result.clone())?;
}
Err(RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id.clone(),
worker_id: worker_ref.worker_id,
operation: result.operation,
outcome: result.outcome,
message: result.message_or_default(),
@@ -1894,15 +1938,64 @@ impl Runtime {
self.stop_worker(worker_ref, reason)
}
/// Stop a Worker. Repeated stops are idempotent.
/// Stop a Worker. Repeated stops are idempotent. The per-Worker lifecycle
/// lock remains held until backend cleanup and the terminal catalog commit
/// have both completed.
pub fn stop_worker(
&self,
worker_ref: &WorkerRef,
reason: Option<String>,
) -> Result<WorkerLifecycleAck, RuntimeError> {
self.dispatch_lifecycle_to_backend(worker_ref, WorkerExecutionOperation::Stop)?;
let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| RuntimeError::StatePoisoned)?;
let (backend, handle) = {
let state = self.lock()?;
state.ensure_running()?;
let worker = state.worker(worker_ref)?;
match worker.execution_handle.clone() {
Some(handle) => {
let backend = state.execution_backend.clone().ok_or_else(|| {
RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id,
message: "runtime has no execution backend".to_string(),
}
})?;
(Some(backend), Some(handle))
}
None if worker.status == WorkerStatus::Stopped => (None, None),
None => {
return Err(RuntimeError::WorkerExecutionUnavailable {
worker_id: worker_ref.worker_id,
message: "worker has no execution handle".to_string(),
});
}
}
};
if let (Some(backend), Some(handle)) = (backend, handle) {
let result = backend.stop_worker(&handle);
if !result.is_accepted() {
return Err(RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id,
operation: result.operation,
outcome: result.outcome,
message: result.message_or_default(),
result,
});
}
self.transition_worker(worker_ref, WorkerStatus::Stopped)?;
}
let _ = reason;
self.transition_worker(worker_ref, WorkerStatus::Stopped)
let state = self.lock()?;
let worker = state.worker(worker_ref)?;
Ok(WorkerLifecycleAck {
worker_ref: worker_ref.clone(),
status: worker.status,
worker_state: worker.worker_state.clone(),
})
}
/// Cancel a Worker through a workspace-scoped Runtime authorization context.
@@ -2185,10 +2278,26 @@ impl Runtime {
#[cfg(feature = "ws-server")]
fn execution_context(&self, worker_ref: WorkerRef) -> crate::execution::WorkerExecutionContext {
let runtime = self.clone();
let runtime = Arc::downgrade(&self.inner);
crate::execution::WorkerExecutionContext::new(
worker_ref,
Arc::new(move |worker_ref, payload| runtime.observe_worker_event(&worker_ref, payload)),
Arc::new(move |worker_ref, payload| {
let runtime = runtime.upgrade().ok_or(RuntimeError::RuntimeStopped)?;
let mut state = runtime.lock().map_err(|_| RuntimeError::StatePoisoned)?;
state.ensure_worker_ref(&worker_ref)?;
let worker_state_changed =
state.project_protocol_event_to_worker_state(&worker_ref, &payload);
let activity_changed =
state.project_internal_worker_activity(&worker_ref, &payload);
if worker_state_changed || activity_changed {
state.publish_worker_upsert(worker_ref.worker_id)?;
}
if worker_state_changed {
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
}
Ok(state.push_worker_observation_event(worker_ref, payload))
}),
)
}
@@ -2199,21 +2308,12 @@ impl Runtime {
#[cfg(feature = "fs-store")]
fn restore_persisted_worker_executions(&self) -> Result<(), RuntimeError> {
#[derive(Clone)]
struct RestoreCandidate {
worker_ref: WorkerRef,
request: CreateWorkerRequest,
run_generation: u64,
previous_working_directory: Option<CatalogWorkingDirectoryStatus>,
config_bundle: Option<ConfigBundle>,
}
let candidates = {
let mut state = self.lock()?;
let state = self.lock()?;
if state.execution_backend.is_none() {
return Ok(());
}
let worker_ids = state
state
.workers
.values()
.filter(|worker| {
@@ -2222,80 +2322,22 @@ impl Runtime {
&& worker.status.is_active()
&& worker.restore_intent == WorkerRestoreIntent::Automatic
})
.map(|worker| worker.worker_id)
.collect::<Vec<_>>();
let mut candidates = Vec::with_capacity(worker_ids.len());
for worker_id in worker_ids {
let (worker_ref, request, previous_working_directory, run_generation) = {
let worker = state
.workers
.get(&worker_id)
.expect("collected Worker exists");
(
worker.worker_ref.clone(),
worker.request.clone(),
worker.working_directory.clone(),
worker.run_generation.saturating_add(1).max(1),
)
};
state
.workers
.get_mut(&worker_id)
.expect("collected Worker exists")
.run_generation = run_generation;
state.persist_worker(&worker_id)?;
candidates.push(RestoreCandidate {
worker_ref,
request,
run_generation,
previous_working_directory,
config_bundle: None,
});
}
candidates
.map(|worker| worker.worker_ref.clone())
.collect::<Vec<_>>()
};
for candidate in candidates {
let (backend, workspace_scope) = {
let state = self.lock()?;
let workspace_scope = candidate.request.workspace_api.as_ref().and_then(|api| {
state
.workspace_owners
.get(&api.workspace_id)
.map(|server_id| RuntimeWorkspaceScope::new(&api.workspace_id, server_id))
});
(state.execution_backend.clone(), workspace_scope)
};
let Some(backend) = backend else {
return Ok(());
};
let request = WorkerExecutionRestoreRequest {
worker_ref: candidate.worker_ref.clone(),
run_generation: candidate.run_generation,
request: candidate.request,
workspace_scope,
context: self.execution_context(candidate.worker_ref.clone()),
previous_working_directory: candidate.previous_working_directory,
working_directory: None,
config_bundle: candidate.config_bundle,
};
match backend.restore_worker(request) {
WorkerExecutionSpawnResult::Connected {
handle,
worker_state,
working_directory,
} => self.commit_restored_worker_execution(
&candidate.worker_ref,
handle,
worker_state,
WorkerStatus::Idle,
working_directory,
)?,
WorkerExecutionSpawnResult::Rejected(result)
| WorkerExecutionSpawnResult::Errored(result) => {
let mut state = self.lock()?;
state.record_restore_failure(&candidate.worker_ref, result)?;
for worker_ref in candidates {
let operation_lock = self.worker_operation_lock(worker_ref.worker_id)?;
let _operation_guard = operation_lock
.lock()
.map_err(|_| RuntimeError::StatePoisoned)?;
match self.restore_worker_under_lock(&worker_ref, WorkerRestoreMode::Automatic) {
Ok(_) => {}
Err(RuntimeError::WorkerExecutionRejected { .. }) => {
// The failed restore and fail-closed terminal status were
// persisted by restore_worker_under_lock.
}
Err(error) => return Err(error),
}
}
Ok(())
@@ -2313,6 +2355,12 @@ impl Runtime {
state.ensure_worker_ref(worker_ref)?;
{
let worker = state.worker_mut(worker_ref)?;
if worker.execution_handle.is_some() {
return Err(RuntimeError::InvalidRequest(format!(
"worker {} already has a current execution",
worker_ref.worker_id
)));
}
worker.execution_handle = Some(handle);
worker.execution_bound = true;
worker.status = status;
@@ -2326,6 +2374,39 @@ impl Runtime {
Ok(())
}
fn cleanup_failed_restore(
&self,
backend: &WorkerExecutionBackendRef,
worker_ref: &WorkerRef,
handle: &WorkerExecutionHandle,
) -> Result<(), RuntimeError> {
let result = backend.stop_worker(handle);
if !result.is_accepted() {
return Err(RuntimeError::WorkerExecutionRejected {
worker_id: worker_ref.worker_id,
operation: result.operation,
outcome: result.outcome,
message: format!(
"restored execution cleanup failed after catalog commit failure: {}",
result.message_or_default()
),
result,
});
}
let mut state = self.lock()?;
let worker = state.worker_mut(worker_ref)?;
worker.execution_handle = None;
worker.status = WorkerStatus::Stopped;
worker.worker_state = None;
worker.restore_intent = WorkerRestoreIntent::Explicit;
worker.internal_workers.clear();
state.publish_worker_upsert(worker_ref.worker_id)?;
state.persist_runtime_snapshot()?;
state.persist_worker(&worker_ref.worker_id)?;
Ok(())
}
/// Bind the Backend registry identity once. Retention evidence fails closed
/// until the Runtime host supplies this trusted configuration.
pub fn bind_runtime_identity(&self, runtime_id: &str) -> Result<(), RuntimeError> {
@@ -3670,6 +3751,7 @@ fn runtime_worker_create_failure_fields(
) -> (&'static str, Option<String>, Option<String>) {
match error {
RuntimeError::RuntimeStopped => ("runtime_stopped", None, None),
RuntimeError::RuntimeStoreAlreadyOpen { .. } => ("runtime_store_already_open", None, None),
RuntimeError::InvalidInitialInputKind { .. } => ("invalid_initial_input_kind", None, None),
RuntimeError::WorkerNotFound { .. } => ("worker_not_found", None, None),
RuntimeError::WorkerExecutionUnavailable { .. } => {
@@ -3854,7 +3936,7 @@ mod tests {
use std::collections::BTreeMap;
#[cfg(feature = "fs-store")]
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
use std::sync::{Arc, Barrier, Condvar, Mutex};
#[test]
fn worker_create_failure_fields_exclude_raw_error_messages() {
@@ -4694,11 +4776,49 @@ mod tests {
.with_computed_digest()
}
#[derive(Default)]
struct RestoreGate {
entered: Mutex<u64>,
entered_changed: Condvar,
released: Mutex<bool>,
released_changed: Condvar,
}
impl RestoreGate {
fn enter_and_wait(&self) {
let mut entered = self.entered.lock().unwrap();
*entered += 1;
self.entered_changed.notify_all();
drop(entered);
let mut released = self.released.lock().unwrap();
while !*released {
released = self.released_changed.wait(released).unwrap();
}
}
fn wait_for_entered(&self, expected: u64, timeout: std::time::Duration) -> bool {
let entered = self.entered.lock().unwrap();
let (entered, _) = self
.entered_changed
.wait_timeout_while(entered, timeout, |entered| *entered < expected)
.unwrap();
*entered >= expected
}
fn release(&self) {
*self.released.lock().unwrap() = true;
self.released_changed.notify_all();
}
}
#[derive(Default)]
struct TestExecutionBackend {
dispatch_result: Mutex<Option<WorkerExecutionResult>>,
stop_result: Mutex<Option<WorkerExecutionResult>>,
stop_gate: Mutex<Option<Arc<RestoreGate>>>,
stop_count: Mutex<u64>,
restore_result: Mutex<Option<WorkerExecutionSpawnResult>>,
restore_gate: Mutex<Option<Arc<RestoreGate>>>,
restore_count: Mutex<u64>,
run_generations: Mutex<Vec<u64>>,
config_bundles: Mutex<Vec<Option<ConfigBundle>>>,
@@ -4828,6 +4948,10 @@ mod tests {
request: WorkerExecutionRestoreRequest,
) -> WorkerExecutionSpawnResult {
*self.restore_count.lock().unwrap() += 1;
let restore_gate = self.restore_gate.lock().unwrap().clone();
if let Some(gate) = restore_gate {
gate.enter_and_wait();
}
self.run_generations
.lock()
.unwrap()
@@ -4888,6 +5012,11 @@ mod tests {
}
fn stop_worker(&self, _handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
*self.stop_count.lock().unwrap() += 1;
let stop_gate = self.stop_gate.lock().unwrap().clone();
if let Some(gate) = stop_gate {
gate.enter_and_wait();
}
self.stop_result
.lock()
.unwrap()
@@ -5660,6 +5789,131 @@ mod tests {
);
}
#[test]
fn failed_stop_cleanup_retains_execution_for_retry() {
let (runtime, backend) = runtime_and_backend();
runtime.store_config_bundle(test_bundle()).unwrap();
let created = runtime
.create_worker(task_request("retry stop cleanup"))
.unwrap();
backend.set_stop_result(WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
"cleanup is still pending",
));
let error = runtime.stop_worker(&created.worker_ref, None).unwrap_err();
assert!(matches!(
error,
RuntimeError::WorkerExecutionRejected {
operation: WorkerExecutionOperation::Stop,
..
}
));
assert_eq!(
runtime.worker_detail(&created.worker_ref).unwrap().status,
WorkerStatus::Idle
);
let stopped = runtime.stop_worker(&created.worker_ref, None).unwrap();
assert_eq!(stopped.status, WorkerStatus::Stopped);
assert_eq!(*backend.stop_count.lock().unwrap(), 2);
}
#[test]
fn restore_waits_for_in_flight_stop_cleanup() {
let backend = Arc::new(TestExecutionBackend::default());
let runtime = Arc::new(
Runtime::with_execution_backend(RuntimeOptions::default(), backend.clone()).unwrap(),
);
runtime.store_config_bundle(test_bundle()).unwrap();
let created = runtime
.create_worker(task_request("stop restore race"))
.unwrap();
let gate = Arc::new(RestoreGate::default());
*backend.stop_gate.lock().unwrap() = Some(gate.clone());
let stop_runtime = runtime.clone();
let stop_ref = created.worker_ref.clone();
let stopping = std::thread::spawn(move || stop_runtime.stop_worker(&stop_ref, None));
assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2)));
let restore_runtime = runtime.clone();
let restore_ref = created.worker_ref.clone();
let restoring = std::thread::spawn(move || restore_runtime.restore_worker(&restore_ref));
std::thread::sleep(std::time::Duration::from_millis(100));
assert_eq!(
*backend.restore_count.lock().unwrap(),
0,
"restore reached the backend before stop cleanup completed"
);
gate.release();
stopping.join().unwrap().unwrap();
let restored = restoring.join().unwrap().unwrap();
assert_eq!(*backend.stop_count.lock().unwrap(), 1);
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(restored.status, WorkerStatus::Idle);
assert_eq!(
restored
.worker_state
.as_ref()
.map(|state| state.execution_generation),
Some(2)
);
}
#[test]
fn concurrent_exact_restores_share_one_execution() {
let backend = Arc::new(TestExecutionBackend::default());
let runtime = Arc::new(
Runtime::with_execution_backend(RuntimeOptions::default(), backend.clone()).unwrap(),
);
runtime.store_config_bundle(test_bundle()).unwrap();
let created = runtime
.create_worker(task_request("concurrent restore"))
.unwrap();
runtime.stop_worker(&created.worker_ref, None).unwrap();
let gate = Arc::new(RestoreGate::default());
*backend.restore_gate.lock().unwrap() = Some(gate.clone());
let start = Arc::new(Barrier::new(3));
let mut threads = Vec::new();
for _ in 0..2 {
let runtime = runtime.clone();
let worker_ref = created.worker_ref.clone();
let start = start.clone();
threads.push(std::thread::spawn(move || {
start.wait();
runtime.restore_worker(&worker_ref)
}));
}
start.wait();
assert!(gate.wait_for_entered(1, std::time::Duration::from_secs(2)));
let duplicate_entered = gate.wait_for_entered(2, std::time::Duration::from_millis(100));
gate.release();
let restored = threads
.into_iter()
.map(|thread| thread.join().unwrap().unwrap())
.collect::<Vec<_>>();
assert!(
!duplicate_entered,
"a second backend restore ran concurrently"
);
assert_eq!(*backend.restore_count.lock().unwrap(), 1);
assert_eq!(restored[0].worker_ref, restored[1].worker_ref);
assert_eq!(
restored[0]
.worker_state
.as_ref()
.map(|state| state.execution_generation),
restored[1]
.worker_state
.as_ref()
.map(|state| state.execution_generation)
);
}
#[test]
fn restored_worker_exposes_the_backend_initial_state_snapshot() {
let (runtime, _) = runtime_and_backend();
@@ -6999,6 +7253,7 @@ mod tests {
)
.unwrap();
drop(corrupt_runtime);
drop(corrupt_store);
let err = Runtime::with_fs_store(crate::fs_store::FsRuntimeStoreOptions {
root: corrupt_root.clone(),
runtime_id: "test-runtime".to_string(),
@@ -7030,6 +7285,7 @@ mod tests {
worker_dirs.sort_by_key(|entry| entry.path());
std::fs::remove_file(worker_dirs[0].path().join("worker.json")).unwrap();
drop(missing_runtime);
drop(missing_store);
let loaded = Runtime::with_fs_store(crate::fs_store::FsRuntimeStoreOptions {
root: missing_root.clone(),
runtime_id: "test-runtime".to_string(),
crates/worker-runtime/src/worker_backend.rs
+261 -83
View File
@@ -10,7 +10,7 @@
use std::collections::HashMap;
use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, RwLock, mpsc};
use std::time::Duration;
@@ -95,6 +95,7 @@ const USER_INPUT_COMMIT_TIMEOUT: Duration = Duration::from_secs(9);
pub struct RuntimeWorkerController {
pub handle: WorkerHandle,
pub shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
pub controller_task: tokio::task::JoinHandle<()>,
pub workspace_client: Arc<dyn WorkspaceClient>,
}
@@ -978,7 +979,8 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
run_dir.display()
),
})?;
let (handle, shutdown_rx) = (started.handle, started.shutdown);
let (handle, shutdown_rx, controller_task) =
(started.handle, started.shutdown, started.controller_task);
if flow_transition_enabled {
handle.shared_state.enable_flow_transition();
}
@@ -990,6 +992,7 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
controller_task,
workspace_client,
})
}
@@ -1172,7 +1175,8 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
run_dir.display()
),
})?;
let (handle, shutdown_rx) = (started.handle, started.shutdown);
let (handle, shutdown_rx, controller_task) =
(started.handle, started.shutdown, started.controller_task);
if flow_transition_enabled {
handle.shared_state.enable_flow_transition();
}
@@ -1184,15 +1188,117 @@ impl RuntimeWorkerFactory for ProfileRuntimeWorkerFactory {
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
controller_task,
workspace_client,
})
}
}
#[derive(Clone)]
struct RuntimeExecutionTaskScope {
tasks: Arc<Mutex<Vec<RuntimeExecutionTask>>>,
terminal_failure: Arc<Mutex<Option<String>>>,
}
struct RuntimeExecutionTask {
name: &'static str,
task: tokio::task::JoinHandle<()>,
abort_before_join: bool,
}
impl RuntimeExecutionTaskScope {
fn new(controller_task: tokio::task::JoinHandle<()>) -> Self {
Self {
tasks: Arc::new(Mutex::new(vec![RuntimeExecutionTask {
name: "controller",
task: controller_task,
abort_before_join: false,
}])),
terminal_failure: Arc::new(Mutex::new(None)),
}
}
fn push(&self, name: &'static str, task: tokio::task::JoinHandle<()>, abort_before_join: bool) {
let mut tasks = match self.tasks.lock() {
Ok(tasks) => tasks,
Err(poisoned) => poisoned.into_inner(),
};
tasks.push(RuntimeExecutionTask {
name,
task,
abort_before_join,
});
}
fn abort_all(&self) {
let mut tasks = match self.tasks.lock() {
Ok(tasks) => tasks,
Err(poisoned) => poisoned.into_inner(),
};
for task in tasks.drain(..) {
task.task.abort();
}
}
async fn join(&self) -> Result<(), String> {
if let Some(message) = self
.terminal_failure
.lock()
.map_err(|_| "execution task failure lock is poisoned".to_string())?
.clone()
{
return Err(message);
}
loop {
let next = {
let mut tasks = self
.tasks
.lock()
.map_err(|_| "execution task registry lock is poisoned".to_string())?;
if tasks.is_empty() {
None
} else {
Some(tasks.remove(0))
}
};
let Some(mut task) = next else {
return Ok(());
};
let name = task.name;
if task.abort_before_join {
task.task.abort();
}
match tokio::time::timeout(Duration::from_secs(5), &mut task.task).await {
Ok(Ok(())) => {}
Ok(Err(error)) if task.abort_before_join && error.is_cancelled() => {}
Ok(Err(error)) => {
let message = format!("{name} task failed while stopping Worker: {error}");
if let Ok(mut failure) = self.terminal_failure.lock() {
*failure = Some(message.clone());
}
return Err(message);
}
Err(_) => {
self.tasks
.lock()
.map_err(|_| "execution task registry lock is poisoned".to_string())?
.insert(0, task);
return Err(format!(
"{name} task did not stop before timeout; stop remains retryable"
));
}
}
}
}
}
#[derive(Clone)]
struct RuntimeWorkerExecution {
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
shutdown_requested: Arc<AtomicBool>,
tasks: RuntimeExecutionTaskScope,
worker_state: Arc<RwLock<protocol::WorkerStateSnapshot>>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
}
@@ -1261,7 +1367,10 @@ where
.map_err(|err| format!("worker adapter task did not complete: {err}"))?
}
fn spawn_on_adapter_runtime<Fut>(&self, task: Fut) -> Result<(), String>
fn spawn_on_adapter_runtime<Fut>(
&self,
task: Fut,
) -> Result<tokio::task::JoinHandle<()>, String>
where
Fut: Future<Output = ()> + Send + 'static,
{
@@ -1272,8 +1381,7 @@ where
let runtime = runtime
.as_ref()
.ok_or_else(|| "worker adapter runtime is shutting down".to_string())?;
runtime.spawn(task);
Ok(())
Ok(runtime.spawn(task))
}
fn run_on_adapter_runtime<T, Fut>(&self, task: Fut) -> Result<T, String>
@@ -1461,6 +1569,44 @@ where
.unwrap_or_else(|message| WorkerExecutionResult::errored(operation, message))
}
fn cleanup_unconnected_controller(
&self,
handle: &WorkerHandle,
shutdown: &Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
tasks: &RuntimeExecutionTaskScope,
worker_state: &Arc<RwLock<protocol::WorkerStateSnapshot>>,
) -> Result<(), String> {
let command = next_internal_command(worker_state)?;
let handle = handle.clone();
let shutdown = shutdown.clone();
let tasks_for_join = tasks.clone();
let cleanup = self.run_on_adapter_runtime(async move {
handle
.send(Method::Shutdown { command })
.await
.map_err(|error| format!("failed to request controller cleanup: {error}"))?;
let mut guard = shutdown.lock().await;
if let Some(mut receiver) = guard.take() {
match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await {
Ok(Ok(())) => {}
Ok(Err(_)) => {
return Err("controller cleanup completion channel closed".to_string());
}
Err(_) => {
*guard = Some(receiver);
return Err("controller cleanup confirmation timed out".to_string());
}
}
}
drop(guard);
tasks_for_join.join().await
});
if cleanup.is_err() {
tasks.abort_all();
}
cleanup
}
fn connect_handle(
&self,
operation: WorkerExecutionOperation,
@@ -1468,17 +1614,19 @@ where
bridge_context: crate::execution::WorkerExecutionContext,
handle: WorkerHandle,
shutdown: Arc<tokio::sync::Mutex<Option<worker::ShutdownReceiver>>>,
controller_task: tokio::task::JoinHandle<()>,
working_directory: Option<WorkingDirectoryBinding>,
workspace_client: Option<Arc<dyn WorkspaceClient>>,
) -> WorkerExecutionSpawnResult {
let worker_state = Arc::new(RwLock::new(handle.shared_state.snapshot()));
let tasks = RuntimeExecutionTaskScope::new(controller_task);
#[cfg(feature = "ws-server")]
{
let streams = subscribe_worker_protocol_session(&handle);
let mut events = streams.events;
let mut entry_events = streams.log_entries;
let bridge_worker_state = worker_state.clone();
if let Err(message) = self.spawn_on_adapter_runtime(async move {
let bridge_task = match self.spawn_on_adapter_runtime(async move {
loop {
tokio::select! {
event = events.recv() => {
@@ -1516,10 +1664,20 @@ where
}
}
}) {
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
operation, message,
));
}
Ok(task) => task,
Err(message) => {
let cleanup = self
.cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state)
.err()
.map(|cleanup| format!("; controller cleanup failed: {cleanup}"))
.unwrap_or_default();
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
operation,
format!("{message}{cleanup}"),
));
}
};
tasks.push("protocol bridge", bridge_task, true);
}
#[cfg(not(feature = "ws-server"))]
{
@@ -1529,12 +1687,29 @@ where
let mut workers = match self.workers.lock() {
Ok(workers) => workers,
Err(_) => {
let cleanup = self
.cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state)
.err()
.map(|cleanup| format!("; controller cleanup failed: {cleanup}"))
.unwrap_or_default();
return WorkerExecutionSpawnResult::Errored(WorkerExecutionResult::errored(
operation,
"worker adapter registry lock is poisoned",
format!("worker adapter registry lock is poisoned{cleanup}"),
));
}
};
if workers.contains_key(&worker_ref) {
drop(workers);
let cleanup = self
.cleanup_unconnected_controller(&handle, &shutdown, &tasks, &worker_state)
.err()
.map(|cleanup| format!("; controller cleanup failed: {cleanup}"))
.unwrap_or_default();
return WorkerExecutionSpawnResult::Rejected(WorkerExecutionResult::busy(
operation,
format!("Worker is already connected to execution backend{cleanup}"),
));
}
let connected_worker_state = worker_state
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
@@ -1544,6 +1719,8 @@ where
RuntimeWorkerExecution {
handle,
shutdown,
shutdown_requested: Arc::new(AtomicBool::new(false)),
tasks,
worker_state,
workspace_client,
},
@@ -1820,6 +1997,7 @@ where
bridge_context,
controller.handle,
controller.shutdown,
controller.controller_task,
working_directory,
Some(controller.workspace_client),
)
@@ -1919,6 +2097,7 @@ where
bridge_context,
controller.handle,
controller.shutdown,
controller.controller_task,
working_directory,
Some(controller.workspace_client),
)
@@ -2103,48 +2282,69 @@ where
}
};
let Some(execution) = execution else {
return WorkerExecutionResult::rejected(
WorkerExecutionOperation::Stop,
"execution handle does not reference a live Worker",
);
// The execution backend cleanup may have committed before the
// Runtime catalog commit failed. Treat the retry as converged so
// the Runtime can durably finish its Stopped transition.
return WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop);
};
let artifact_cleanup = execution.handle.clone();
let shutdown = execution.shutdown.clone();
let command = match next_internal_command(&execution.worker_state) {
Ok(command) => command,
Err(error) => {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error);
}
};
let result = self.send_method(
WorkerExecutionOperation::Stop,
execution.handle.clone(),
Method::Shutdown { command },
);
if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
return result;
}
let shutdown_wait = self.run_on_adapter_runtime(async move {
let mut guard = shutdown.lock().await;
let Some(mut receiver) = guard.take() else {
return Ok(());
let first_request = !execution.shutdown_requested.swap(true, Ordering::AcqRel);
let result = if first_request {
let command = match next_internal_command(&execution.worker_state) {
Ok(command) => command,
Err(error) => {
execution.shutdown_requested.store(false, Ordering::Release);
return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, error);
}
};
match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await {
Ok(Ok(())) => Ok(()),
Ok(Err(_)) => Err("Worker shutdown completion channel closed".to_string()),
Err(_) => {
*guard = Some(receiver);
Err("Worker shutdown confirmation timed out; stop remains retryable".into())
let result = self.send_method(
WorkerExecutionOperation::Stop,
execution.handle.clone(),
Method::Shutdown { command },
);
if result.outcome != crate::execution::WorkerExecutionOutcome::Accepted {
execution.shutdown_requested.store(false, Ordering::Release);
return result;
}
result
} else {
WorkerExecutionResult::accepted(WorkerExecutionOperation::Stop)
};
let shutdown = execution.shutdown.clone();
let tasks = execution.tasks.clone();
let shutdown_wait = self.run_on_adapter_runtime(async move {
{
let mut guard = shutdown.lock().await;
if let Some(mut receiver) = guard.take() {
match tokio::time::timeout(Duration::from_secs(5), &mut receiver).await {
Ok(Ok(())) => {}
Ok(Err(_)) => {
return Err("Worker shutdown completion channel closed".to_string());
}
Err(_) => {
*guard = Some(receiver);
return Err(
"Worker shutdown confirmation timed out; stop remains retryable"
.to_string(),
);
}
}
}
}
tasks.join().await
});
if let Err(message) = shutdown_wait {
return WorkerExecutionResult::errored(WorkerExecutionOperation::Stop, message);
}
let artifact_cleanup_error = artifact_cleanup
.delete_uncommitted_uploaded_files()
.err()
.map(|error| format!("uploaded_file_cleanup_failed: {error}"));
if let Err(error) = execution.handle.delete_uncommitted_uploaded_files() {
return WorkerExecutionResult::errored(
WorkerExecutionOperation::Stop,
format!("uploaded_file_cleanup_failed: {error}; stop remains retryable"),
);
}
match self.workers.lock() {
Ok(mut workers) => {
workers.remove(handle.worker_ref());
@@ -2153,13 +2353,7 @@ where
poisoned.into_inner().remove(handle.worker_ref());
}
}
if let Some(message) = artifact_cleanup_error {
let mut result = result;
result.message = Some(message);
result
} else {
result
}
result
}
fn cancel_worker(&self, handle: &WorkerExecutionHandle) -> WorkerExecutionResult {
@@ -2353,14 +2547,6 @@ mod tests {
);
}
fn test_command() -> WorkerCommandEnvelope {
WorkerCommandEnvelope {
command_id: 1,
expected_execution_generation: 1,
expected_worker_state_revision: 0,
}
}
fn adapter_command(
backend: &WorkerRuntimeExecutionBackend<MockFactory>,
worker_ref: &WorkerRef,
@@ -2828,6 +3014,7 @@ mod tests {
Ok(RuntimeWorkerController {
handle,
shutdown: Arc::new(tokio::sync::Mutex::new(Some(shutdown_rx))),
controller_task: tokio::spawn(async {}),
workspace_client,
})
}
@@ -3343,7 +3530,7 @@ mod tests {
let root = tempfile::tempdir().unwrap();
let long_component = "embedded-workspace-store-segment".repeat(4);
let runtime_store_dir = root.path().join(long_component);
let worker_ref = WorkerRef::new(crate::identity::WorkerId::from_legacy_u64(1));
let worker_ref = WorkerRef::new(crate::identity::WorkerId::now_v7());
let worker_aggregate_dir = runtime_store_dir
.join("workers")
.join(worker_ref.worker_id.to_string());
@@ -3417,16 +3604,24 @@ mod tests {
assert!(!socket_path.exists());
assert!(run_dir.join("worker.out.log").is_file());
assert!(run_dir.join("worker.err.log").is_file());
let worker_state = Arc::new(RwLock::new(controller.handle.shared_state.snapshot()));
controller
.handle
.send(Method::Shutdown {
command: test_command(),
command: next_internal_command(&worker_state).unwrap(),
})
.await
.unwrap();
if let Some(receiver) = controller.shutdown.lock().await.take() {
receiver.await.unwrap();
tokio::time::timeout(Duration::from_secs(5), receiver)
.await
.expect("controller shutdown signal timed out")
.unwrap();
}
tokio::time::timeout(Duration::from_secs(5), controller.controller_task)
.await
.expect("controller task join timed out")
.unwrap();
assert!(!socket_path.exists());
}
@@ -3529,25 +3724,8 @@ mod tests {
);
assert!(!first_run_socket.exists());
let (handle, shutdown) = {
let workers = backend.workers.lock().unwrap();
let execution = workers.get(&worker.worker_ref).unwrap();
(execution.handle.clone(), execution.shutdown.clone())
};
backend
.run_on_adapter_runtime(async move {
handle
.send(Method::Shutdown {
command: test_command(),
})
.await
.map_err(|error| error.to_string())?;
if let Some(receiver) = shutdown.lock().await.take() {
receiver.await.map_err(|error| error.to_string())?;
}
Ok(())
})
.unwrap();
let handle = WorkerExecutionHandle::new(worker.worker_ref.clone(), backend.backend_id());
assert!(backend.stop_worker(&handle).is_accepted());
drop(runtime);
drop(backend);
crates/worker/src/bootstrap.rs
+8 -3
View File
@@ -67,6 +67,7 @@ pub struct PreparedWorker<C: LlmClient, St: Store> {
pub struct BootstrappedWorker {
pub handle: WorkerHandle,
pub shutdown: ShutdownReceiver,
pub controller_task: tokio::task::JoinHandle<()>,
}
#[derive(Debug, Error)]
@@ -227,7 +228,7 @@ where
runtime_base,
bash_output_dir,
} => {
WorkerController::spawn_with_transport(
WorkerController::spawn_with_transport_owned(
worker,
&runtime_base,
&bash_output_dir,
@@ -239,7 +240,7 @@ where
run_dir,
bash_output_dir,
} => {
WorkerController::spawn_runtime_managed_run_with_transport(
WorkerController::spawn_runtime_managed_run_with_transport_owned(
worker,
&run_dir,
&bash_output_dir,
@@ -250,7 +251,11 @@ where
};
match controller {
Ok((handle, shutdown)) => Ok(BootstrappedWorker { handle, shutdown }),
Ok((handle, shutdown, controller_task)) => Ok(BootstrappedWorker {
handle,
shutdown,
controller_task,
}),
Err(source) => {
let cleanup_failed = match cleanup_session {
Some(session) => session.close().await.is_err(),
crates/worker/src/controller.rs
+90 -29
View File
@@ -486,7 +486,7 @@ impl WorkerController {
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
Self::spawn_inner(
let (handle, shutdown, _task) = Self::spawn_inner(
worker,
runtime_base,
bash_output_dir,
@@ -494,7 +494,8 @@ impl WorkerController {
None,
WorkerControllerTransport::UnixSocket,
)
.await
.await?;
Ok((handle, shutdown))
}
/// Spawn a direct Worker while letting an in-process host select the
@@ -505,6 +506,22 @@ impl WorkerController {
bash_output_dir: &Path,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let (handle, shutdown, _task) =
Self::spawn_with_transport_owned(worker, runtime_base, bash_output_dir, transport)
.await?;
Ok((handle, shutdown))
}
pub(crate) async fn spawn_with_transport_owned<C, St>(
worker: Worker<C, St>,
runtime_base: &Path,
bash_output_dir: &Path,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
@@ -535,7 +552,7 @@ impl WorkerController {
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
Self::spawn_inner(
let (handle, shutdown, _task) = Self::spawn_inner(
worker,
runtime_base,
bash_output_dir,
@@ -543,7 +560,8 @@ impl WorkerController {
None,
WorkerControllerTransport::UnixSocket,
)
.await
.await?;
Ok((handle, shutdown))
}
/// Spawn into an exact persistent `runs/<generation>` directory.
@@ -573,6 +591,26 @@ impl WorkerController {
bash_output_dir: &Path,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
{
let (handle, shutdown, _task) = Self::spawn_runtime_managed_run_with_transport_owned(
worker,
run_dir,
bash_output_dir,
transport,
)
.await?;
Ok((handle, shutdown))
}
pub(crate) async fn spawn_runtime_managed_run_with_transport_owned<C, St>(
worker: Worker<C, St>,
run_dir: &Path,
bash_output_dir: &Path,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
@@ -598,7 +636,7 @@ impl WorkerController {
runtime_managed: bool,
runtime_run: Option<&Path>,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
@@ -629,7 +667,7 @@ impl WorkerController {
runtime_managed: bool,
runtime_run: Option<&Path>,
transport: WorkerControllerTransport,
) -> Result<(WorkerHandle, ShutdownReceiver), std::io::Error>
) -> Result<(WorkerHandle, ShutdownReceiver, tokio::task::JoinHandle<()>), std::io::Error>
where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + Send + Sync + 'static,
@@ -729,9 +767,9 @@ impl WorkerController {
None,
)
.await?;
if let Some(session) = fs_for_view.as_ref() {
wire_workdir_command_events(session, &in_flight);
}
let command_observer = fs_for_view
.as_ref()
.and_then(|session| wire_workdir_command_events(session, &in_flight));
// Intake role Workers self-terminate only after a successful
// TicketIntakeReady turn has fully settled back to Idle. The request
@@ -805,7 +843,7 @@ impl WorkerController {
let pause_tx = worker.engine_mut().pause_sender();
let notify_buffer = worker.notify_buffer_handle();
tokio::spawn(controller_loop(
let controller_task = tokio::spawn(controller_loop(
worker,
method_rx,
working_event_tx,
@@ -820,26 +858,27 @@ impl WorkerController {
shutdown_tx,
socket_server,
shutdown_after_idle,
command_observer,
));
Ok((handle, shutdown_rx))
Ok((handle, shutdown_rx, controller_task))
}
}
pub(crate) fn wire_workdir_command_events(
session: &Arc<dyn WorkdirSession>,
in_flight: &InFlightEvents,
) {
) -> Option<tokio::task::JoinHandle<()>> {
in_flight.replace_command_snapshot(protocol_command_snapshots(session.as_ref()));
let Some(mut events) = session.subscribe_command_events() else {
return;
return None;
};
// Keep only a weak reference in the observer task. Holding the session
// strongly here would keep its broadcast sender alive forever and prevent
// the receiver from observing closure during Worker teardown.
let session = Arc::downgrade(session);
let in_flight = in_flight.clone();
tokio::spawn(async move {
Some(tokio::spawn(async move {
loop {
match events.recv().await {
Ok(event) => in_flight.publish_command_event(protocol_command_event(event)),
@@ -853,7 +892,7 @@ pub(crate) fn wire_workdir_command_events(
Err(broadcast::error::RecvError::Closed) => break,
}
}
});
}))
}
fn protocol_command_snapshots(session: &dyn WorkdirSession) -> Vec<ProtocolCommandSnapshot> {
@@ -1507,6 +1546,7 @@ async fn controller_loop<C, St>(
shutdown_tx: oneshot::Sender<()>,
socket_server: Option<SocketServer>,
shutdown_after_idle: ShutdownAfterIdleRequest,
mut command_observer: Option<tokio::task::JoinHandle<()>>,
) where
C: LlmClient + Clone + 'static,
St: Store + WorkerMetadataStore + Clone + 'static,
@@ -2343,28 +2383,49 @@ async fn controller_loop<C, St>(
}
}
drop(_socket_server);
if let Err(error) = runtime_dir.close_socket().await {
tracing::warn!(%error, "Worker runtime socket cleanup failed");
let had_socket_server = _socket_server.is_some();
if let Some(socket_server) = _socket_server {
socket_server.shutdown().await;
}
while had_socket_server {
match runtime_dir.close_socket().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Worker runtime socket cleanup failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
}
}
// Feature callbacks and tasks share the Worker scope. Stop them before
// Memory/Workdir teardown so they cannot observe a partially closed Worker.
worker.stop_feature_runtime("controller shutdown").await;
let child_cleanup_succeeded = match spawned_registry.shutdown_internal().await {
Ok(()) => true,
Err(error) => {
tracing::warn!(%error, "Internal SubWorker cleanup failed before Workdir shutdown");
false
loop {
match spawned_registry.shutdown_internal().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Internal SubWorker cleanup failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
}
};
}
if child_cleanup_succeeded
&& let Some(session) = worker.workdir_session()
&& let Err(error) = session.close().await
{
tracing::warn!(%error, "Workdir session close failed");
if let Some(session) = worker.workdir_session() {
loop {
match session.close().await {
Ok(()) => break,
Err(error) => {
tracing::warn!(%error, "Workdir session close failed; retrying");
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
}
}
}
if let Some(observer) = command_observer.take() {
observer.abort();
let _ = observer.await;
}
// Report upward that this Worker is stopping before the controller
crates/worker/src/ipc/server.rs
+40 -9
View File
@@ -4,7 +4,8 @@ use std::path::PathBuf;
use protocol::stream::{JsonLineReader, JsonLineWriter};
use tokio::net::UnixListener;
use tokio::task::JoinHandle;
use tokio::sync::oneshot;
use tokio::task::{JoinHandle, JoinSet};
use crate::controller::WorkerHandle;
use crate::ipc::protocol_session::{
@@ -19,7 +20,8 @@ use protocol::{ErrorCode, Event};
/// - Client writes Method lines → forwarded to WorkerController
/// - Worker events → written as Event lines to all connected clients
pub struct SocketServer {
_accept_task: JoinHandle<()>,
accept_task: Option<JoinHandle<()>>,
shutdown: Option<oneshot::Sender<()>>,
path: PathBuf,
}
@@ -33,20 +35,45 @@ impl SocketServer {
let listener = UnixListener::bind(&path)?;
let handle = handle.clone();
let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
let _accept_task = tokio::spawn(async move {
let accept_task = tokio::spawn(async move {
let mut connections = JoinSet::new();
loop {
match listener.accept().await {
Ok((stream, _)) => {
let handle = handle.clone();
tokio::spawn(handle_connection(stream, handle));
tokio::select! {
_ = &mut shutdown_rx => break,
accepted = listener.accept() => match accepted {
Ok((stream, _)) => {
let handle = handle.clone();
connections.spawn(handle_connection(stream, handle));
}
Err(_) => break,
},
completed = connections.join_next(), if !connections.is_empty() => {
let _ = completed;
}
Err(_) => break,
}
}
connections.shutdown().await;
});
Ok(Self { _accept_task, path })
Ok(Self {
accept_task: Some(accept_task),
shutdown: Some(shutdown_tx),
path,
})
}
/// Stop accepting connections and join the server task. Dropping its
/// `JoinSet` cancels every active connection task before this returns.
pub async fn shutdown(mut self) {
if let Some(shutdown) = self.shutdown.take() {
let _ = shutdown.send(());
}
if let Some(task) = self.accept_task.take() {
let _ = task.await;
}
let _ = tokio::fs::remove_file(&self.path).await;
}
/// The socket file path.
@@ -57,6 +84,10 @@ impl SocketServer {
impl Drop for SocketServer {
fn drop(&mut self) {
if let Some(shutdown) = self.shutdown.take() {
let _ = shutdown.send(());
}
let _ = self.accept_task.take();
let _ = std::fs::remove_file(&self.path);
}
}
crates/worker/tests/controller_test.rs
+27
View File
@@ -2219,6 +2219,33 @@ async fn status_json_reflects_worker_name() {
// Socket transport tests
// ---------------------------------------------------------------------------
#[tokio::test]
async fn shutdown_joins_socket_server_with_active_connection() {
use tokio::net::UnixStream;
let worker = make_worker(MockClient::new(simple_text_events())).await;
let runtime_base = tempfile::tempdir().unwrap();
let bash_output_dir = runtime_base.path().join("bash-output");
let (handle, shutdown_rx) =
WorkerController::spawn(worker, runtime_base.path(), &bash_output_dir)
.await
.unwrap();
let socket_path = handle.runtime_dir.socket_path();
let _connection = UnixStream::connect(&socket_path).await.unwrap();
handle
.send(Method::Shutdown {
command: worker_command(&handle),
})
.await
.unwrap();
tokio::time::timeout(std::time::Duration::from_secs(5), shutdown_rx)
.await
.expect("controller should join its socket tasks")
.expect("controller shutdown signal should remain open");
assert!(!socket_path.exists());
}
#[tokio::test]
async fn socket_run_receives_events() {
use protocol::stream::{JsonLineReader, JsonLineWriter};
crates/workspace-server/src/hosts.rs
+5
View File
@@ -5064,6 +5064,11 @@ fn embedded_runtime_diagnostic(error: &EmbeddedRuntimeError) -> RuntimeDiagnosti
DiagnosticSeverity::Warning,
"Embedded Runtime rejected the request".to_string(),
),
EmbeddedRuntimeError::RuntimeStoreAlreadyOpen { .. } => diagnostic(
"embedded_runtime_store_already_open",
DiagnosticSeverity::Error,
"Embedded Runtime store is already owned by another Runtime process".to_string(),
),
EmbeddedRuntimeError::StoreIo { .. }
| EmbeddedRuntimeError::StoreMissing { .. }
| EmbeddedRuntimeError::StoreCorrupt { .. } => diagnostic(