use std::ffi::OsString;
use std::io;
use std::process::{Child, Command, Stdio};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use std::thread::{self, JoinHandle};
use std::time::Duration;
pub type SpawnFn = dyn Fn() -> io::Result<Child> + Send + Sync + 'static;
const POLL_INTERVAL: Duration = Duration::from_millis(40);
const MIN_BACKOFF: Duration = Duration::from_millis(50);
const MAX_BACKOFF: Duration = Duration::from_secs(2);
struct SupState {
child: Option<Child>,
last_pid: Option<u32>,
restarts: u64,
}
pub type OnSpawnFn = dyn FnMut(&mut Child) + Send + 'static;
struct Shared {
spawn: Box<SpawnFn>,
on_spawn: Mutex<Box<OnSpawnFn>>,
state: Mutex<SupState>,
shutdown: AtomicBool,
suspended: AtomicBool,
}
pub struct Supervisor {
shared: Arc<Shared>,
monitor: Option<JoinHandle<()>>,
}
impl Supervisor {
pub fn start<F>(spawn: F) -> io::Result<Self>
where
F: Fn() -> io::Result<Child> + Send + Sync + 'static,
{
Self::start_with_hook(spawn, |_child: &mut Child| {})
}
pub fn start_with_hook<F, H>(spawn: F, on_spawn: H) -> io::Result<Self>
where
F: Fn() -> io::Result<Child> + Send + Sync + 'static,
H: FnMut(&mut Child) + Send + 'static,
{
let shared = Arc::new(Shared {
spawn: Box::new(spawn),
on_spawn: Mutex::new(Box::new(on_spawn)),
state: Mutex::new(SupState {
child: None,
last_pid: None,
restarts: 0,
}),
shutdown: AtomicBool::new(false),
suspended: AtomicBool::new(false),
});
let mut first = (shared.spawn)()?;
invoke_on_spawn(&shared, &mut first);
{
let mut st = lock(&shared.state);
st.last_pid = Some(first.id());
st.child = Some(first);
}
let mon_shared = Arc::clone(&shared);
let monitor = thread::Builder::new()
.name("tf-ra-supervisor".into())
.spawn(move || monitor_loop(mon_shared))
.expect("spawn tf-ra-supervisor thread");
Ok(Self {
shared,
monitor: Some(monitor),
})
}
pub fn suspend_handle(&self) -> SuspendHandle {
SuspendHandle {
shared: Arc::clone(&self.shared),
}
}
pub fn current_pid(&self) -> Option<u32> {
lock(&self.shared.state).last_pid
}
pub fn restart_count(&self) -> u64 {
lock(&self.shared.state).restarts
}
pub fn is_alive(&self) -> bool {
let mut st = lock(&self.shared.state);
match st.child.as_mut() {
Some(c) => matches!(c.try_wait(), Ok(None)),
None => false,
}
}
pub fn shutdown(mut self) {
self.do_shutdown();
}
fn do_shutdown(&mut self) {
self.shared.shutdown.store(true, Ordering::SeqCst);
if let Some(t) = self.monitor.take() {
let _ = t.join();
}
let mut st = lock(&self.shared.state);
if let Some(mut c) = st.child.take() {
let _ = c.kill();
let _ = c.wait();
}
}
}
impl Drop for Supervisor {
fn drop(&mut self) {
self.do_shutdown();
}
}
#[derive(Clone)]
pub struct SuspendHandle {
shared: Arc<Shared>,
}
impl SuspendHandle {
pub fn suspend(&self) {
self.shared.suspended.store(true, Ordering::SeqCst);
let mut st = lock(&self.shared.state);
if let Some(c) = st.child.as_mut() {
let _ = c.kill();
}
}
pub fn resume(&self) {
self.shared.suspended.store(false, Ordering::SeqCst);
}
pub fn is_suspended(&self) -> bool {
self.shared.suspended.load(Ordering::SeqCst)
}
pub fn child_alive(&self) -> bool {
let mut st = lock(&self.shared.state);
matches!(st.child.as_mut().map(|c| c.try_wait()), Some(Ok(None)))
}
}
fn invoke_on_spawn(shared: &Shared, child: &mut Child) {
let mut hook = shared.on_spawn.lock().unwrap_or_else(|e| e.into_inner());
(*hook)(child);
}
fn lock<T>(m: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
m.lock().unwrap_or_else(|e| e.into_inner())
}
fn monitor_loop(shared: Arc<Shared>) {
let mut backoff = MIN_BACKOFF;
loop {
if shared.shutdown.load(Ordering::SeqCst) {
break;
}
let dead = {
let mut st = lock(&shared.state);
match st.child.as_mut() {
Some(c) => match c.try_wait() {
Ok(Some(_status)) => true, Ok(None) => false, Err(_) => true, },
None => true,
}
};
if !dead {
thread::sleep(POLL_INTERVAL);
continue;
}
if shared.shutdown.load(Ordering::SeqCst) {
break;
}
{
let mut st = lock(&shared.state);
if let Some(mut old) = st.child.take() {
let _ = old.wait();
}
}
while shared.suspended.load(Ordering::SeqCst) && !shared.shutdown.load(Ordering::SeqCst) {
thread::sleep(POLL_INTERVAL);
}
thread::sleep(backoff);
if shared.shutdown.load(Ordering::SeqCst) {
break;
}
match (shared.spawn)() {
Ok(mut child) => {
invoke_on_spawn(&shared, &mut child);
let mut st = lock(&shared.state);
st.last_pid = Some(child.id());
st.child = Some(child);
st.restarts += 1;
backoff = MIN_BACKOFF;
}
Err(_) => {
backoff = (backoff * 2).min(MAX_BACKOFF);
}
}
}
let mut st = lock(&shared.state);
if let Some(mut c) = st.child.take() {
let _ = c.kill();
let _ = c.wait();
}
}
pub fn rust_analyzer_command() -> io::Result<Command> {
let exe = resolve_rust_analyzer()?;
let mut cmd = Command::new(exe);
cmd.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null());
#[cfg(unix)]
{
use std::os::unix::process::CommandExt as _;
cmd.process_group(0);
unsafe {
cmd.pre_exec(|| {
unsafe extern "C" {
fn setsid() -> i32;
}
let _ = setsid();
Ok(())
});
}
}
apply_ra_allocator_env(&mut cmd);
Ok(cmd)
}
fn apply_ra_allocator_env(cmd: &mut Command) {
if matches!(std::env::var("TF_RA_ALLOC").as_deref(), Ok("off")) {
return;
}
if std::env::var_os("MALLOC_ARENA_MAX").is_none() {
cmd.env("MALLOC_ARENA_MAX", "2");
}
let want_jemalloc = matches!(std::env::var("TF_RA_JEMALLOC").as_deref(), Ok("1"))
&& std::env::var_os("LD_PRELOAD").is_none();
let preload = if want_jemalloc { find_jemalloc() } else { None };
if let Some(so) = preload {
cmd.env("LD_PRELOAD", so);
}
}
fn find_jemalloc() -> Option<std::ffi::OsString> {
if let Some(p) =
std::env::var_os("TF_RA_JEMALLOC_SO").filter(|p| std::path::Path::new(p).exists())
{
return Some(p);
}
const CANDIDATES: &[&str] = &[
"/usr/lib/x86_64-linux-gnu/libjemalloc.so.2",
"/usr/lib/aarch64-linux-gnu/libjemalloc.so.2",
"/usr/local/lib/libjemalloc.so.2",
"/usr/lib/libjemalloc.so.2",
"/lib/x86_64-linux-gnu/libjemalloc.so.2",
];
CANDIDATES
.iter()
.find(|p| std::path::Path::new(p).exists())
.map(std::ffi::OsString::from)
}
pub struct ReapOnDrop(Option<std::process::Child>);
impl ReapOnDrop {
pub fn new(child: std::process::Child) -> Self {
Self(Some(child))
}
pub fn take_stdio(&mut self) -> Option<(std::process::ChildStdin, std::process::ChildStdout)> {
let child = self.0.as_mut()?;
let stdin = child.stdin.take()?;
let stdout = child.stdout.take()?;
Some((stdin, stdout))
}
}
impl Drop for ReapOnDrop {
fn drop(&mut self) {
let Some(mut child) = self.0.take() else {
return;
};
#[cfg(unix)]
{
let pid = child.id() as i32;
let session_members = snapshot_session_members(pid);
unsafe {
unsafe extern "C" {
fn kill(pid: i32, sig: i32) -> i32;
}
const SIGKILL: i32 = 9;
let _ = kill(-pid, SIGKILL);
for m in session_members {
if m != pid {
let _ = kill(m, SIGKILL);
}
}
}
}
let _ = child.kill();
let _ = child.wait();
}
}
#[cfg(unix)]
fn snapshot_session_members(sid: i32) -> Vec<i32> {
let Ok(output) = Command::new("pgrep")
.arg("-s")
.arg(sid.to_string())
.output()
else {
return Vec::new();
};
if !output.status.success() {
return Vec::new();
}
String::from_utf8_lossy(&output.stdout)
.lines()
.filter_map(|l| l.trim().parse::<i32>().ok())
.collect()
}
fn resolve_rust_analyzer() -> io::Result<OsString> {
if let Some(p) = std::env::var_os("RUST_ANALYZER") {
return Ok(p);
}
if let Some(p) = rustup_which_rust_analyzer() {
return Ok(p);
}
Ok(OsString::from("rust-analyzer"))
}
fn rustup_which_rust_analyzer() -> Option<OsString> {
let out = Command::new("rustup")
.args(["which", "rust-analyzer"])
.output()
.ok()?;
if !out.status.success() {
return None;
}
let path = String::from_utf8_lossy(&out.stdout).trim().to_string();
if path.is_empty() {
None
} else {
Some(OsString::from(path))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(unix)]
#[test]
fn reap_on_drop_kills_the_child_on_scope_exit() {
let child = Command::new("sleep")
.arg("30")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.expect("spawn sleep");
let pid = child.id() as i32;
{
let _guard = ReapOnDrop::new(child);
assert!(
pid_is_alive(pid),
"child should be alive while ReapOnDrop guard exists"
);
}
for _ in 0..40 {
if !pid_is_alive(pid) {
return;
}
thread::sleep(Duration::from_millis(5));
}
panic!("ReapOnDrop guard exited but pid {pid} still alive");
}
#[cfg(unix)]
#[test]
fn reap_on_drop_take_stdio_returns_pipes_once() {
let child = Command::new("sleep")
.arg("30")
.stdin(Stdio::piped())
.stdout(Stdio::piped())
.stderr(Stdio::null())
.spawn()
.expect("spawn sleep with piped stdio");
let mut guard = ReapOnDrop::new(child);
let first = guard.take_stdio();
assert!(first.is_some(), "first take_stdio yields the pipes");
let second = guard.take_stdio();
assert!(second.is_none(), "second take_stdio is None");
drop(first);
drop(guard);
}
#[cfg(unix)]
fn pid_is_alive(pid: i32) -> bool {
unsafe {
unsafe extern "C" {
fn kill(pid: i32, sig: i32) -> i32;
}
kill(pid, 0) == 0
}
}
#[cfg(unix)]
#[test]
fn snapshot_session_members_returns_self_for_own_session() {
let my_pid = std::process::id() as i32;
let Ok(out) = Command::new("ps")
.arg("-o")
.arg("sid=")
.arg("-p")
.arg(my_pid.to_string())
.output()
else {
return;
};
let Some(sid) = String::from_utf8_lossy(&out.stdout)
.trim()
.parse::<i32>()
.ok()
else {
return;
};
let members = snapshot_session_members(sid);
if !members.is_empty() {
assert!(
members.contains(&my_pid),
"session snapshot {members:?} should include our own pid {my_pid}"
);
}
}
#[cfg(unix)]
#[test]
fn snapshot_session_members_for_unknown_sid_is_empty_not_panic() {
let v = snapshot_session_members(0x7FFF_FFFF);
assert!(v.is_empty(), "nonsense sid → empty Vec, got {v:?}");
}
#[cfg(unix)]
#[test]
fn reap_on_drop_with_session_walk_still_kills_immediate_child() {
let child = Command::new("sleep")
.arg("30")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
.expect("spawn sleep");
let pid = child.id() as i32;
{
let _g = ReapOnDrop::new(child);
assert!(pid_is_alive(pid), "alive while guard in scope");
}
for _ in 0..40 {
if !pid_is_alive(pid) {
return;
}
thread::sleep(Duration::from_millis(5));
}
panic!("deepened ReapOnDrop still must kill the immediate child; pid {pid} alive");
}
#[test]
fn rust_analyzer_command_is_resolvable_and_piped() {
let cmd = rust_analyzer_command().expect("command resolves");
assert!(!format!("{cmd:?}").is_empty());
}
#[test]
fn supervisor_reports_initial_pid_and_zero_restarts() {
let sup = Supervisor::start(|| {
Command::new("sleep")
.arg("30")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
})
.expect("start");
assert!(sup.current_pid().is_some());
assert_eq!(sup.restart_count(), 0);
assert!(sup.is_alive());
sup.shutdown();
}
#[cfg(unix)]
#[test]
fn on_spawn_hook_fires_on_initial_and_after_kill9_restart() {
use std::sync::atomic::AtomicUsize;
let calls = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&calls);
let sup = Supervisor::start_with_hook(
|| {
Command::new("sleep")
.arg("30")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
},
move |_child: &mut Child| {
counter.fetch_add(1, Ordering::SeqCst);
},
)
.expect("start_with_hook");
assert_eq!(
calls.load(Ordering::SeqCst),
1,
"hook must fire once on the initial spawn"
);
let pid1 = sup.current_pid().expect("first pid");
assert!(
Command::new("kill")
.arg("-9")
.arg(pid1.to_string())
.status()
.expect("invoke kill(1)")
.success()
);
let deadline = std::time::Instant::now() + Duration::from_secs(15);
while std::time::Instant::now() < deadline {
if sup.restart_count() >= 1 && calls.load(Ordering::SeqCst) >= 2 {
break;
}
thread::sleep(Duration::from_millis(20));
}
assert!(sup.restart_count() >= 1, "supervisor must have restarted");
assert!(
calls.load(Ordering::SeqCst) >= 2,
"hook must re-fire on the transparent restart (re-init LSP)"
);
sup.shutdown();
}
#[cfg(unix)]
#[test]
fn supervisor_suspend_reclaims_then_resume_respawns() {
use std::sync::atomic::AtomicUsize;
let calls = Arc::new(AtomicUsize::new(0));
let counter = Arc::clone(&calls);
let sup = Supervisor::start_with_hook(
|| {
Command::new("sleep")
.arg("60")
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()
},
move |_child: &mut Child| {
counter.fetch_add(1, Ordering::SeqCst);
},
)
.expect("start_with_hook");
let h = sup.suspend_handle();
assert!(h.child_alive(), "alive on initial spawn");
assert!(!h.is_suspended());
assert_eq!(calls.load(Ordering::SeqCst), 1, "hook fired once (spawn)");
h.suspend();
assert!(h.is_suspended());
let deadline = std::time::Instant::now() + Duration::from_secs(10);
while h.child_alive() && std::time::Instant::now() < deadline {
thread::sleep(Duration::from_millis(20));
}
assert!(!h.child_alive(), "suspend() must reap the child");
let restarts_at_suspend = sup.restart_count();
thread::sleep(Duration::from_millis(400));
assert!(
!h.child_alive(),
"a suspended RA must NOT be auto-respawned (RAM stays reclaimed)"
);
h.resume();
assert!(!h.is_suspended());
let deadline = std::time::Instant::now() + Duration::from_secs(15);
while !h.child_alive() && std::time::Instant::now() < deadline {
thread::sleep(Duration::from_millis(20));
}
assert!(h.child_alive(), "resume() must respawn RA");
assert!(
sup.restart_count() > restarts_at_suspend,
"resume respawn must count as a restart"
);
assert!(
calls.load(Ordering::SeqCst) >= 2,
"on_spawn hook must re-fire on resume (LSP re-init/re-did_open)"
);
sup.shutdown();
}
}