use std::cell::RefCell;
use std::sync::Arc;
use std::sync::Once;
use std::sync::atomic::{AtomicBool, Ordering};
use vmm_sys_util::eventfd::EventFd;
#[cfg(test)]
use std::sync::atomic::AtomicUsize;
type PanicHook = dyn Fn(&std::panic::PanicHookInfo<'_>) + Send + Sync + 'static;
fn make_hook(prev: Box<PanicHook>) -> Box<PanicHook> {
Box::new(move |info| {
VCPU_PANIC_CTX.with(|slot| {
if let Some(ctx) = slot.borrow().as_ref() {
if let Some(ref alive) = ctx.alive {
alive.store(false, Ordering::Release);
}
ctx.kill.store(true, Ordering::Release);
ctx.exited.store(true, Ordering::Release);
if let Some(ref evt) = ctx.kill_evt {
let _ = evt.write(1);
}
if let Some(ref evt) = ctx.exited_evt {
let _ = evt.write(1);
}
}
});
prev(info);
})
}
#[cfg(test)]
static INSTALL_COUNT: AtomicUsize = AtomicUsize::new(0);
#[cfg(test)]
static HOOK_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
thread_local! {
static VCPU_PANIC_CTX: RefCell<Option<VcpuPanicCtx>> = const { RefCell::new(None) };
}
#[derive(Clone)]
pub(crate) struct VcpuPanicCtx {
pub(crate) kill: Arc<AtomicBool>,
pub(crate) exited: Arc<AtomicBool>,
pub(crate) kill_evt: Option<Arc<EventFd>>,
pub(crate) exited_evt: Option<Arc<EventFd>>,
pub(crate) alive: Option<Arc<AtomicBool>>,
}
static HOOK_ONCE: Once = Once::new();
pub(crate) fn install_once() {
HOOK_ONCE.call_once(|| {
#[cfg(test)]
INSTALL_COUNT.fetch_add(1, Ordering::Relaxed);
let prev = std::panic::take_hook();
std::panic::set_hook(make_hook(prev));
});
}
#[cfg(test)]
fn install_hook_with_prev_for_test(prev: Box<PanicHook>) {
std::panic::set_hook(make_hook(prev));
}
pub(crate) fn with_vcpu_panic_ctx<R>(ctx: VcpuPanicCtx, body: impl FnOnce() -> R) -> R {
struct CtxGuard;
impl Drop for CtxGuard {
fn drop(&mut self) {
VCPU_PANIC_CTX.with(|slot| {
*slot.borrow_mut() = None;
});
}
}
VCPU_PANIC_CTX.with(|slot| {
*slot.borrow_mut() = Some(ctx);
});
let _guard = CtxGuard;
body()
}
#[cfg(test)]
mod tests {
use super::*;
use std::panic::{AssertUnwindSafe, catch_unwind};
#[test]
fn install_once_is_idempotent() {
for _ in 0..10 {
install_once();
}
}
#[test]
fn with_vcpu_panic_ctx_clears_thread_local_on_return() {
install_once();
let ctx = VcpuPanicCtx {
kill: Arc::new(AtomicBool::new(false)),
exited: Arc::new(AtomicBool::new(false)),
kill_evt: None,
exited_evt: None,
alive: None,
};
std::thread::spawn(move || {
with_vcpu_panic_ctx(ctx, || {});
VCPU_PANIC_CTX.with(|slot| {
assert!(
slot.borrow().is_none(),
"thread-local must be None after normal return",
);
});
})
.join()
.unwrap();
}
#[test]
fn with_vcpu_panic_ctx_clears_thread_local_on_unwind() {
install_once();
let ctx = VcpuPanicCtx {
kill: Arc::new(AtomicBool::new(false)),
exited: Arc::new(AtomicBool::new(false)),
kill_evt: None,
exited_evt: None,
alive: None,
};
std::thread::spawn(move || {
let _ = catch_unwind(AssertUnwindSafe(|| {
with_vcpu_panic_ctx(ctx, || panic!("test: intended panic"));
}));
VCPU_PANIC_CTX.with(|slot| {
assert!(
slot.borrow().is_none(),
"thread-local must be None after unwind — RAII guard must clear on drop, not just after body()",
);
});
})
.join()
.unwrap();
}
#[test]
fn panic_inside_ctx_flips_flags() {
install_once();
let kill = Arc::new(AtomicBool::new(false));
let exited = Arc::new(AtomicBool::new(false));
let ctx = VcpuPanicCtx {
kill: kill.clone(),
exited: exited.clone(),
kill_evt: None,
exited_evt: None,
alive: None,
};
let kill_c = kill.clone();
let exited_c = exited.clone();
let (kill_r, exited_r) = std::thread::spawn(move || {
let _ = catch_unwind(AssertUnwindSafe(|| {
with_vcpu_panic_ctx(ctx, || panic!("test: intended panic"));
}));
(
kill_c.load(Ordering::Acquire),
exited_c.load(Ordering::Acquire),
)
})
.join()
.unwrap();
assert!(kill_r, "kill must be flipped by the panic hook");
assert!(exited_r, "exited must be flipped by the panic hook");
}
#[test]
fn install_once_body_runs_exactly_once() {
install_once();
let after_first = INSTALL_COUNT.load(Ordering::Relaxed);
for _ in 0..20 {
install_once();
}
let after_many = INSTALL_COUNT.load(Ordering::Relaxed);
assert_eq!(
after_many, after_first,
"HOOK_ONCE body ran more than once under repeated install_once calls",
);
assert!(
after_many >= 1,
"INSTALL_COUNT must reach 1 after install_once",
);
}
#[test]
fn panic_inside_ctx_still_runs_prev_hook() {
let _guard = HOOK_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let saved = std::panic::take_hook();
let prev_ran = Arc::new(AtomicBool::new(false));
let prev_ran_c = prev_ran.clone();
install_hook_with_prev_for_test(Box::new(move |_info| {
prev_ran_c.store(true, Ordering::Release);
}));
let kill = Arc::new(AtomicBool::new(false));
let exited = Arc::new(AtomicBool::new(false));
let ctx = VcpuPanicCtx {
kill: kill.clone(),
exited: exited.clone(),
kill_evt: None,
exited_evt: None,
alive: None,
};
std::thread::spawn(move || {
let _ = catch_unwind(AssertUnwindSafe(|| {
with_vcpu_panic_ctx(ctx, || panic!("test: prev-hook chain"));
}));
})
.join()
.unwrap();
std::panic::set_hook(saved);
assert!(
prev_ran.load(Ordering::Acquire),
"prev hook must run after our hook in the chain",
);
assert!(
kill.load(Ordering::Acquire),
"our hook must flip kill before delegating to prev",
);
assert!(
exited.load(Ordering::Acquire),
"our hook must flip exited before delegating to prev",
);
}
#[test]
fn panic_outside_ctx_leaves_flags_alone() {
install_once();
let kill = Arc::new(AtomicBool::new(false));
let exited = Arc::new(AtomicBool::new(false));
let kill_c = kill.clone();
let exited_c = exited.clone();
let (kill_r, exited_r) = std::thread::spawn(move || {
let _ = catch_unwind(AssertUnwindSafe(|| {
panic!("test: intended panic without registered ctx");
}));
(
kill_c.load(Ordering::Acquire),
exited_c.load(Ordering::Acquire),
)
})
.join()
.unwrap();
assert!(!kill_r, "kill must stay false when no ctx registered");
assert!(!exited_r, "exited must stay false when no ctx registered");
}
#[test]
fn panic_inside_ctx_flips_alive_before_prev() {
let _guard = HOOK_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
let saved = std::panic::take_hook();
let alive = Arc::new(AtomicBool::new(true));
let alive_seen_by_prev = Arc::new(AtomicBool::new(true));
let alive_for_prev = alive.clone();
let alive_seen_clone = alive_seen_by_prev.clone();
install_hook_with_prev_for_test(Box::new(move |_info| {
alive_seen_clone.store(alive_for_prev.load(Ordering::Acquire), Ordering::Release);
}));
let ctx = VcpuPanicCtx {
kill: Arc::new(AtomicBool::new(false)),
exited: Arc::new(AtomicBool::new(false)),
kill_evt: None,
exited_evt: None,
alive: Some(alive.clone()),
};
std::thread::spawn(move || {
let _ = catch_unwind(AssertUnwindSafe(|| {
with_vcpu_panic_ctx(ctx, || panic!("test: alive flip"));
}));
})
.join()
.unwrap();
std::panic::set_hook(saved);
assert!(
!alive_seen_by_prev.load(Ordering::Acquire),
"prev hook must observe alive == false — our hook \
must flip alive synchronously before delegating",
);
assert!(
!alive.load(Ordering::Acquire),
"alive must remain false post-panic",
);
}
}