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

プロトコル経由のshutdow経路

Browse Source
This commit is contained in:
Hare
2026-04-16 13:49:53 +09:00
parent 710220c920
commit aa138e6583
11 changed files with 183 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/pod/examples/pod_protocol.rs
+1 -1
View File
@@ -43,7 +43,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let pod = pod::Pod::from_manifest_toml(&toml, store).await?;
let runtime_tmp = tempfile::tempdir()?;
let handle = PodController::spawn(pod, runtime_tmp.path()).await?;
let (handle, _shutdown_rx) = PodController::spawn(pod, runtime_tmp.path()).await?;
// Check initial status via shared state
println!("[shared_state] {}", handle.shared_state.status_json());
crates/pod/src/controller.rs
+37 -15
View File
@@ -4,7 +4,7 @@ use std::sync::Arc;
use llm_worker::WorkerError;
use llm_worker::llm_client::client::LlmClient;
use session_store::Store;
use tokio::sync::{broadcast, mpsc};
use tokio::sync::{broadcast, mpsc, oneshot};
use crate::notifier::Notifier;
use crate::pod::{Pod, PodError, PodRunResult};
@@ -50,17 +50,20 @@ impl PodHandle {
// PodController — actor that owns a Pod
// ---------------------------------------------------------------------------
pub type ShutdownReceiver = oneshot::Receiver<()>;
pub struct PodController;
impl PodController {
pub async fn spawn<C, St>(
mut pod: Pod<C, St>,
runtime_base: &Path,
) -> Result<PodHandle, std::io::Error>
) -> Result<(PodHandle, ShutdownReceiver), std::io::Error>
where
C: LlmClient + 'static,
St: Store + 'static,
{
let (shutdown_tx, shutdown_rx) = oneshot::channel::<()>();
let (method_tx, mut method_rx) = mpsc::channel::<Method>(32);
let (event_tx, _) = broadcast::channel::<Event>(256);
let notifier = Notifier::new(event_tx.clone());
@@ -225,7 +228,7 @@ impl PodController {
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
let new_status = run_with_cancel_support(
let (new_status, shutdown) = run_with_cancel_support(
pod.run(&input),
&mut method_rx,
&event_tx,
@@ -234,7 +237,6 @@ impl PodController {
)
.await;
// Proactive post-run compaction (best-effort).
if new_status == PodStatus::Idle {
if let Err(e) = pod.try_post_run_compact().await {
tracing::warn!(error = %e, "Post-run compaction error");
@@ -251,6 +253,11 @@ impl PodController {
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
break;
}
}
Method::Resume => {
@@ -264,7 +271,7 @@ impl PodController {
shared_state.set_status(PodStatus::Running);
let _ = runtime_dir.write_status(&shared_state).await;
let new_status = run_with_cancel_support(
let (new_status, shutdown) = run_with_cancel_support(
pod.resume(),
&mut method_rx,
&event_tx,
@@ -273,7 +280,6 @@ impl PodController {
)
.await;
// Proactive post-run compaction (best-effort).
if new_status == PodStatus::Idle {
if let Err(e) = pod.try_post_run_compact().await {
tracing::warn!(error = %e, "Post-run compaction error");
@@ -290,6 +296,11 @@ impl PodController {
shared_state.set_status(new_status);
let _ = runtime_dir.write_status(&shared_state).await;
let _ = runtime_dir.write_history(&shared_state).await;
if shutdown {
let _ = event_tx.send(Event::Shutdown);
break;
}
}
Method::Cancel => {
@@ -299,30 +310,39 @@ impl PodController {
});
}
Method::Shutdown => {
let _ = event_tx.send(Event::Shutdown);
break;
}
// GetHistory is handled at the socket layer (direct response).
// If it somehow reaches the controller, ignore it.
Method::GetHistory => {}
}
}
let _ = shutdown_tx.send(());
});
Ok(handle)
Ok((handle, shutdown_rx))
}
}
/// Runs a Pod future while concurrently processing incoming methods.
/// Only `Cancel` is handled during execution; `Run` and `Resume` get errors.
///
/// Returns `(final_status, shutdown_requested)`.
async fn run_with_cancel_support<F>(
pod_future: F,
method_rx: &mut mpsc::Receiver<Method>,
event_tx: &broadcast::Sender<Event>,
cancel_tx: &mpsc::Sender<()>,
shared_state: &Arc<PodSharedState>,
) -> PodStatus
) -> (PodStatus, bool)
where
F: std::future::Future<Output = Result<PodRunResult, PodError>>,
{
tokio::pin!(pod_future);
let mut shutdown_requested = false;
loop {
tokio::select! {
@@ -335,7 +355,7 @@ where
PodRunResult::LimitReached => (PodStatus::Idle, RunResult::LimitReached),
};
let _ = event_tx.send(Event::RunEnd { result: run_result });
status
(status, shutdown_requested)
}
Err(e) => {
let code = worker_error_code(&e);
@@ -343,7 +363,7 @@ where
code,
message: e.to_string(),
});
PodStatus::Idle
(PodStatus::Idle, shutdown_requested)
}
};
}
@@ -352,19 +372,21 @@ where
Some(Method::Cancel) => {
let _ = cancel_tx.try_send(());
}
Some(Method::Shutdown) => {
shutdown_requested = true;
let _ = cancel_tx.try_send(());
}
Some(Method::Run { .. } | Method::Resume) => {
let _ = event_tx.send(Event::Error {
code: ErrorCode::AlreadyRunning,
message: "Pod is already executing a turn".into(),
});
}
Some(Method::GetHistory) => {
// Handled at socket layer; ignore here.
}
Some(Method::GetHistory) => {}
None => {
let _ = cancel_tx.try_send(());
shared_state.set_status(PodStatus::Idle);
return PodStatus::Idle;
return (PodStatus::Idle, false);
}
}
}
crates/pod/src/lib.rs
+1 -1
View File
@@ -19,7 +19,7 @@ mod usage_tracker;
pub use token_counter::{EstimateSource, SplitPoint, TokenEstimate};
pub use controller::{PodController, PodHandle};
pub use controller::{PodController, PodHandle, ShutdownReceiver};
pub use factory::{FactoryError, PodFactory};
pub use notifier::Notifier;
pub use hook::{Hook, HookEventKind, HookRegistryBuilder};
crates/pod/src/main.rs
+7 -8
View File
@@ -157,8 +157,8 @@ async fn main() -> ExitCode {
return ExitCode::FAILURE;
}
};
let handle = match PodController::spawn(pod, &runtime_base).await {
Ok(h) => h,
let (handle, shutdown_rx) = match PodController::spawn(pod, &runtime_base).await {
Ok(pair) => pair,
Err(e) => {
eprintln!("error: failed to start pod controller: {e}");
return ExitCode::FAILURE;
@@ -170,13 +170,12 @@ async fn main() -> ExitCode {
handle.runtime_dir.socket_path()
);
// Wait for shutdown signal
match tokio::signal::ctrl_c().await {
Ok(()) => {
eprintln!("pod: {pod_name} shutting down");
tokio::select! {
_ = tokio::signal::ctrl_c() => {
eprintln!("pod: {pod_name} shutting down (signal)");
}
Err(e) => {
eprintln!("error: failed to listen for signal: {e}");
_ = shutdown_rx => {
eprintln!("pod: {pod_name} shutting down (client request)");
}
}
crates/pod/tests/controller_test.rs
+2 -2
View File
@@ -111,9 +111,9 @@ use pod::PodHandle;
async fn spawn_controller(pod: Pod<MockClient, FsStore>) -> PodHandle {
let tmp = tempfile::tempdir().unwrap();
let runtime_base = tmp.path().to_owned();
// Leak tempdir so it survives the test
std::mem::forget(tmp);
PodController::spawn(pod, &runtime_base).await.unwrap()
let (handle, _shutdown_rx) = PodController::spawn(pod, &runtime_base).await.unwrap();
handle
}
// ---------------------------------------------------------------------------