Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 29 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions crates/vm/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down
119 changes: 119 additions & 0 deletions crates/vm/src/vm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,8 @@ pub struct VirtualMachine {
pub state: PyRc<PyGlobalState>,
pub initialized: bool,
recursion_depth: Cell<usize>,
/// C stack soft limit for detecting stack overflow (like c_stack_soft_limit)
c_stack_soft_limit: Cell<usize>,
/// Async generator firstiter hook (per-thread, set via sys.set_asyncgen_hooks)
pub async_gen_firstiter: RefCell<Option<PyObjectRef>>,
/// Async generator finalizer hook (per-thread, set via sys.set_asyncgen_hooks)
Expand Down Expand Up @@ -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),
};
Expand Down Expand Up @@ -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::<usize>();

/// 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<R, F: FnOnce() -> PyResult<R>>(&self, _where: &str, f: F) -> PyResult<R> {
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);
Expand Down
1 change: 1 addition & 0 deletions crates/vm/src/vm/thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
};
Expand Down
Loading