lean_ctx/core/
interrupt.rs

1//! User-initiated cooperative cancellation (Ctrl-C / SIGINT).
2//!
3//! Long-running CLI builds (notably dense embedding via CUDA) spend most of
4//! their time inside ONNX Runtime FFI (`session.run()`). If the default SIGINT
5//! disposition kills the process mid-kernel, the CUDA context is never torn
6//! down cleanly: the process lingers as a zombie and its VRAM stays allocated
7//! until the driver reclaims it (observed under WSL2 GPU passthrough).
8//!
9//! This module installs a *cooperative* SIGINT handler: the first Ctrl-C only
10//! flips a global flag. Cancellable loops poll [`is_cancelled`] **between** FFI
11//! calls and return early, so control is back in Rust code — not inside a CUDA
12//! kernel — when the process exits. That lets the driver reclaim VRAM on a
13//! clean exit. A second Ctrl-C forces an immediate `_exit` for the impatient.
14//!
15//! The handler body only performs async-signal-safe operations (atomic stores
16//! and, on the second signal, `libc::_exit`).
17
18use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
19
20static CANCELLED: AtomicBool = AtomicBool::new(false);
21static HANDLER_INSTALLED: AtomicBool = AtomicBool::new(false);
22static SIGNAL_COUNT: AtomicU8 = AtomicU8::new(0);
23
24/// `true` once the user has requested cancellation via Ctrl-C.
25pub fn is_cancelled() -> bool {
26    CANCELLED.load(Ordering::Relaxed)
27}
28
29/// Clear the cancellation state before starting a fresh cancellable operation.
30pub fn reset() {
31    CANCELLED.store(false, Ordering::SeqCst);
32    SIGNAL_COUNT.store(0, Ordering::SeqCst);
33}
34
35#[cfg(unix)]
36extern "C" fn handle_sigint(_sig: libc::c_int) {
37    CANCELLED.store(true, Ordering::SeqCst);
38    let count = SIGNAL_COUNT
39        .fetch_add(1, Ordering::SeqCst)
40        .saturating_add(1);
41    if count >= 2 {
42        // Second Ctrl-C: the user wants out now. `_exit` is async-signal-safe
43        // and terminates the process without running (unsafe-in-a-handler)
44        // destructors; the OS/driver reclaims the CUDA context on death.
45        // SAFETY: `_exit` is async-signal-safe (POSIX.1-2017 §2.4.3) and does
46        // not run C++ destructors, Rust `Drop` impls, or atexit handlers.
47        // We call it only after the second SIGINT, where graceful shutdown has
48        // already been requested and the user explicitly wants immediate exit.
49        unsafe { libc::_exit(130) };
50    }
51}
52
53/// Install the cooperative SIGINT handler (idempotent).
54///
55/// Call this at the start of a long-running, cancellable CLI operation. The
56/// daemon/MCP server never calls it, so [`is_cancelled`] stays `false` there
57/// and background embedding is unaffected.
58pub fn install_ctrlc_handler() {
59    if HANDLER_INSTALLED.swap(true, Ordering::SeqCst) {
60        return;
61    }
62    // SAFETY: `handle_sigint` only performs async-signal-safe work (atomic
63    // stores, and `libc::_exit` on the second signal). Registering it replaces
64    // the default terminate-immediately disposition with a cooperative one.
65    #[cfg(unix)]
66    unsafe {
67        libc::signal(
68            libc::SIGINT,
69            handle_sigint as *const () as libc::sighandler_t,
70        );
71    }
72}
lean_ctx/core/interrupt.rs

lean_ctx/core/
interrupt.rs
