diff --git a/lib/crates/fabro-server/src/lib.rs b/lib/crates/fabro-server/src/lib.rs index 76805781c..e9c18aad1 100644 --- a/lib/crates/fabro-server/src/lib.rs +++ b/lib/crates/fabro-server/src/lib.rs @@ -49,6 +49,7 @@ pub mod static_files; pub mod test_support; pub mod web_auth; mod worker_control; +mod worker_runtime; mod worker_token; pub use error::{ApiError, Error, Result}; diff --git a/lib/crates/fabro-server/src/serve.rs b/lib/crates/fabro-server/src/serve.rs index 3042cbae1..2b71e6d4a 100644 --- a/lib/crates/fabro-server/src/serve.rs +++ b/lib/crates/fabro-server/src/serve.rs @@ -778,6 +778,8 @@ where shutdown: shutdown.clone(), #[cfg(test)] worker_control_bus: None, + #[cfg(test)] + worker_runtime: None, #[cfg(any(test, feature = "test-support"))] automation_materializer_override: None, })?; diff --git a/lib/crates/fabro-server/src/server.rs b/lib/crates/fabro-server/src/server.rs index 69de23079..11963ba1c 100644 --- a/lib/crates/fabro-server/src/server.rs +++ b/lib/crates/fabro-server/src/server.rs @@ -123,8 +123,8 @@ use futures_util::future::join_all; use sha2::{Digest, Sha256}; use tempfile::NamedTempFile; use tokio::fs; -use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader}; -use tokio::process::{ChildStderr, Command}; +use tokio::io::{AsyncBufReadExt, AsyncRead, AsyncWriteExt, BufReader}; +use tokio::process::Command; use tokio::sync::broadcast::error::RecvError; use tokio::sync::{ Mutex as AsyncMutex, Notify, OwnedMutexGuard, RwLock as AsyncRwLock, Semaphore, broadcast, @@ -158,9 +158,12 @@ use crate::principal_middleware::{ use crate::request_id::{self, RequestId}; use crate::run_files::{FilesInFlight, new_files_in_flight}; use crate::server_secrets::{LlmClientResult, ServerSecrets}; -use crate::spawn_env::{apply_render_graph_env, apply_worker_env}; +use crate::spawn_env::apply_render_graph_env; use crate::startup::load_startup_vault; use crate::worker_control::{LocalWorkerControlBus, WorkerControlBus, WorkerControlBusError}; +use crate::worker_runtime::{ + LocalWorkerRuntime, WorkerExit, WorkerLaunchSpec, WorkerRef, WorkerRuntime, +}; use crate::worker_token::{WorkerScopeSet, WorkerTokenKeys, issue_worker_token_with_scopes}; use crate::{ canonical_host, demo, diagnostics, run_manifest, security_headers, static_files, web_auth, @@ -253,8 +256,7 @@ struct ManagedRun { checkpoint: Option, cancel_tx: Option>, cancel_token: Option, - worker_pid: Option, - worker_pgid: Option, + worker_ref: Option, run_dir: Option, execution_mode: RunExecutionMode, } @@ -1068,6 +1070,7 @@ pub struct AppState { resource_sampler: resource_sampler::ResourceSampler, max_concurrent_runs: usize, pub(crate) worker_control_bus: Arc, + pub(crate) worker_runtime: Arc, scheduler_notify: Notify, global_event_tx: broadcast::Sender, /// Per-run coalescing registry for `GET /runs/{id}/files`. Concurrent @@ -1225,6 +1228,8 @@ pub(crate) struct AppStateConfig { pub(crate) shutdown: CancellationToken, #[cfg(test)] pub(crate) worker_control_bus: Option>, + #[cfg(test)] + pub(crate) worker_runtime: Option>, #[cfg(any(test, feature = "test-support"))] pub(crate) automation_materializer_override: Option>, } @@ -2262,6 +2267,8 @@ pub(crate) fn build_app_state(config: AppStateConfig) -> anyhow::Result anyhow::Result = { + #[cfg(test)] + { + worker_runtime.unwrap_or_else(|| Arc::new(LocalWorkerRuntime::new())) + } + #[cfg(not(test))] + { + Arc::new(LocalWorkerRuntime::new()) + } + }; Ok(Arc::new(AppState { runs: Mutex::new(HashMap::new()), aggregate_billing: Mutex::new(BillingAccumulator::default()), @@ -2371,6 +2388,7 @@ pub(crate) fn build_app_state(config: AppStateConfig) -> anyhow::Result S } async fn terminate_worker_for_deletion( - worker_pid: Option, - worker_pgid: Option, + worker_runtime: &Arc, + worker_ref: Option, grace: Duration, ) { - #[cfg(unix)] - if let Some(process_group_id) = worker_pgid.or(worker_pid) { - fabro_proc::sigterm_process_group(process_group_id); + let Some(worker_ref) = worker_ref else { + return; + }; - let deadline = Instant::now() + grace; - while Instant::now() < deadline && fabro_proc::process_group_alive(process_group_id) { - sleep(Duration::from_millis(50)).await; - } + worker_runtime.request_stop(&worker_ref).await; - if fabro_proc::process_group_alive(process_group_id) { - fabro_proc::sigkill_process_group(process_group_id); - - let kill_deadline = Instant::now() + Duration::from_secs(1); - while Instant::now() < kill_deadline - && fabro_proc::process_group_alive(process_group_id) - { - sleep(Duration::from_millis(50)).await; - } - } + let deadline = Instant::now() + grace; + while Instant::now() < deadline && worker_runtime.is_alive(&worker_ref).await { + sleep(Duration::from_millis(50)).await; } - #[cfg(not(unix))] - if let Some(worker_pid) = worker_pid { - fabro_proc::sigterm(worker_pid); - - let deadline = Instant::now() + grace; - while Instant::now() < deadline && fabro_proc::process_running(worker_pid) { + if worker_runtime.is_alive(&worker_ref).await { + worker_runtime.force_stop(&worker_ref).await; + let kill_deadline = Instant::now() + Duration::from_secs(1); + while Instant::now() < kill_deadline && worker_runtime.is_alive(&worker_ref).await { sleep(Duration::from_millis(50)).await; } - - if fabro_proc::process_running(worker_pid) { - fabro_proc::sigkill(worker_pid); - - let kill_deadline = Instant::now() + Duration::from_secs(1); - while Instant::now() < kill_deadline && fabro_proc::process_running(worker_pid) { - sleep(Duration::from_millis(50)).await; - } - } } } @@ -2815,8 +2812,7 @@ fn clear_live_run_state(run: &mut ManagedRun) { run.event_tx = None; run.cancel_tx = None; run.cancel_token = None; - run.worker_pid = None; - run.worker_pgid = None; + run.worker_ref = None; } fn cleanup_worker_control_bus_for_run(state: &AppState, run_id: RunId) { @@ -2870,10 +2866,10 @@ fn release_run_answer_claim(state: &AppState, run_id: RunId, qid: &str) { } } -#[derive(Clone, Copy)] +#[derive(Clone)] struct LiveWorkerProcess { - run_id: RunId, - process_group_id: u32, + run_id: RunId, + worker_ref: WorkerRef, } fn failure_for_incomplete_run( @@ -2941,11 +2937,11 @@ fn live_worker_processes(state: &AppState) -> Vec { runs.iter() .filter_map(|(run_id, managed_run)| { managed_run - .worker_pgid - .or(managed_run.worker_pid) - .map(|process_group_id| LiveWorkerProcess { + .worker_ref + .clone() + .map(|worker_ref| LiveWorkerProcess { run_id: *run_id, - process_group_id, + worker_ref, }) }) .collect() @@ -2998,49 +2994,64 @@ async fn shutdown_active_workers_with_grace( state.begin_shutdown(); let workers = live_worker_processes(state.as_ref()); - #[cfg(unix)] - { - let process_groups = workers + join_all( + workers .iter() - .map(|worker| worker.process_group_id) - .collect::>(); + .map(|worker| state.worker_runtime.request_stop(&worker.worker_ref)), + ) + .await; - for process_group_id in &process_groups { - fabro_proc::sigterm_process_group(*process_group_id); - } + let survivors = poll_until_dead(state.as_ref(), &workers, grace, poll_interval).await; - let deadline = Instant::now() + grace; - while Instant::now() < deadline - && process_groups + if !survivors.is_empty() { + join_all( + survivors .iter() - .any(|process_group_id| fabro_proc::process_group_alive(*process_group_id)) + .map(|worker_ref| state.worker_runtime.force_stop(worker_ref)), + ) + .await; + // Wait for the kernel to reap the killed workers so callers can + // assume the processes are actually gone when shutdown returns. + let kill_deadline = Instant::now() + Duration::from_secs(1); + while Instant::now() < kill_deadline + && !alive_refs(state.as_ref(), &survivors).await.is_empty() { sleep(poll_interval).await; } - - let survivors = process_groups - .into_iter() - .filter(|process_group_id| fabro_proc::process_group_alive(*process_group_id)) - .collect::>(); - for process_group_id in &survivors { - fabro_proc::sigkill_process_group(*process_group_id); - } - if !survivors.is_empty() { - let kill_deadline = Instant::now() + Duration::from_secs(1); - while Instant::now() < kill_deadline - && survivors - .iter() - .any(|process_group_id| fabro_proc::process_group_alive(*process_group_id)) - { - sleep(poll_interval).await; - } - } } persist_shutdown_run_failures(state, &workers).await?; Ok(workers.len()) } +/// Poll until either the deadline expires or every worker is dead, returning +/// the set of workers still alive when polling stopped. +async fn poll_until_dead( + state: &AppState, + workers: &[LiveWorkerProcess], + grace: Duration, + poll_interval: Duration, +) -> Vec { + let refs: Vec = workers.iter().map(|w| w.worker_ref.clone()).collect(); + let deadline = Instant::now() + grace; + loop { + let alive = alive_refs(state, &refs).await; + if alive.is_empty() || Instant::now() >= deadline { + return alive; + } + sleep(poll_interval).await; + } +} + +async fn alive_refs(state: &AppState, refs: &[WorkerRef]) -> Vec { + let liveness = join_all(refs.iter().map(|r| state.worker_runtime.is_alive(r))).await; + refs.iter() + .zip(liveness) + .filter(|(_, alive)| *alive) + .map(|(r, _)| r.clone()) + .collect() +} + async fn persist_cancelled_run_status(state: &AppState, run_id: RunId) -> anyhow::Result<()> { let run_store = state.store.open_run(&run_id).await?; let run_state = run_store.state().await?; @@ -3167,8 +3178,7 @@ fn managed_run( checkpoint: None, cancel_tx: None, cancel_token: None, - worker_pid: None, - worker_pgid: None, + worker_ref: None, run_dir: Some(run_dir), execution_mode, } @@ -3373,7 +3383,10 @@ fn update_live_run_from_event(state: &AppState, run_id: RunId, event: &RunEvent) } } -async fn drain_worker_stderr(run_id: RunId, stderr: ChildStderr) -> anyhow::Result<()> { +async fn drain_worker_stderr( + run_id: RunId, + stderr: std::pin::Pin>, +) -> anyhow::Result<()> { let mut lines = BufReader::new(stderr).lines(); while let Some(line) = lines.next_line().await? { @@ -3383,10 +3396,32 @@ async fn drain_worker_stderr(run_id: RunId, stderr: ChildStderr) -> anyhow::Resu Ok(()) } +async fn fail_worker_launch( + state: &Arc, + run_store: &fabro_store::RunDatabase, + run_id: RunId, + err: anyhow::Error, +) { + tracing::error!(run_id = %run_id, error = %err, "Failed to spawn worker"); + let message = format!("Failed to spawn worker: {err}"); + let failure_event = workflow_event::Event::workflow_run_failed_from_error( + &WorkflowError::engine_with_anyhow("Failed to spawn worker", err), + fabro_types::RunTiming::default(), + FailureReason::LaunchFailed, + None, + None, + None, + None, + ); + let _ = workflow_event::append_event(run_store, &run_id, &failure_event).await; + fail_managed_run(state, run_id, FailureReason::LaunchFailed, message); + state.scheduler_notify.notify_one(); +} + async fn append_worker_exit_failure( run_store: &fabro_store::RunDatabase, run_id: RunId, - wait_status: &std::process::ExitStatus, + worker_exit: &WorkerExit, ) { let state = match run_store.state().await { Ok(state) => state, @@ -3403,7 +3438,10 @@ async fn append_worker_exit_failure( let (error, reason) = failure_for_incomplete_run( state.pending_control, - format!("Worker exited before emitting a terminal run event: {wait_status}"), + format!( + "Worker exited before emitting a terminal run event: {}", + worker_exit.detail + ), ); let failure_event = workflow_event::Event::workflow_run_failed_from_error( &error, @@ -3424,15 +3462,16 @@ async fn append_worker_exit_failure( clippy::disallowed_methods, reason = "Worker subprocess startup resolves Cargo's test binary env override when present." )] -fn worker_command( +fn worker_launch_spec( state: &AppState, run_id: RunId, mode: RunExecutionMode, run_dir: &std::path::Path, agent_fabro_tools_enabled: bool, -) -> anyhow::Result { +) -> anyhow::Result { let current_exe = std::env::current_exe().context("reading current executable path")?; - let exe = std::env::var_os(EnvVars::CARGO_BIN_EXE_FABRO).map_or(current_exe, PathBuf::from); + let executable = + std::env::var_os(EnvVars::CARGO_BIN_EXE_FABRO).map_or(current_exe, PathBuf::from); let storage_dir = state.server_storage_dir(); let runtime_directory = Storage::new(&storage_dir).runtime_directory(); let daemon = ServerDaemon::read(&runtime_directory)?.with_context(|| { @@ -3441,7 +3480,6 @@ fn worker_command( runtime_directory.record_path().display() ) })?; - let server_target = daemon.bind.to_target(); let scopes = if agent_fabro_tools_enabled { WorkerScopeSet::run_worker_with_agent_run_tools() } else { @@ -3449,46 +3487,26 @@ fn worker_command( }; let worker_token = issue_worker_token_with_scopes(state.worker_token_keys(), &run_id, scopes) .map_err(|_| anyhow::anyhow!("failed to sign worker token"))?; - let server_destination = resolved_log_destination(state)?; - let worker_stdout = match server_destination { - LogDestination::Stdout => Stdio::inherit(), - LogDestination::File => Stdio::null(), + let log_destination = resolved_log_destination(state)?; + let fabro_log = if (state.env_lookup)(EnvVars::FABRO_LOG).is_none() { + state.server_settings().server.logging.level.clone() + } else { + None }; - let mut cmd = Command::new(exe); - cmd.arg("__run-worker") - .arg("--server") - .arg(server_target) - .arg("--storage-dir") - .arg(&storage_dir) - .arg("--run-dir") - .arg(run_dir) - .arg("--run-id") - .arg(run_id.to_string()) - .arg("--mode") - .arg(worker_mode_arg(mode)) - .stdin(Stdio::null()) - .stdout(worker_stdout) - .stderr(Stdio::piped()); - apply_worker_env(&mut cmd); - if (state.env_lookup)(EnvVars::FABRO_LOG).is_none() { - if let Some(level) = state.server_settings().server.logging.level.as_deref() { - cmd.env(EnvVars::FABRO_LOG, level); - } - } - let value: &'static str = server_destination.into(); - cmd.env(EnvVars::FABRO_LOG_DESTINATION, value); - cmd.env(EnvVars::FABRO_CONFIG, state.active_config_path()); - cmd.env_remove(EnvVars::FABRO_WORKER_TOKEN); - cmd.env(EnvVars::FABRO_WORKER_TOKEN, worker_token); - if let Some(pem) = state.vault_secret(EnvVars::GITHUB_APP_PRIVATE_KEY) { - cmd.env(EnvVars::GITHUB_APP_PRIVATE_KEY, pem); - } - - #[cfg(unix)] - fabro_proc::pre_exec_setpgid(cmd.as_std_mut()); - - Ok(cmd) + Ok(WorkerLaunchSpec { + executable, + server_target: daemon.bind.to_target(), + storage_dir, + run_dir: run_dir.to_path_buf(), + run_id, + mode: worker_mode_arg(mode), + worker_token, + log_destination, + fabro_log, + active_config_path: state.active_config_path().to_path_buf(), + github_app_private_key: state.vault_secret(EnvVars::GITHUB_APP_PRIVATE_KEY), + }) } fn resolved_log_destination(state: &AppState) -> anyhow::Result { @@ -4081,8 +4099,8 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { let state_for_build = Arc::clone(&state); let run_dir_for_build = run_dir.clone(); - let build_cmd_result = spawn_blocking(move || { - worker_command( + let start_result = spawn_blocking(move || { + worker_launch_spec( state_for_build.as_ref(), run_id, execution_mode, @@ -4090,83 +4108,28 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { agent_fabro_tools_enabled, ) }) - .await; + .await + .context("worker_launch_spec task failed") + .and_then(|inner| inner); - let mut child = match build_cmd_result - .context("worker_command task failed") - .and_then(|inner| inner) - .and_then(|mut cmd| cmd.spawn().context("spawning run worker process")) - { - Ok(child) => child, + let launch_result = match start_result { + Ok(spec) => state.worker_runtime.start(spec).await, + Err(err) => Err(err), + }; + let started_worker = match launch_result { + Ok(worker) => worker, Err(err) => { - tracing::error!(run_id = %run_id, error = %err, "Failed to spawn worker"); - let message = format!("Failed to spawn worker: {err}"); - let failure_event = workflow_event::Event::workflow_run_failed_from_error( - &WorkflowError::engine_with_anyhow("Failed to spawn worker", err), - fabro_types::RunTiming::default(), - FailureReason::LaunchFailed, - None, - None, - None, - None, - ); - let _ = workflow_event::append_event(&run_store, &run_id, &failure_event).await; - fail_managed_run(&state, run_id, FailureReason::LaunchFailed, message); - state.scheduler_notify.notify_one(); + fail_worker_launch(&state, &run_store, run_id, err).await; return; } }; - - let Some(worker_pid) = child.id() else { - let message = "Worker process did not report a PID".to_string(); - tracing::error!(run_id = %run_id, "{message}"); - let _ = child.start_kill(); - let failure_event = workflow_event::Event::workflow_run_failed_from_error( - &WorkflowError::engine(message.clone()), - fabro_types::RunTiming::default(), - FailureReason::LaunchFailed, - None, - None, - None, - None, - ); - let _ = workflow_event::append_event(&run_store, &run_id, &failure_event).await; - fail_managed_run(&state, run_id, FailureReason::LaunchFailed, message); - state.scheduler_notify.notify_one(); - return; - }; + let worker_ref = started_worker.worker_ref.clone(); { let mut runs = state.runs.lock().expect("runs lock poisoned"); if let Some(managed_run) = runs.get_mut(&run_id) { - managed_run.worker_pid = Some(worker_pid); - managed_run.worker_pgid = Some(worker_pid); + managed_run.worker_ref = Some(worker_ref.clone()); managed_run.run_dir = Some(run_dir.clone()); - } - } - - let Some(stderr) = child.stderr.take() else { - let message = "Worker stderr pipe was unavailable".to_string(); - tracing::error!(run_id = %run_id, "{message}"); - let _ = child.start_kill(); - let failure_event = workflow_event::Event::workflow_run_failed_from_error( - &WorkflowError::engine(message.clone()), - fabro_types::RunTiming::default(), - FailureReason::LaunchFailed, - None, - None, - None, - None, - ); - let _ = workflow_event::append_event(&run_store, &run_id, &failure_event).await; - fail_managed_run(&state, run_id, FailureReason::LaunchFailed, message); - state.scheduler_notify.notify_one(); - return; - }; - - { - let mut runs = state.runs.lock().expect("runs lock poisoned"); - if let Some(managed_run) = runs.get_mut(&run_id) { managed_run.answer_transport = Some(RunAnswerTransport::Worker { run_id, bus: Arc::clone(&state.worker_control_bus), @@ -4174,14 +4137,14 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { } } - let stderr_task = tokio::spawn(drain_worker_stderr(run_id, stderr)); + let stderr_task = tokio::spawn(drain_worker_stderr(run_id, started_worker.stderr)); - let wait_status = match child.wait().await { - Ok(status) => status, + let worker_exit = match started_worker.wait.await { + Ok(exit) => exit, Err(err) => { tracing::error!(run_id = %run_id, error = %err, "Failed while waiting on worker"); let message = format!("Worker wait failed: {err}"); - let _ = child.start_kill(); + state.worker_runtime.force_stop(&worker_ref).await; let failure_event = workflow_event::Event::workflow_run_failed_from_error( &WorkflowError::engine_with_source("Worker wait failed", err), fabro_types::RunTiming::default(), @@ -4211,18 +4174,18 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { let superseded = { let runs = state.runs.lock().expect("runs lock poisoned"); runs.get(&run_id) - .is_some_and(|managed_run| managed_run.worker_pid != Some(worker_pid)) + .is_some_and(|managed_run| managed_run.worker_ref.as_ref() != Some(&worker_ref)) }; if superseded { tracing::info!( run_id = %run_id, - worker_pid, + worker_ref = ?worker_ref, "Skipping stale worker cleanup for superseded run execution" ); return; } - append_worker_exit_failure(&run_store, run_id, &wait_status).await; + append_worker_exit_failure(&run_store, run_id, &worker_exit).await; let final_state = match run_store.state().await { Ok(state) => state, @@ -4254,7 +4217,7 @@ async fn execute_run_subprocess(state: Arc, run_id: RunId) { if let Some(managed_run) = runs.get_mut(&run_id) { if final_state.status != managed_run.status { managed_run.status = final_state.status; - } else if !wait_status.success() { + } else if !worker_exit.success { managed_run.status = RunStatus::Failed { reason: FailureReason::Terminated, }; diff --git a/lib/crates/fabro-server/src/server/handler/lifecycle.rs b/lib/crates/fabro-server/src/server/handler/lifecycle.rs index d8a2589ac..d8735beff 100644 --- a/lib/crates/fabro-server/src/server/handler/lifecycle.rs +++ b/lib/crates/fabro-server/src/server/handler/lifecycle.rs @@ -18,6 +18,7 @@ use super::super::{ reject_if_archived, sleep, update_live_run_from_event, workflow_event, }; use super::runs::run_provenance; +use crate::worker_runtime::WorkerRef; pub(super) fn routes() -> Router> { Router::new() @@ -348,16 +349,17 @@ async fn deny_run( run_response(state.as_ref(), id, StatusCode::OK).await } -fn schedule_worker_kill(state: Arc, run_id: RunId, worker_pid: u32) { +fn schedule_worker_force_stop(state: Arc, run_id: RunId, worker_ref: WorkerRef) { tokio::spawn(async move { sleep(WORKER_CANCEL_GRACE).await; - let current_pid = { + let current_ref = { let runs = state.runs.lock().expect("runs lock poisoned"); - runs.get(&run_id).and_then(|run| run.worker_pid) + runs.get(&run_id).and_then(|run| run.worker_ref.clone()) }; - if current_pid == Some(worker_pid) && fabro_proc::process_group_alive(worker_pid) { - #[cfg(unix)] - fabro_proc::sigkill_process_group(worker_pid); + if current_ref.as_ref() == Some(&worker_ref) + && state.worker_runtime.is_alive(&worker_ref).await + { + state.worker_runtime.force_stop(&worker_ref).await; } }); } @@ -401,7 +403,7 @@ async fn cancel_run( managed_run.answer_transport.clone(), managed_run.cancel_token.clone(), managed_run.cancel_tx.take(), - managed_run.worker_pid, + managed_run.worker_ref.clone(), )) } _ => { @@ -412,7 +414,7 @@ async fn cancel_run( None => None, } }; - let Some((persist_cancelled_status, answer_transport, cancel_token, cancel_tx, worker_pid)) = + let Some((persist_cancelled_status, answer_transport, cancel_token, cancel_tx, worker_ref)) = cancel_target else { return unmanaged_cancel_response(state.as_ref(), id, actor, pending_control).await; @@ -448,10 +450,9 @@ async fn cancel_run( false }; if !delivered_control { - if let Some(worker_pid) = worker_pid { - #[cfg(unix)] - fabro_proc::sigterm(worker_pid); - schedule_worker_kill(Arc::clone(&state), id, worker_pid); + if let Some(worker_ref) = worker_ref { + state.worker_runtime.request_stop(&worker_ref).await; + schedule_worker_force_stop(Arc::clone(&state), id, worker_ref); } } diff --git a/lib/crates/fabro-server/src/server/tests.rs b/lib/crates/fabro-server/src/server/tests.rs index 866db6757..4543f9e2f 100644 --- a/lib/crates/fabro-server/src/server/tests.rs +++ b/lib/crates/fabro-server/src/server/tests.rs @@ -4,6 +4,7 @@ use std::os::unix::fs::PermissionsExt; use std::path::{Path, PathBuf}; #[cfg(unix)] use std::process::Stdio; +use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc as StdArc, Mutex as StdMutex}; use axum::body::Body; @@ -51,6 +52,9 @@ use crate::test_support::*; use crate::worker_control::{ LocalWorkerControlBus, WorkerControlBus, WorkerControlCursor, WorkerControlReceiver, }; +use crate::worker_runtime::{ + LocalWorkerRuntime, StartedWorker, WorkerLaunchSpec, WorkerRef, WorkerRuntime, +}; const MINIMAL_DOT: &str = r#"digraph Test { graph [goal="Test"] @@ -1944,6 +1948,7 @@ fn slack_app_state_with_secret_sources( sandbox_provider_registry: None, shutdown: tokio_util::sync::CancellationToken::new(), worker_control_bus: None, + worker_runtime: None, automation_materializer_override: None, }) .expect("slack test app state should build") @@ -2045,6 +2050,7 @@ fn slack_service_respects_disabled_server_config_even_with_vault_tokens() { sandbox_provider_registry: None, shutdown: tokio_util::sync::CancellationToken::new(), worker_control_bus: None, + worker_runtime: None, automation_materializer_override: None, }) .expect("slack disabled test app state should build"); @@ -2069,6 +2075,43 @@ fn worker_command_uses_null_stdin_and_token_env() { assert_worker_command_passes_token_only_by_env(&cmd); } +#[cfg(unix)] +#[test] +fn worker_command_sets_worker_args() { + let storage_dir = tempfile::tempdir().unwrap(); + let run_dir = storage_dir.path().join("run-scratch"); + let state = worker_command_test_state(storage_dir.path(), &["dev-token"], Some(TEST_DEV_TOKEN)); + let run_id = RunId::new(); + + let cmd = worker_command( + state.as_ref(), + run_id, + RunExecutionMode::Resume, + &run_dir, + false, + ) + .unwrap(); + + let args = cmd + .as_std() + .get_args() + .map(|arg| arg.to_string_lossy().into_owned()) + .collect::>(); + assert_eq!(args, vec![ + "__run-worker".to_string(), + "--server".to_string(), + "http://127.0.0.1:32276".to_string(), + "--storage-dir".to_string(), + storage_dir.path().display().to_string(), + "--run-dir".to_string(), + run_dir.display().to_string(), + "--run-id".to_string(), + run_id.to_string(), + "--mode".to_string(), + "resume".to_string(), + ]); +} + #[cfg(unix)] #[test] fn worker_command_default_token_omits_agent_run_tools_scope() { @@ -2368,6 +2411,7 @@ methods = ["dev-token"] sandbox_provider_registry: None, shutdown: tokio_util::sync::CancellationToken::new(), worker_control_bus: None, + worker_runtime: None, automation_materializer_override: None, }) else { panic!("build_app_state should require SESSION_SECRET") @@ -2497,10 +2541,27 @@ fn build_test_app_state_with_vault_path(vault_path: &Path) -> anyhow::Result WorkerRef { + WorkerRef::Local { pid } +} + +#[cfg(unix)] +fn worker_command( + state: &AppState, + run_id: RunId, + mode: RunExecutionMode, + run_dir: &Path, + agent_fabro_tools_enabled: bool, +) -> anyhow::Result { + let spec = worker_launch_spec(state, run_id, mode, run_dir, agent_fabro_tools_enabled)?; + Ok(LocalWorkerRuntime::command_for_spec(&spec)) +} + fn worker_command_test_state( storage_dir: &Path, methods: &[&str], @@ -2677,6 +2738,56 @@ fn worker_token_claims(cmd: &Command, state: &AppState) -> crate::worker_token:: .claims } +#[derive(Default)] +struct RecordingWorkerRuntime { + requested: StdMutex>, + forced: StdMutex>, + alive: AtomicBool, +} + +impl RecordingWorkerRuntime { + fn requested_refs(&self) -> Vec { + self.requested + .lock() + .expect("requested lock poisoned") + .clone() + } + + fn forced_refs(&self) -> Vec { + self.forced.lock().expect("forced lock poisoned").clone() + } + + fn set_alive(&self, alive: bool) { + self.alive.store(alive, Ordering::Relaxed); + } +} + +#[async_trait::async_trait] +impl WorkerRuntime for RecordingWorkerRuntime { + async fn start(&self, _spec: WorkerLaunchSpec) -> anyhow::Result { + anyhow::bail!("recording runtime does not start workers") + } + + async fn request_stop(&self, worker_ref: &WorkerRef) { + self.requested + .lock() + .expect("requested lock poisoned") + .push(worker_ref.clone()); + } + + async fn force_stop(&self, worker_ref: &WorkerRef) { + self.forced + .lock() + .expect("forced lock poisoned") + .push(worker_ref.clone()); + self.alive.store(false, Ordering::Relaxed); + } + + async fn is_alive(&self, _worker_ref: &WorkerRef) -> bool { + self.alive.load(Ordering::Relaxed) + } +} + async fn worker_transport_with_receiver( run_id: RunId, ) -> (RunAnswerTransport, WorkerControlReceiver) { @@ -5775,6 +5886,7 @@ fn create_github_token_app_state_with_env_lookup_and_llm_catalog_settings( sandbox_provider_registry: None, shutdown: tokio_util::sync::CancellationToken::new(), worker_control_bus: None, + worker_runtime: None, automation_materializer_override: None, }; let state = build_app_state(config).expect("test app state should build"); @@ -13442,7 +13554,7 @@ async fn cancel_run_overwrites_pending_pause_request() { let mut runs = state.runs.lock().expect("runs lock poisoned"); let managed_run = runs.get_mut(&run_id).expect("run should exist"); managed_run.status = RunStatus::Running; - managed_run.worker_pid = Some(u32::MAX); + managed_run.worker_ref = Some(test_worker_ref(u32::MAX)); } append_control_request(state.as_ref(), run_id, RunControlAction::Pause, None) .await @@ -13464,6 +13576,38 @@ async fn cancel_run_overwrites_pending_pause_request() { ); } +#[tokio::test] +async fn cancel_run_requests_worker_runtime_stop_when_control_unavailable() { + let runtime = StdArc::new(RecordingWorkerRuntime::default()); + let state = TestAppStateBuilder::new() + .worker_runtime(runtime.clone()) + .build(); + let app = crate::test_support::build_test_router(Arc::clone(&state)); + let run_id = create_and_start_run(&app, MINIMAL_DOT) + .await + .parse::() + .unwrap(); + let worker_ref = test_worker_ref(u32::MAX); + + { + let mut runs = state.runs.lock().expect("runs lock poisoned"); + let managed_run = runs.get_mut(&run_id).expect("run should exist"); + managed_run.status = RunStatus::Running; + managed_run.answer_transport = None; + managed_run.worker_ref = Some(worker_ref.clone()); + } + + let req = Request::builder() + .method("POST") + .uri(api(&format!("/runs/{run_id}/cancel"))) + .body(Body::empty()) + .unwrap(); + let response = app.oneshot(req).await.unwrap(); + assert_status!(response, StatusCode::OK).await; + + assert_eq!(runtime.requested_refs(), vec![worker_ref]); +} + #[tokio::test] async fn pause_run_rejects_when_control_is_already_pending() { let state = test_app_state(); @@ -13475,7 +13619,7 @@ async fn pause_run_rejects_when_control_is_already_pending() { let mut runs = state.runs.lock().expect("runs lock poisoned"); let managed_run = runs.get_mut(&run_id).expect("run should exist"); managed_run.status = RunStatus::Running; - managed_run.worker_pid = Some(u32::MAX); + managed_run.worker_ref = Some(test_worker_ref(u32::MAX)); } append_control_request(state.as_ref(), run_id, RunControlAction::Cancel, None) .await @@ -13508,7 +13652,7 @@ async fn pause_run_sets_pending_control_on_board_response() { let mut runs = state.runs.lock().expect("runs lock poisoned"); let managed_run = runs.get_mut(&run_id).expect("run should exist"); managed_run.status = RunStatus::Running; - managed_run.worker_pid = Some(u32::MAX); + managed_run.worker_ref = Some(test_worker_ref(u32::MAX)); managed_run.answer_transport = Some(transport); } @@ -13594,7 +13738,7 @@ async fn pause_run_immediately_pauses_blocked_run() { managed_run.status = RunStatus::Blocked { blocked_reason: BlockedReason::HumanInputRequired, }; - managed_run.worker_pid = Some(u32::MAX); + managed_run.worker_ref = Some(test_worker_ref(u32::MAX)); } let req = Request::builder() @@ -13630,7 +13774,7 @@ async fn unpause_run_sets_pending_control() { let mut runs = state.runs.lock().expect("runs lock poisoned"); let managed_run = runs.get_mut(&run_id).expect("run should exist"); managed_run.status = RunStatus::Paused { prior_block: None }; - managed_run.worker_pid = Some(u32::MAX); + managed_run.worker_ref = Some(test_worker_ref(u32::MAX)); managed_run.answer_transport = Some(transport); } @@ -13705,7 +13849,7 @@ async fn unpause_run_returns_blocked_when_human_gate_is_still_unresolved() { managed_run.status = RunStatus::Paused { prior_block: Some(BlockedReason::HumanInputRequired), }; - managed_run.worker_pid = Some(u32::MAX); + managed_run.worker_ref = Some(test_worker_ref(u32::MAX)); } let req = Request::builder() @@ -13799,6 +13943,55 @@ async fn startup_reconciliation_marks_inflight_runs_terminal() { assert_eq!(run_3.pending_control, None); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn shutdown_active_workers_uses_worker_runtime_for_live_refs() { + let runtime = StdArc::new(RecordingWorkerRuntime::default()); + runtime.set_alive(true); + let state = TestAppStateBuilder::new() + .worker_runtime(runtime.clone()) + .build(); + let worker_refs = [test_worker_ref(u32::MAX - 1), test_worker_ref(u32::MAX)]; + let run_ids = [RunId::new(), RunId::new()]; + let temp_dir = tempfile::tempdir().unwrap(); + + for (run_id, worker_ref) in run_ids.iter().zip(worker_refs.iter()) { + create_durable_run_with_events(&state, *run_id, &[ + workflow_event::Event::RunSubmitted { + definition_blob: None, + }, + workflow_event::Event::RunStarting, + workflow_event::Event::RunRunning, + ]) + .await; + + let mut run = managed_run( + String::new(), + RunStatus::Running, + chrono::Utc::now(), + temp_dir.path().join(run_id.to_string()), + RunExecutionMode::Start, + ); + run.worker_ref = Some(worker_ref.clone()); + state + .runs + .lock() + .expect("runs lock poisoned") + .insert(*run_id, run); + } + + let terminated = shutdown_active_workers_with_grace( + &state, + Duration::from_millis(0), + Duration::from_millis(1), + ) + .await + .unwrap(); + + assert_eq!(terminated, 2); + assert_eq!(runtime.requested_refs().len(), 2); + assert_eq!(runtime.forced_refs().len(), 2); +} + #[cfg(unix)] #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn shutdown_active_workers_terminates_process_groups() { @@ -13824,7 +14017,7 @@ async fn shutdown_active_workers_terminates_process_groups() { .stderr(Stdio::null()); fabro_proc::pre_exec_setpgid(child.as_std_mut()); let mut child = child.spawn().unwrap(); - let worker_pid = child.id().expect("worker pid should be available"); + let worker_process_id = child.id().expect("worker pid should be available"); { let mut runs = state.runs.lock().expect("runs lock poisoned"); @@ -13835,8 +14028,7 @@ async fn shutdown_active_workers_terminates_process_groups() { temp_dir.path().join(run_id.to_string()), RunExecutionMode::Start, ); - run.worker_pid = Some(worker_pid); - run.worker_pgid = Some(worker_pid); + run.worker_ref = Some(test_worker_ref(worker_process_id)); runs.insert(run_id, run); } @@ -13854,7 +14046,7 @@ async fn shutdown_active_workers_terminates_process_groups() { .expect("worker should exit after shutdown") .expect("wait should succeed"); assert!(!exit_status.success()); - assert!(!fabro_proc::process_group_alive(worker_pid)); + assert!(!fabro_proc::process_group_alive(worker_process_id)); let run_state = state .store diff --git a/lib/crates/fabro-server/src/test_support.rs b/lib/crates/fabro-server/src/test_support.rs index 9267375b3..de5003629 100644 --- a/lib/crates/fabro-server/src/test_support.rs +++ b/lib/crates/fabro-server/src/test_support.rs @@ -39,6 +39,8 @@ use crate::server::{ RouterOptions, build_app_state, process_env_var, }; use crate::server_secrets::ServerSecrets; +#[cfg(test)] +use crate::worker_runtime::WorkerRuntime; pub const TEST_DEV_TOKEN: &str = "fabro_dev_abababababababababababababababababababababababababababababababab"; @@ -73,25 +75,29 @@ pub struct TestAppStateBuilder { env_lookup: EnvLookup, llm_catalog_settings: LlmCatalogSettings, automation_materializer: Option>, + #[cfg(test)] + worker_runtime: Option>, } impl Default for TestAppStateBuilder { fn default() -> Self { Self { - server_settings: default_test_server_settings(), - manifest_run_defaults: RunLayer::default(), - max_concurrent_runs: 5, - registry_factory_override: None, - sandbox_provider_registry: None, - store_bundle: None, - vault_path: None, - vault_entries: Vec::new(), - server_env_path: None, - active_config_path: None, - server_secret_env: HashMap::new(), - env_lookup: default_env_lookup(), - llm_catalog_settings: LlmCatalogSettings::default(), - automation_materializer: None, + server_settings: default_test_server_settings(), + manifest_run_defaults: RunLayer::default(), + max_concurrent_runs: 5, + registry_factory_override: None, + sandbox_provider_registry: None, + store_bundle: None, + vault_path: None, + vault_entries: Vec::new(), + server_env_path: None, + active_config_path: None, + server_secret_env: HashMap::new(), + env_lookup: default_env_lookup(), + llm_catalog_settings: LlmCatalogSettings::default(), + automation_materializer: None, + #[cfg(test)] + worker_runtime: None, } } } @@ -153,6 +159,12 @@ impl TestAppStateBuilder { self } + #[cfg(test)] + pub(crate) fn worker_runtime(mut self, worker_runtime: Arc) -> Self { + self.worker_runtime = Some(worker_runtime); + self + } + pub fn provider_base_url( mut self, provider: impl Into, @@ -250,6 +262,8 @@ impl TestAppStateBuilder { shutdown: CancellationToken::new(), #[cfg(test)] worker_control_bus: None, + #[cfg(test)] + worker_runtime: self.worker_runtime, automation_materializer_override: self.automation_materializer, }) } diff --git a/lib/crates/fabro-server/src/worker_runtime.rs b/lib/crates/fabro-server/src/worker_runtime.rs new file mode 100644 index 000000000..90f595a5e --- /dev/null +++ b/lib/crates/fabro-server/src/worker_runtime.rs @@ -0,0 +1,166 @@ +use std::path::PathBuf; +use std::pin::Pin; +use std::process::Stdio; + +use anyhow::{Context, Result}; +use async_trait::async_trait; +use fabro_static::EnvVars; +use fabro_types::RunId; +use fabro_types::settings::server::LogDestination; +use futures_util::future::BoxFuture; +use tokio::io::AsyncRead; +use tokio::process::Command; + +use crate::spawn_env::apply_worker_env; + +#[async_trait] +pub(crate) trait WorkerRuntime: Send + Sync { + async fn start(&self, spec: WorkerLaunchSpec) -> Result; + async fn request_stop(&self, worker_ref: &WorkerRef); + async fn force_stop(&self, worker_ref: &WorkerRef); + async fn is_alive(&self, worker_ref: &WorkerRef) -> bool; +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub(crate) enum WorkerRef { + /// A worker running as a local subprocess. `pre_exec_setpgid` ensures the + /// child is the leader of its own process group with `pgid == pid`, so a + /// single PID identifies both the process and its group. + Local { pid: u32 }, +} + +pub(crate) struct WorkerLaunchSpec { + pub(crate) executable: PathBuf, + pub(crate) server_target: String, + pub(crate) storage_dir: PathBuf, + pub(crate) run_dir: PathBuf, + pub(crate) run_id: RunId, + pub(crate) mode: &'static str, + pub(crate) worker_token: String, + pub(crate) log_destination: LogDestination, + pub(crate) fabro_log: Option, + pub(crate) active_config_path: PathBuf, + pub(crate) github_app_private_key: Option, +} + +pub(crate) struct StartedWorker { + pub(crate) worker_ref: WorkerRef, + pub(crate) stderr: Pin>, + pub(crate) wait: BoxFuture<'static, Result>, +} + +#[derive(Debug)] +pub(crate) struct WorkerExit { + pub(crate) success: bool, + pub(crate) detail: String, +} + +#[derive(Default)] +pub(crate) struct LocalWorkerRuntime; + +impl LocalWorkerRuntime { + pub(crate) fn new() -> Self { + Self + } + + pub(crate) fn command_for_spec(spec: &WorkerLaunchSpec) -> Command { + let worker_stdout = match spec.log_destination { + LogDestination::Stdout => Stdio::inherit(), + LogDestination::File => Stdio::null(), + }; + let log_destination_env: &'static str = spec.log_destination.into(); + + let mut cmd = Command::new(&spec.executable); + cmd.arg("__run-worker") + .arg("--server") + .arg(&spec.server_target) + .arg("--storage-dir") + .arg(&spec.storage_dir) + .arg("--run-dir") + .arg(&spec.run_dir) + .arg("--run-id") + .arg(spec.run_id.to_string()) + .arg("--mode") + .arg(spec.mode) + .stdin(Stdio::null()) + .stdout(worker_stdout) + .stderr(Stdio::piped()); + + apply_worker_env(&mut cmd); + if let Some(level) = spec.fabro_log.as_deref() { + cmd.env(EnvVars::FABRO_LOG, level); + } + cmd.env(EnvVars::FABRO_LOG_DESTINATION, log_destination_env); + cmd.env(EnvVars::FABRO_CONFIG, &spec.active_config_path); + cmd.env_remove(EnvVars::FABRO_WORKER_TOKEN); + cmd.env(EnvVars::FABRO_WORKER_TOKEN, &spec.worker_token); + if let Some(pem) = spec.github_app_private_key.as_deref() { + cmd.env(EnvVars::GITHUB_APP_PRIVATE_KEY, pem); + } + + #[cfg(unix)] + fabro_proc::pre_exec_setpgid(cmd.as_std_mut()); + + cmd + } +} + +#[async_trait] +impl WorkerRuntime for LocalWorkerRuntime { + async fn start(&self, spec: WorkerLaunchSpec) -> Result { + let mut child = Self::command_for_spec(&spec) + .spawn() + .context("spawning run worker process")?; + + let Some(pid) = child.id() else { + let _ = child.start_kill(); + anyhow::bail!("worker process did not report a PID"); + }; + let Some(stderr) = child.stderr.take() else { + let _ = child.start_kill(); + anyhow::bail!("worker child stderr should be piped"); + }; + let stderr: Pin> = Box::pin(stderr); + let wait: BoxFuture<'static, Result> = Box::pin(async move { + let status = child.wait().await.context("worker wait failed")?; + Ok(WorkerExit { + success: status.success(), + detail: status.to_string(), + }) + }); + + Ok(StartedWorker { + worker_ref: WorkerRef::Local { pid }, + stderr, + wait, + }) + } + + async fn request_stop(&self, worker_ref: &WorkerRef) { + let WorkerRef::Local { pid } = worker_ref; + #[cfg(unix)] + fabro_proc::sigterm_process_group(*pid); + #[cfg(not(unix))] + let _ = pid; + } + + async fn force_stop(&self, worker_ref: &WorkerRef) { + let WorkerRef::Local { pid } = worker_ref; + #[cfg(unix)] + fabro_proc::sigkill_process_group(*pid); + #[cfg(not(unix))] + let _ = pid; + } + + async fn is_alive(&self, worker_ref: &WorkerRef) -> bool { + let WorkerRef::Local { pid } = worker_ref; + #[cfg(unix)] + { + fabro_proc::process_group_alive(*pid) + } + #[cfg(not(unix))] + { + fabro_proc::process_running(*pid) + } + } +}