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
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};