diff --git a/.alfonso/release-notes/next/lsp-exit-sweep.md b/.alfonso/release-notes/next/lsp-exit-sweep.md new file mode 100644 index 0000000000..1cc3ff9177 --- /dev/null +++ b/.alfonso/release-notes/next/lsp-exit-sweep.md @@ -0,0 +1 @@ +- When the daemon exits, it now stops every running language server within its shared shutdown budget and reports the true count, instead of leaving some to per-project background threads. diff --git a/benchmarks/aft-search/slice-descriptors/land-lsp-exit-sweep.json b/benchmarks/aft-search/slice-descriptors/land-lsp-exit-sweep.json new file mode 100644 index 0000000000..ebd574711e --- /dev/null +++ b/benchmarks/aft-search/slice-descriptors/land-lsp-exit-sweep.json @@ -0,0 +1,8 @@ +{ + "targeted_mechanism": "none", + "kind": "harness", + "slice_class": "non_ranking", + "fixtures": [ + "harness-goldens" + ] +} diff --git a/crates/aft/src/lsp/manager.rs b/crates/aft/src/lsp/manager.rs index f85dc1ec97..63650195fa 100644 --- a/crates/aft/src/lsp/manager.rs +++ b/crates/aft/src/lsp/manager.rs @@ -4333,6 +4333,17 @@ impl LspManager { self.take_all_clients() } + /// Refuse new server starts like [`Self::quiesce`], but keep the running + /// clients in this manager. Used when the module's connection has ended + /// and the exit phase that follows takes every manager's clients and shuts + /// them down together under one shared budget. Detaching them here would + /// hand each root's servers to its own background reap thread, which can + /// stop them before the exit phase sees them. + pub(crate) fn close_start_admission(&self) { + self.quiesced + .store(true, std::sync::atomic::Ordering::Release); + } + pub(crate) fn resume(&self) { self.quiesced .store(false, std::sync::atomic::Ordering::Release); diff --git a/crates/aft/src/subc/mod.rs b/crates/aft/src/subc/mod.rs index 25d369081a..fb50a52bf1 100644 --- a/crates/aft/src/subc/mod.rs +++ b/crates/aft/src/subc/mod.rs @@ -1865,10 +1865,32 @@ fn allocator_pressure_relief_after_idle_sweep( } } +/// Who stops a root's running language servers once the root is quiesced. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum QuiescedLsp { + /// Detach the clients and shut them down on a background thread. Right for + /// a route that closes while the module keeps running: nothing else would + /// stop those servers. + ReapInBackground, + /// Refuse new starts but leave the clients in the manager, because the + /// caller's exit phase shuts down every manager's servers together under + /// one shared budget right after the module loop returns. + LeaveForExitSweep, +} + fn quiesce_unbound_root( root_id: &ProjectRootId, live_roots: &mut HashMap, executor: &Arc, +) { + quiesce_unbound_root_with(root_id, live_roots, executor, QuiescedLsp::ReapInBackground); +} + +fn quiesce_unbound_root_with( + root_id: &ProjectRootId, + live_roots: &mut HashMap, + executor: &Arc, + lsp: QuiescedLsp, ) { let Some(meta) = live_roots.get_mut(root_id) else { return; @@ -1887,10 +1909,16 @@ fn quiesce_unbound_root( ctx.stop_watcher_runtime(); ctx.invalidate_artifacts_after_watcher_gap(); } - // Detach under the root's manager lock before allowing a rebind. The - // reaper owns only the old clients, never servers a new route starts. - let clients = ctx.lsp().quiesce(); - crate::lsp::manager::LspManager::spawn_idle_lsp_reap(clients); + match lsp { + QuiescedLsp::ReapInBackground => { + // Detach under the root's manager lock before allowing a + // rebind. The reaper owns only the old clients, never servers + // a new route starts. + let clients = ctx.lsp().quiesce(); + crate::lsp::manager::LspManager::spawn_idle_lsp_reap(clients); + } + QuiescedLsp::LeaveForExitSweep => ctx.lsp().close_start_admission(), + } ctx.bash_background() .replace_live_delivery_sessions(HashSet::new()); } @@ -1931,6 +1959,7 @@ fn quiesce_connection_roots( route_bash_cancels: &mut HashMap, active_tool_calls: &ActiveToolCalls, executor: &Arc, + lsp: QuiescedLsp, ) { cancel_all_active_tool_calls(active_tool_calls, executor, "connection teardown"); for cancel in route_bash_cancels.values() { @@ -1950,7 +1979,7 @@ fn quiesce_connection_roots( // root activity after the loop-owned route tables disappear. for root_id in roots { if live_roots.contains_key(&root_id) { - quiesce_unbound_root(&root_id, live_roots, executor); + quiesce_unbound_root_with(&root_id, live_roots, executor, lsp); } else if let Some(ctx) = executor.actor_context(&root_id) { ctx.mark_subc_unbound(); executor.cancel_queued_maintenance(&root_id); @@ -3949,6 +3978,8 @@ fn run_subc_mode_inner( tool_response_body_limit, lifecycle_probe, &storage_dir, + // The exit phase below takes and shuts down every root's servers. + QuiescedLsp::LeaveForExitSweep, #[cfg(test)] None, #[cfg(test)] @@ -3980,6 +4011,9 @@ fn run_subc_mode_inner( }); let mut clients = Vec::new(); for actor_ctx in &actor_contexts { + // Connection teardown left every root's running servers in its + // manager (`QuiescedLsp::LeaveForExitSweep`), so they are all stopped + // here under the one shared budget below. clients.extend(actor_ctx.lsp().take_all_clients()); actor_ctx.bash_background().detach(); } @@ -4638,6 +4672,7 @@ async fn run_module_loop( tool_response_body_limit: usize, lifecycle_probe: Option, storage_dir: &Path, + connection_end_lsp: QuiescedLsp, #[cfg(test)] standing_actor_override: Option>, #[cfg(test)] mut watchdog_test: Option, ) -> Result @@ -6074,6 +6109,7 @@ where &mut route_bash_cancels, &active_tool_calls, &executor, + connection_end_lsp, ); fleet_status_client.set_route_live(false); @@ -12418,6 +12454,7 @@ mod tests { usize::MAX, None, &fixture_path.join("storage"), + QuiescedLsp::ReapInBackground, Some(actor), None, )) @@ -12581,6 +12618,7 @@ mod tests { usize::MAX, None, &fixture_path.join("storage"), + QuiescedLsp::ReapInBackground, Some(actor), None, )) @@ -15014,6 +15052,7 @@ mod tests { &mut route_bash_cancels, &active_tool_calls, &executor, + QuiescedLsp::ReapInBackground, ); assert!(live_roots[&root].unbound_quiesced); assert!(!live_roots[&root].maintenance_pending); @@ -15166,6 +15205,142 @@ mod tests { ctx.lsp().shutdown_all(); } + /// A root with one running fake rust-analyzer that ignores Shutdown, bound + /// through `live_roots` as the module loop would hold it. + fn root_with_running_fake_server( + name: &str, + ) -> ( + tempfile::TempDir, + ProjectRootId, + Arc, + Arc, + HashMap, + u32, + ) { + use crate::lsp::registry::ServerKind; + + let (root_dir, root) = test_root(name); + std::fs::write( + root.as_path().join("Cargo.toml"), + "[package]\nname='quiesce'\nversion='0.1.0'\n", + ) + .unwrap(); + let file = root.as_path().join("lib.rs"); + std::fs::write(&file, "pub fn f() {}\n").unwrap(); + let ctx = test_ctx(); + ctx.mark_subc_bound(); + ctx.update_config(|config| config.project_root = Some(root.as_path().to_path_buf())); + let config = ctx.config().as_ref().clone(); + ctx.lsp().override_binary( + ServerKind::Rust, + crate::fake_lsp_test_helper::fake_server_binary(), + ); + ctx.lsp() + .set_extra_env("AFT_FAKE_LSP_IGNORE_SHUTDOWN".into(), "1".into()); + assert_eq!(ctx.lsp().ensure_server_for_file(&file, &config).len(), 1); + assert_eq!(ctx.lsp().active_client_count(), 1); + let pids = ctx.lsp_child_registry().pids(); + assert_eq!(pids.len(), 1, "one tracked server child: {pids:?}"); + let executor = Arc::new(Executor::new()); + assert!(executor.register_actor(root.clone(), Arc::clone(&ctx))); + let live_roots = HashMap::from([(root.clone(), RootMeta::new(Instant::now()))]); + (root_dir, root, ctx, executor, live_roots, pids[0]) + } + + fn quiesce_whole_connection( + live_roots: &mut HashMap, + executor: &Arc, + lsp: QuiescedLsp, + ) { + quiesce_connection_roots( + live_roots, + &mut HashMap::new(), + &mut HashMap::new(), + &mut HashMap::new(), + &mut HashMap::new(), + &mut HashMap::new(), + &Arc::new(StdMutex::new(HashMap::new())), + executor, + lsp, + ); + } + + /// When the module's connection ends, the exit phase that follows must be + /// the one owner of every running server, so it can stop them all under + /// its single shared budget and report each one. Connection teardown used + /// to detach them to per-root background reap threads, which could stop + /// every server first and leave the exit phase reporting `servers=0`. + #[test] + fn connection_end_quiesce_leaves_running_servers_for_the_exit_sweep() { + use crate::lsp::manager::ServerAttemptResult; + + let (_root_dir, root, ctx, executor, mut live_roots, pid) = + root_with_running_fake_server("connection-end-exit-sweep"); + quiesce_whole_connection(&mut live_roots, &executor, QuiescedLsp::LeaveForExitSweep); + assert!(live_roots[&root].unbound_quiesced); + assert_eq!( + ctx.lsp().active_client_count(), + 1, + "the exit sweep, not a background reap, must own the running server" + ); + assert!(crate::bash_background::process::is_process_alive(pid)); + + // The exit phase of `run_subc_mode_inner`: take every client and stop + // them under one budget, killing the server that ignores Shutdown. + let clients = ctx.lsp().take_all_clients(); + assert_eq!(clients.len(), 1); + let outcome = crate::lsp::manager::LspManager::shutdown_at_process_exit( + clients, + ctx.lsp_child_registry(), + ); + assert_eq!(outcome.graceful + outcome.forced, 1, "{outcome:?}"); + assert_eq!(outcome.unreaped, 0, "{outcome:?}"); + // The forced kill does not wait; its worker thread reaps the child + // right after. Hang guard only. + let deadline = Instant::now() + Duration::from_secs(30); + while crate::bash_background::process::is_process_alive(pid) { + assert!( + Instant::now() < deadline, + "exit sweep left server {pid} running" + ); + std::thread::sleep(Duration::from_millis(20)); + } + + // Admission stays closed: a late request (an inspect job that outlived + // its route) cannot start a replacement server after the sweep. + let config = ctx.config().as_ref().clone(); + let refused = ctx + .lsp() + .ensure_server_for_file_detailed(&root.as_path().join("lib.rs"), &config); + assert!(refused + .attempts + .iter() + .any(|attempt| matches!(attempt.result, ServerAttemptResult::ProjectClosed))); + assert_eq!(ctx.lsp().active_client_count(), 0); + assert!(ctx.lsp_child_registry().pids().is_empty()); + } + + /// A teardown that the process outlives has no exit sweep behind it, so + /// its servers must still be stopped by the background reap. + #[test] + fn connection_end_quiesce_without_exit_sweep_reaps_servers_in_background() { + let (_root_dir, _root, ctx, executor, mut live_roots, pid) = + root_with_running_fake_server("connection-end-background-reap"); + quiesce_whole_connection(&mut live_roots, &executor, QuiescedLsp::ReapInBackground); + assert_eq!(ctx.lsp().active_client_count(), 0); + // Hang guard only: the reap waits 25 ms for Shutdown, then terminates. + let deadline = Instant::now() + Duration::from_secs(30); + while crate::bash_background::process::is_process_alive(pid) + || ctx.lsp_child_registry().pids().contains(&pid) + { + assert!( + Instant::now() < deadline, + "background reap left server {pid} running" + ); + std::thread::sleep(Duration::from_millis(20)); + } + } + #[test] fn unbound_root_quiesces_maintenance_without_removing_actor() { let (_root_dir, root) = test_root("unbound-root-quiesce"); diff --git a/crates/aft/src/subc/remote_policy_tests.rs b/crates/aft/src/subc/remote_policy_tests.rs index 64f7c1a3af..3843677211 100644 --- a/crates/aft/src/subc/remote_policy_tests.rs +++ b/crates/aft/src/subc/remote_policy_tests.rs @@ -1761,6 +1761,7 @@ async fn exec_remote_module_loop_captures_its_authenticated_daemon_endpoint() { 1024 * 1024, None, dir.path(), + QuiescedLsp::ReapInBackground, None, None, ); diff --git a/crates/aft/src/subc/stall_watchdog.rs b/crates/aft/src/subc/stall_watchdog.rs index e745d83494..d3cf7be97e 100644 --- a/crates/aft/src/subc/stall_watchdog.rs +++ b/crates/aft/src/subc/stall_watchdog.rs @@ -1033,6 +1033,7 @@ mod tests { usize::MAX, None, &path, + QuiescedLsp::ReapInBackground, None, Some(hooks), )) diff --git a/crates/aft/tests/integration/subc_detach_test.rs b/crates/aft/tests/integration/subc_detach_test.rs index 211fdf7888..cba79879c4 100644 --- a/crates/aft/tests/integration/subc_detach_test.rs +++ b/crates/aft/tests/integration/subc_detach_test.rs @@ -758,6 +758,13 @@ fn drain_with_live_lsp_servers_and_writer( shutdowns[0].contains(" budget_ms=1500 "), "all servers must be admitted under one 1500 ms shutdown budget: {log}" ); + // Connection teardown leaves every running server to the exit phase, + // so the phase takes all of them however late the runner schedules it. + assert!( + shutdowns[0].contains(&format!(" servers={server_count} ")), + "the exit phase must take every root's running server: {}", + shutdowns[0] + ); let durable_path = data_home.path().join("aft").join("logs").join(format!("aft-{}.log", module.child.id())); let durable = std::fs::read_to_string(&durable_path) .unwrap_or_else(|error| panic!("{}: {error}", durable_path.display())); @@ -769,9 +776,8 @@ fn drain_with_live_lsp_servers_and_writer( if active_ort { assert!(durable.contains("phase=ort_quiesce active=1 budget_ms=0"), "{}", log_tail(&durable)); } - // Draining roots may already have stopped some servers before the final - // sweep, so the summary's split depends on timing. What must hold is one - // summary, a bounded exit and no orphans. + // One summary, covering every server: the exit phase owns them all, so + // its count does not depend on timing. The graceful/killed split does. let summaries: Vec<&str> = log .lines() .filter(|line| line.contains("lsp shutdown_all: servers=")) @@ -791,7 +797,11 @@ fn drain_with_live_lsp_servers_and_writer( .unwrap_or_else(|| panic!("missing {field} in shutdown summary: {}", summaries[0])) }; let servers = summary_field("servers="); - assert!(servers >= 3, "shared-deadline proof needs at least three shutdowns: {}", summaries[0]); + assert_eq!( + servers, server_count as u128, + "shared-deadline proof needs every server stopped by the exit phase: {}", + summaries[0] + ); let lsp_elapsed_ms = summary_field("elapsed_ms="); let lsp_ceiling_ms = 2 * aft::lsp::manager::LSP_SHUTDOWN_ALL_BUDGET.as_millis(); // Measure AFT's phase, not the parent thread's wakeup or process exit.