diff --git a/stdlib/src/faulthandler.rs b/stdlib/src/faulthandler.rs index 2fb93ecc884..e6b183678ef 100644 --- a/stdlib/src/faulthandler.rs +++ b/stdlib/src/faulthandler.rs @@ -2,7 +2,24 @@ pub(crate) use decl::make_module; #[pymodule(name = "faulthandler")] mod decl { - use crate::vm::{frame::Frame, function::OptionalArg, stdlib::sys::PyStderr, VirtualMachine}; + use crate::vm::{ + frame::{Frame, FrameRef}, + function::OptionalArg, + stdlib::sys::PyStderr, + VirtualMachine, + }; + use rustpython_common::lock::PyRwLock; + use std::io::Write; + + struct FaultHandlerThread { + enabled: bool, + repeat: bool, + exit: bool, + timeout: std::time::Duration, + frames: Vec, + } + + static FAULT_HANDLER_THREAD: PyRwLock> = PyRwLock::new(None); fn dump_frame(frame: &Frame, vm: &VirtualMachine) { let stderr = PyStderr(vm); @@ -29,6 +46,82 @@ mod decl { } } + #[pyfunction] + fn dump_traceback_later( + timeout: f32, + repeat: OptionalArg, + _file: OptionalArg, + exit: OptionalArg, + vm: &VirtualMachine, + ) { + let timeout_duration = std::time::Duration::from_secs_f32(timeout); + let repeat = repeat.unwrap_or(false); + let exit = exit.unwrap_or(false); + let frames = vm.frames.borrow().clone(); + let t_data = FaultHandlerThread { + enabled: true, + repeat, + timeout: timeout_duration, + frames, + exit, + }; + if let Some(t) = FAULT_HANDLER_THREAD.write().as_mut() { + *t = t_data; + } else { + std::thread::spawn(move || { + loop { + let thread_info = FAULT_HANDLER_THREAD.read(); + let thread_info = match thread_info.as_ref() { + Some(t) => t, + None => return, + }; + if !thread_info.enabled { + *FAULT_HANDLER_THREAD.write() = None; + return; + } + + std::thread::sleep(thread_info.timeout); + + let thread_info = FAULT_HANDLER_THREAD.read(); + let thread_info = match thread_info.as_ref() { + Some(t) => t, + None => return, + }; + if !thread_info.enabled { + *FAULT_HANDLER_THREAD.write() = None; + return; + } + for frame in thread_info.frames.iter() { + // TODO: Fix + let _ = writeln!(std::io::stderr(), "Stack (most recent call first):"); + let _ = writeln!( + std::io::stderr(), + " File \"{}\", line {} in {}", + frame.code.source_path, + frame.current_location().row.to_usize(), + frame.code.obj_name + ); + } + if thread_info.exit { + std::process::exit(1); + } + if !thread_info.repeat { + *FAULT_HANDLER_THREAD.write() = None; + return; + } + } + }); + } + } + + #[pyfunction] + fn cancel_dump_traceback_later() { + FAULT_HANDLER_THREAD + .write() + .as_mut() + .map(|t| t.enabled = false); + } + #[derive(FromArgs)] #[allow(unused)] struct EnableArgs {