stack soft limit (#6754)

This commit is contained in:
Jeong, YunWon
2026-01-18 18:09:55 +09:00
committed by GitHub
parent 941a7bb784
commit f8a78e6ab6
4 changed files with 150 additions and 0 deletions

29
Cargo.lock generated
View File

@@ -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",

View File

@@ -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"

View File

@@ -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)
@@ -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::<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);

View File

@@ -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),
};