diff --git a/hf_xet/Cargo.lock b/hf_xet/Cargo.lock index 7089f4ff..ccf2a954 100644 --- a/hf_xet/Cargo.lock +++ b/hf_xet/Cargo.lock @@ -753,16 +753,6 @@ version = "0.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e2931af7e13dc045d8e9d26afccc6fa115d64e115c9c84b1166288b46f6782c2" -[[package]] -name = "ctrlc" -version = "3.4.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "46f93780a459b7d656ef7f071fe699c4d3d2cb201c4b24d085b6ddc505276e73" -dependencies = [ - "nix 0.30.1", - "windows-sys 0.59.0", -] - [[package]] name = "data" version = "0.14.5" @@ -1410,7 +1400,6 @@ dependencies = [ "bipbuffer", "cas_client", "chrono", - "ctrlc", "data", "error_printer", "itertools 0.14.0", @@ -1422,8 +1411,10 @@ dependencies = [ "serde", "serde_json", "signal-hook", + "signal-hook-registry", "tracing", "utils", + "winapi", "xet_config", "xet_logging", "xet_runtime", @@ -2269,18 +2260,6 @@ dependencies = [ "libc", ] -[[package]] -name = "nix" -version = "0.30.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" -dependencies = [ - "bitflags 2.9.2", - "cfg-if 1.0.1", - "cfg_aliases", - "libc", -] - [[package]] name = "nom" version = "7.1.3" @@ -2638,7 +2617,7 @@ dependencies = [ "inferno", "libc", "log", - "nix 0.26.4", + "nix", "once_cell", "prost 0.12.6", "protobuf 2.28.0", diff --git a/hf_xet/Cargo.toml b/hf_xet/Cargo.toml index dd8b5a45..e0b7b240 100644 --- a/hf_xet/Cargo.toml +++ b/hf_xet/Cargo.toml @@ -42,10 +42,11 @@ tracing = "0.1" # Unix-specific dependencies [target.'cfg(unix)'.dependencies] signal-hook = "0.3" +signal-hook-registry = "1.4" # Windows-specific dependencies [target.'cfg(windows)'.dependencies] -ctrlc = "3.4" +winapi = { version = "0.3", features = ["consoleapi", "wincon", "errhandlingapi"] } [features] default = ["no-default-cache"] # By default, hf_xet disables the disk cache. diff --git a/hf_xet/src/runtime.rs b/hf_xet/src/runtime.rs index a47a9bac..4fbe3ec2 100644 --- a/hf_xet/src/runtime.rs +++ b/hf_xet/src/runtime.rs @@ -30,17 +30,57 @@ fn install_sigint_handler() -> Result<(), MultithreadedRuntimeError> { Ok(()) } +#[cfg(windows)] +extern "system" fn console_ctrl_handler( + ctrl_type: winapi::shared::minwindef::DWORD, +) -> winapi::shared::minwindef::BOOL { + use winapi::um::wincon; + + // Only handle CTRL_C_EVENT + if ctrl_type == wincon::CTRL_C_EVENT { + // Check if we have active operations + let has_active_ops = { + let guard = MULTITHREADED_RUNTIME.read().unwrap(); + if let Some((runtime_pid, ref runtime)) = *guard { + runtime_pid == std::process::id() && runtime.external_executor_count() > 0 + } else { + false + } + }; + + if has_active_ops { + // We have active operations, handle it ourselves + SIGINT_DETECTED.store(true, Ordering::SeqCst); + winapi::shared::minwindef::TRUE + } else { + // No active operations, let Python's handler (or default) handle it + // Return FALSE to continue handler chain + winapi::shared::minwindef::FALSE + } + } else { + // For other control events, let default handler process them + winapi::shared::minwindef::FALSE + } +} + #[cfg(windows)] fn install_sigint_handler() -> Result<(), MultithreadedRuntimeError> { - // On Windows, use ctrlc crate. - // This sets a callback to run on Ctrl-C: - let sigint_detected_flag = SIGINT_DETECTED.clone(); - ctrlc::set_handler(move || { - sigint_detected_flag.store(true, Ordering::SeqCst); - }) - .map_err(|e| { - MultithreadedRuntimeError::Other(format!("Initialization Error: Unable to register SIGINT handler {e:?}")) - })?; + use winapi::um::consoleapi::SetConsoleCtrlHandler; + use winapi::um::wincon::CTRL_C_EVENT; + + // Install our handler using Windows API directly (instead of ctrlc) + // Our handler checks if operations are active: + // - If active: handles Ctrl+C and returns TRUE (stops propagation) + // - If not active: returns FALSE (allows Python's handler or default to handle it) + // This way we don't need to save/restore Python's handler - we just delegate to it + unsafe { + if SetConsoleCtrlHandler(Some(console_ctrl_handler), winapi::shared::minwindef::TRUE) == 0 { + let error = winapi::um::errhandlingapi::GetLastError(); + return Err(MultithreadedRuntimeError::Other(format!( + "Initialization Error: Unable to register SIGINT handler. Windows error: {error}" + ))); + } + } Ok(()) }