From 4bdf5c4ab6c0eb498ed47f587a3c5d7c79bfd71f Mon Sep 17 00:00:00 2001 From: "Jeong, YunWon" Date: Fri, 16 Jan 2026 19:51:15 +0900 Subject: [PATCH] stack soft limit --- Cargo.lock | 29 +++++++++ crates/vm/Cargo.toml | 1 + crates/vm/src/vm/mod.rs | 119 +++++++++++++++++++++++++++++++++++++ crates/vm/src/vm/thread.rs | 1 + 4 files changed, 150 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 41ab3280428..8df208e4aaa 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -142,6 +142,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "ar_archive_writer" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c269894b6fe5e9d7ada0cf69b5bf847ff35bc25fc271f08e1d080fce80339a" +dependencies = [ + "object", +] + [[package]] name = "arbitrary" version = "1.4.2" @@ -2081,6 +2090,15 @@ dependencies = [ "syn", ] +[[package]] +name = "object" +version = "0.32.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6a622008b6e321afc04970976f62ee297fdbaa6f95318ca343e3eebb9648441" +dependencies = [ + "memchr", +] + [[package]] name = "oid-registry" version = "0.8.1" @@ -2448,6 +2466,16 @@ dependencies = [ "unicode-ident", ] +[[package]] +name = "psm" +version = "0.1.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d11f2fedc3b7dafdc2851bc52f277377c5473d378859be234bc7ebb593144d01" +dependencies = [ + "ar_archive_writer", + "cc", +] + [[package]] name = "pymath" version = "0.1.5" @@ -3280,6 +3308,7 @@ dependencies = [ "optional", "parking_lot", "paste", + "psm", "result-like", "ruff_python_ast", "ruff_python_parser", diff --git a/crates/vm/Cargo.toml b/crates/vm/Cargo.toml index b74aba41145..da01eff65b9 100644 --- a/crates/vm/Cargo.toml +++ b/crates/vm/Cargo.toml @@ -77,6 +77,7 @@ memchr = { workspace = true } caseless = "0.2.2" flamer = { version = "0.5", optional = true } half = "2" +psm = "0.1" optional = { workspace = true } result-like = "0.5.0" timsort = "0.1.2" diff --git a/crates/vm/src/vm/mod.rs b/crates/vm/src/vm/mod.rs index 0c29d34ef88..5bd2b9e8297 100644 --- a/crates/vm/src/vm/mod.rs +++ b/crates/vm/src/vm/mod.rs @@ -86,6 +86,8 @@ pub struct VirtualMachine { pub state: PyRc, pub initialized: bool, recursion_depth: Cell, + /// C stack soft limit for detecting stack overflow (like c_stack_soft_limit) + c_stack_soft_limit: Cell, /// Async generator firstiter hook (per-thread, set via sys.set_asyncgen_hooks) pub async_gen_firstiter: RefCell>, /// Async generator finalizer hook (per-thread, set via sys.set_asyncgen_hooks) @@ -228,6 +230,7 @@ impl VirtualMachine { }), initialized: false, recursion_depth: Cell::new(0), + c_stack_soft_limit: Cell::new(Self::calculate_c_stack_soft_limit()), async_gen_firstiter: RefCell::new(None), async_gen_finalizer: RefCell::new(None), }; @@ -689,11 +692,127 @@ impl VirtualMachine { self.recursion_depth.get() } + /// Stack margin bytes (like _PyOS_STACK_MARGIN_BYTES). + /// 2048 * sizeof(void*) = 16KB for 64-bit. + const STACK_MARGIN_BYTES: usize = 2048 * std::mem::size_of::(); + + /// Get the stack boundaries using platform-specific APIs. + /// Returns (base, top) where base is the lowest address and top is the highest. + #[cfg(all(not(miri), windows))] + fn get_stack_bounds() -> (usize, usize) { + use windows_sys::Win32::System::Threading::{ + GetCurrentThreadStackLimits, SetThreadStackGuarantee, + }; + let mut low: usize = 0; + let mut high: usize = 0; + unsafe { + GetCurrentThreadStackLimits(&mut low as *mut usize, &mut high as *mut usize); + // Add the guaranteed stack space (reserved for exception handling) + let mut guarantee: u32 = 0; + SetThreadStackGuarantee(&mut guarantee); + low += guarantee as usize; + } + (low, high) + } + + /// Get stack boundaries on non-Windows platforms. + /// Falls back to estimating based on current stack pointer. + #[cfg(all(not(miri), not(windows)))] + fn get_stack_bounds() -> (usize, usize) { + // Use pthread_attr_getstack on platforms that support it + #[cfg(any(target_os = "linux", target_os = "android"))] + { + use libc::{ + pthread_attr_destroy, pthread_attr_getstack, pthread_attr_t, pthread_getattr_np, + pthread_self, + }; + let mut attr: pthread_attr_t = unsafe { std::mem::zeroed() }; + unsafe { + if pthread_getattr_np(pthread_self(), &mut attr) == 0 { + let mut stack_addr: *mut libc::c_void = std::ptr::null_mut(); + let mut stack_size: libc::size_t = 0; + if pthread_attr_getstack(&attr, &mut stack_addr, &mut stack_size) == 0 { + pthread_attr_destroy(&mut attr); + let base = stack_addr as usize; + let top = base + stack_size; + return (base, top); + } + pthread_attr_destroy(&mut attr); + } + } + } + + #[cfg(target_os = "macos")] + { + use libc::{pthread_get_stackaddr_np, pthread_get_stacksize_np, pthread_self}; + unsafe { + let thread = pthread_self(); + let stack_top = pthread_get_stackaddr_np(thread) as usize; + let stack_size = pthread_get_stacksize_np(thread); + let stack_base = stack_top - stack_size; + return (stack_base, stack_top); + } + } + + // Fallback: estimate based on current SP and a default stack size + #[allow(unreachable_code)] + { + let current_sp = psm::stack_pointer() as usize; + // Assume 8MB stack, estimate base + let estimated_size = 8 * 1024 * 1024; + let base = current_sp.saturating_sub(estimated_size); + let top = current_sp + 1024 * 1024; // Assume we're not at the very top + (base, top) + } + } + + /// Calculate the C stack soft limit based on actual stack boundaries. + /// soft_limit = base + 2 * margin (for downward-growing stacks) + #[cfg(not(miri))] + fn calculate_c_stack_soft_limit() -> usize { + let (base, _top) = Self::get_stack_bounds(); + // Soft limit is 2 margins above the base + base + Self::STACK_MARGIN_BYTES * 2 + } + + /// Miri doesn't support inline assembly, so disable C stack checking. + #[cfg(miri)] + fn calculate_c_stack_soft_limit() -> usize { + 0 + } + + /// Check if we're near the C stack limit (like _Py_MakeRecCheck). + /// Returns true only when stack pointer is in the "danger zone" between + /// soft_limit and hard_limit (soft_limit - 2*margin). + #[cfg(not(miri))] + #[inline(always)] + fn check_c_stack_overflow(&self) -> bool { + let current_sp = psm::stack_pointer() as usize; + let soft_limit = self.c_stack_soft_limit.get(); + // Stack grows downward: check if we're below soft limit but above hard limit + // This matches CPython's _Py_MakeRecCheck behavior + current_sp < soft_limit + && current_sp >= soft_limit.saturating_sub(Self::STACK_MARGIN_BYTES * 2) + } + + /// Miri doesn't support inline assembly, so always return false. + #[cfg(miri)] + #[inline(always)] + fn check_c_stack_overflow(&self) -> bool { + false + } + /// Used to run the body of a (possibly) recursive function. It will raise a /// RecursionError if recursive functions are nested far too many times, /// preventing a stack overflow. pub fn with_recursion PyResult>(&self, _where: &str, f: F) -> PyResult { self.check_recursive_call(_where)?; + + // Native stack guard: check C stack like _Py_MakeRecCheck + if self.check_c_stack_overflow() { + return Err(self.new_recursion_error(_where.to_string())); + } + self.recursion_depth.set(self.recursion_depth.get() + 1); let result = f(); self.recursion_depth.set(self.recursion_depth.get() - 1); diff --git a/crates/vm/src/vm/thread.rs b/crates/vm/src/vm/thread.rs index 7188aa6d270..bfbfc4f6a04 100644 --- a/crates/vm/src/vm/thread.rs +++ b/crates/vm/src/vm/thread.rs @@ -234,6 +234,7 @@ impl VirtualMachine { state: self.state.clone(), initialized: self.initialized, recursion_depth: Cell::new(0), + c_stack_soft_limit: Cell::new(VirtualMachine::calculate_c_stack_soft_limit()), async_gen_firstiter: RefCell::new(None), async_gen_finalizer: RefCell::new(None), };