Visitar URL original
backup · RustPython/RustPython@db453fc · GitHub
Skip to content

Commit db453fc

Browse files
youknowoneyouknowone
committed
backup
1 parent b29c3b8 commit db453fc

2 files changed

Lines changed: 179 additions & 77 deletions

File tree

‎crates/vm/src/stdlib/posix.rs‎

Lines changed: 115 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -872,57 +872,129 @@ pub mod module {
872872
run_at_forkers(after_forkers_parent, false, vm);
873873
}
874874

875-
/// Warn if forking from a multi-threaded process
876-
fn warn_if_multi_threaded(name: &str, vm: &VirtualMachine) {
877-
// Only check threading if it was already imported
878-
// Avoid vm.import() which can execute arbitrary Python code in the fork path
879-
let threading = match vm
880-
.sys_module
881-
.get_attr("modules", vm)
882-
.and_then(|m| m.get_item("threading", vm))
875+
/// Best-effort number of OS threads in this process.
876+
/// Returns <= 0 when unavailable, mirroring CPython fallback behavior.
877+
fn get_number_of_os_threads() -> isize {
878+
#[cfg(target_os = "macos")]
883879
{
884-
Ok(m) => m,
885-
Err(_) => return,
886-
};
887-
let active = threading.get_attr("_active", vm).ok();
888-
let limbo = threading.get_attr("_limbo", vm).ok();
880+
type MachPortT = libc::c_uint;
881+
type KernReturnT = libc::c_int;
882+
type MachMsgTypeNumberT = libc::c_uint;
883+
type ThreadActArrayT = *mut MachPortT;
884+
const KERN_SUCCESS: KernReturnT = 0;
885+
unsafe extern "C" {
886+
fn mach_task_self() -> MachPortT;
887+
fn task_for_pid(
888+
task: MachPortT,
889+
pid: libc::c_int,
890+
target_task: *mut MachPortT,
891+
) -> KernReturnT;
892+
fn task_threads(
893+
target_task: MachPortT,
894+
act_list: *mut ThreadActArrayT,
895+
act_list_cnt: *mut MachMsgTypeNumberT,
896+
) -> KernReturnT;
897+
fn vm_deallocate(
898+
target_task: MachPortT,
899+
address: libc::uintptr_t,
900+
size: libc::uintptr_t,
901+
) -> KernReturnT;
902+
}
889903

890-
let count_dict = |obj: Option<crate::PyObjectRef>| -> usize {
891-
obj.and_then(|o| o.length_opt(vm))
892-
.and_then(|r| r.ok())
893-
.unwrap_or(0)
894-
};
904+
let self_task = unsafe { mach_task_self() };
905+
let mut proc_task: MachPortT = 0;
906+
if unsafe { task_for_pid(self_task, libc::getpid(), &mut proc_task) } == KERN_SUCCESS {
907+
let mut threads: ThreadActArrayT = core::ptr::null_mut();
908+
let mut n_threads: MachMsgTypeNumberT = 0;
909+
if unsafe { task_threads(proc_task, &mut threads, &mut n_threads) } == KERN_SUCCESS
910+
{
911+
if !threads.is_null() {
912+
let _ = unsafe {
913+
vm_deallocate(
914+
self_task,
915+
threads as libc::uintptr_t,
916+
(n_threads as usize * core::mem::size_of::<MachPortT>())
917+
as libc::uintptr_t,
918+
)
919+
};
920+
}
921+
return n_threads as isize;
922+
}
923+
}
924+
0
925+
}
926+
#[cfg(target_os = "linux")]
927+
{
928+
use std::io::Read as _;
929+
let mut file = match std::fs::File::open("/proc/self/stat") {
930+
Ok(f) => f,
931+
Err(_) => return 0,
932+
};
933+
let mut buf = [0u8; 160];
934+
let n = match file.read(&mut buf) {
935+
Ok(n) => n,
936+
Err(_) => return 0,
937+
};
938+
let line = match std::str::from_utf8(&buf[..n]) {
939+
Ok(s) => s,
940+
Err(_) => return 0,
941+
};
942+
if let Some(field) = line.split_whitespace().nth(19) {
943+
return field.parse::<isize>().unwrap_or(0);
944+
}
945+
return 0;
946+
}
947+
#[cfg(not(any(target_os = "macos", target_os = "linux")))]
948+
{
949+
0
950+
}
951+
}
895952

896-
let num_threads = count_dict(active) + count_dict(limbo);
897-
if num_threads > 1 {
898-
// Use Python warnings module to ensure filters are applied correctly
899-
let Ok(warnings) = vm.import("warnings", 0) else {
900-
return;
953+
/// Warn if forking from a multi-threaded process.
954+
/// `num_os_threads` should be captured before parent after-fork hooks run.
955+
fn warn_if_multi_threaded(name: &str, num_os_threads: isize, vm: &VirtualMachine) {
956+
let num_threads = if num_os_threads > 0 {
957+
num_os_threads as usize
958+
} else {
959+
// CPython fallback: if OS-level count isn't available, use the
960+
// threading module's active+limbo view.
961+
// Only check threading if it was already imported. Avoid vm.import()
962+
// which can execute arbitrary Python code in the fork path.
963+
let threading = match vm
964+
.sys_module
965+
.get_attr("modules", vm)
966+
.and_then(|m| m.get_item("threading", vm))
967+
{
968+
Ok(m) => m,
969+
Err(_) => return,
901970
};
902-
let Ok(warn_fn) = warnings.get_attr("warn", vm) else {
903-
return;
971+
let active = threading.get_attr("_active", vm).ok();
972+
let limbo = threading.get_attr("_limbo", vm).ok();
973+
974+
// Match threading module internals and avoid sequence overcounting:
975+
// count only dict-backed _active/_limbo containers.
976+
let count_dict = |obj: Option<crate::PyObjectRef>| -> usize {
977+
obj.and_then(|o| {
978+
o.downcast_ref::<crate::builtins::PyDict>()
979+
.map(|d| d.__len__())
980+
})
981+
.unwrap_or(0)
904982
};
905983

984+
count_dict(active) + count_dict(limbo)
985+
};
986+
987+
if num_threads > 1 {
906988
let pid = unsafe { libc::getpid() };
907989
let msg = format!(
908990
"This process (pid={}) is multi-threaded, use of {}() may lead to deadlocks in the child.",
909991
pid, name
910992
);
911993

912-
// Call warnings.warn(message, DeprecationWarning, stacklevel=2)
913-
// stacklevel=2 to point to the caller of fork()
914-
let args = crate::function::FuncArgs::new(
915-
vec![
916-
vm.ctx.new_str(msg).into(),
917-
vm.ctx.exceptions.deprecation_warning.as_object().to_owned(),
918-
],
919-
crate::function::KwArgs::new(
920-
[("stacklevel".to_owned(), vm.ctx.new_int(2).into())]
921-
.into_iter()
922-
.collect(),
923-
),
924-
);
925-
let _ = warn_fn.call(args, vm);
994+
// Match PyErr_WarnFormat(..., stacklevel=1) in CPython.
995+
// Best effort: ignore failures like CPython does in this path.
996+
let _ =
997+
crate::stdlib::warnings::warn(vm.ctx.exceptions.deprecation_warning, msg, 1, vm);
926998
}
927999
}
9281000

@@ -953,9 +1025,12 @@ pub mod module {
9531025
if pid == 0 {
9541026
py_os_after_fork_child(vm);
9551027
} else {
1028+
// Match CPython timing: capture this before parent after-fork hooks
1029+
// in case those hooks start threads.
1030+
let num_os_threads = get_number_of_os_threads();
9561031
py_os_after_fork_parent(vm);
9571032
// Match CPython timing: warn only after parent callback path resumes world.
958-
warn_if_multi_threaded("fork", vm);
1033+
warn_if_multi_threaded("fork", num_os_threads, vm);
9591034
}
9601035
if pid == -1 {
9611036
Err(nix::Error::from_raw(saved_errno).into_pyexception(vm))

‎crates/vm/src/stdlib/thread.rs‎

Lines changed: 64 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ pub(crate) mod _thread {
2727
RawMutex, RawThreadId,
2828
lock_api::{RawMutex as RawMutexT, RawMutexTimed, RawReentrantMutex},
2929
};
30+
use rustpython_common::str::levenshtein::{MOVE_COST, levenshtein_distance};
3031
use std::thread;
3132

3233
// PYTHREAD_NAME: show current thread name
@@ -1349,39 +1350,47 @@ pub(crate) mod _thread {
13491350
}
13501351

13511352
fn parse_join_timeout(
1352-
timeout: OptionalArg<Option<Either<f64, i64>>>,
1353+
timeout: OptionalArg<Option<crate::PyObjectRef>>,
13531354
vm: &VirtualMachine,
13541355
) -> PyResult<Option<Duration>> {
13551356
const JOIN_TIMEOUT_MAX_SECONDS: i64 = TIMEOUT_MAX_IN_MICROSECONDS / 1_000_000;
1356-
match timeout.flatten() {
1357-
Some(Either::A(t)) => {
1358-
if t.is_nan() {
1359-
return Err(vm.new_value_error("Invalid value NaN (not a number)"));
1360-
}
1361-
if !t.is_finite() || !(-TIMEOUT_MAX..=TIMEOUT_MAX).contains(&t) {
1362-
return Err(
1363-
vm.new_overflow_error("timestamp out of range for platform time_t")
1364-
);
1365-
}
1366-
if t < 0.0 {
1367-
return Ok(None);
1368-
}
1369-
Ok(Some(Duration::from_secs_f64(t)))
1357+
let Some(timeout_obj) = timeout.flatten() else {
1358+
return Ok(None);
1359+
};
1360+
1361+
if let Some(t) = timeout_obj.try_index_opt(vm) {
1362+
let t: i64 = t?.try_to_primitive(vm).map_err(|_| {
1363+
vm.new_overflow_error("timestamp too large to convert to C PyTime_t")
1364+
})?;
1365+
if !(-JOIN_TIMEOUT_MAX_SECONDS..=JOIN_TIMEOUT_MAX_SECONDS).contains(&t) {
1366+
return Err(
1367+
vm.new_overflow_error("timestamp too large to convert to C PyTime_t")
1368+
);
13701369
}
1371-
Some(Either::B(t)) => {
1372-
if !(-JOIN_TIMEOUT_MAX_SECONDS..=JOIN_TIMEOUT_MAX_SECONDS).contains(&t) {
1373-
return Err(
1374-
vm.new_overflow_error("timestamp too large to convert to C PyTime_t")
1375-
);
1376-
}
1377-
if t < 0 {
1378-
Ok(None)
1379-
} else {
1380-
Ok(Some(Duration::from_secs(t as u64)))
1381-
}
1370+
if t < 0 {
1371+
return Ok(None);
1372+
}
1373+
return Ok(Some(Duration::from_secs(t as u64)));
1374+
}
1375+
1376+
if let Some(t) = timeout_obj.try_float_opt(vm) {
1377+
let t = t?.to_f64();
1378+
if t.is_nan() {
1379+
return Err(vm.new_value_error("Invalid value NaN (not a number)"));
13821380
}
1383-
None => Ok(None),
1381+
if !t.is_finite() || !(-TIMEOUT_MAX..=TIMEOUT_MAX).contains(&t) {
1382+
return Err(vm.new_overflow_error("timestamp out of range for platform time_t"));
1383+
}
1384+
if t < 0.0 {
1385+
return Ok(None);
1386+
}
1387+
return Ok(Some(Duration::from_secs_f64(t)));
13841388
}
1389+
1390+
Err(vm.new_type_error(format!(
1391+
"'{}' object cannot be interpreted as an integer or float",
1392+
timeout_obj.class().name()
1393+
)))
13851394
}
13861395

13871396
#[pygetset]
@@ -1412,7 +1421,7 @@ pub(crate) mod _thread {
14121421
#[pymethod]
14131422
fn join(
14141423
&self,
1415-
timeout: OptionalArg<Option<Either<f64, i64>>>,
1424+
timeout: OptionalArg<Option<crate::PyObjectRef>>,
14161425
vm: &VirtualMachine,
14171426
) -> PyResult<()> {
14181427
let timeout_duration = Self::parse_join_timeout(timeout, vm)?;
@@ -1438,15 +1447,6 @@ pub(crate) mod _thread {
14381447
"start_joinable_thread() takes at most 3 arguments ({given} given)"
14391448
)));
14401449
}
1441-
if let Some(unexpected) = f_args
1442-
.kwargs
1443-
.keys()
1444-
.find(|k| !matches!(k.as_str(), "function" | "handle" | "daemon"))
1445-
{
1446-
return Err(vm.new_type_error(format!(
1447-
"start_joinable_thread() got an unexpected keyword argument '{unexpected}'"
1448-
)));
1449-
}
14501450

14511451
let function_pos = f_args.take_positional();
14521452
let function_kw = f_args.take_keyword("function");
@@ -1481,6 +1481,33 @@ pub(crate) mod _thread {
14811481
.or(daemon_kw)
14821482
.map_or(Ok(true), |obj| obj.try_to_bool(vm))?;
14831483

1484+
// Match CPython parser precedence:
1485+
// - required positional/keyword argument errors are raised before
1486+
// unknown keyword errors when `function` is missing.
1487+
if let Some(unexpected) = f_args.kwargs.keys().next() {
1488+
let suggestion = ["function", "handle", "daemon"]
1489+
.iter()
1490+
.filter_map(|candidate| {
1491+
let max_distance = (unexpected.len() + candidate.len() + 3) * MOVE_COST / 6;
1492+
let distance = levenshtein_distance(
1493+
unexpected.as_bytes(),
1494+
candidate.as_bytes(),
1495+
max_distance,
1496+
);
1497+
(distance <= max_distance).then_some((distance, *candidate))
1498+
})
1499+
.min_by_key(|(distance, _)| *distance)
1500+
.map(|(_, candidate)| candidate);
1501+
let msg = if let Some(suggestion) = suggestion {
1502+
format!(
1503+
"start_joinable_thread() got an unexpected keyword argument '{unexpected}'. Did you mean '{suggestion}'?"
1504+
)
1505+
} else {
1506+
format!("start_joinable_thread() got an unexpected keyword argument '{unexpected}'")
1507+
};
1508+
return Err(vm.new_type_error(msg));
1509+
}
1510+
14841511
if function_obj.to_callable().is_none() {
14851512
return Err(vm.new_type_error("thread function must be callable"));
14861513
}

0 commit comments

Comments
 (0)