implement dump_traceback_later and cancel_dump_traceback_later

arihant2math · arihant2math · commit 82b2a898baa8 · 2025-02-26T11:44:09.000-08:00
Signed-off-by: Ashwin Naren &lt;arihant2math@gmail.com&gt;
diff --git 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<FrameRef>,
+    }
+
+    static FAULT_HANDLER_THREAD: PyRwLock<Option<FaultHandlerThread>> = 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<bool>,
+        _file: OptionalArg<i32>,
+        exit: OptionalArg<bool>,
+        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 {
