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
4 changes: 3 additions & 1 deletion Lib/test/libregrtest/save_env.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,7 +240,9 @@ def get_multiprocessing_process__dangling(self):
# Unjoined process objects can survive after process exits
multiprocessing_process._cleanup()
# This copies the weakrefs without making any strong reference
return multiprocessing_process._dangling.copy()
# TODO: RUSTPYTHON - filter out dead processes since gc doesn't clean WeakSet. Revert this line when we have a GC
# return multiprocessing_process._dangling.copy()
return {p for p in multiprocessing_process._dangling if p.is_alive()}
def restore_multiprocessing_process__dangling(self, saved):
multiprocessing_process = self.get_module('multiprocessing.process')
multiprocessing_process._dangling.clear()
Expand Down
12 changes: 8 additions & 4 deletions crates/vm/src/signal.rs
Original file line number Diff line number Diff line change
@@ -1,13 +1,17 @@
#![cfg_attr(target_os = "wasi", allow(dead_code))]
use crate::{PyResult, VirtualMachine};
use crate::{PyObjectRef, PyResult, VirtualMachine};
use alloc::fmt;
use core::cell::Cell;
use core::cell::{Cell, RefCell};
#[cfg(windows)]
use core::sync::atomic::AtomicIsize;
use core::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;

pub(crate) const NSIG: usize = 64;

pub(crate) fn new_signal_handlers() -> Box<RefCell<[Option<PyObjectRef>; NSIG]>> {
Box::new(const { RefCell::new([const { None }; NSIG]) })
}
static ANY_TRIGGERED: AtomicBool = AtomicBool::new(false);
// hack to get around const array repeat expressions, rust issue #79270
#[allow(
Expand Down Expand Up @@ -37,7 +41,7 @@ impl Drop for SignalHandlerGuard {
#[cfg_attr(feature = "flame-it", flame)]
#[inline(always)]
pub fn check_signals(vm: &VirtualMachine) -> PyResult<()> {
if vm.signal_handlers.is_none() {
if vm.signal_handlers.get().is_none() {
return Ok(());
}

Expand All @@ -58,7 +62,7 @@ fn trigger_signals(vm: &VirtualMachine) -> PyResult<()> {
let _guard = SignalHandlerGuard;

// unwrap should never fail since we check above
let signal_handlers = vm.signal_handlers.as_ref().unwrap().borrow();
let signal_handlers = vm.signal_handlers.get().unwrap().borrow();
for (signum, trigger) in TRIGGERS.iter().enumerate().skip(1) {
let triggered = trigger.swap(false, Ordering::Relaxed);
if triggered
Expand Down
5 changes: 5 additions & 0 deletions crates/vm/src/stdlib/posix.rs
Original file line number Diff line number Diff line change
Expand Up @@ -711,6 +711,11 @@ pub mod module {
#[cfg(feature = "threading")]
crate::stdlib::thread::after_fork_child(vm);

// Initialize signal handlers for the child's main thread.
// When forked from a worker thread, the OnceCell is empty.
vm.signal_handlers
.get_or_init(crate::signal::new_signal_handlers);

let after_forkers_child: Vec<PyObjectRef> = vm.state.after_forkers_child.lock().clone();
run_at_forkers(after_forkers_child, false, vm);
}
Expand Down
21 changes: 10 additions & 11 deletions crates/vm/src/stdlib/signal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,9 @@ pub(crate) mod _signal {
} else {
None
};
vm.signal_handlers.as_deref().unwrap().borrow_mut()[signum] = py_handler;
vm.signal_handlers
.get_or_init(signal::new_signal_handlers)
.borrow_mut()[signum] = py_handler;
}

let int_handler = module
Expand Down Expand Up @@ -220,10 +222,9 @@ pub(crate) mod _signal {
return Err(vm.new_value_error(format!("signal number {} out of range", signalnum)));
}
}
let signal_handlers = vm
.signal_handlers
.as_deref()
.ok_or_else(|| vm.new_value_error("signal only works in main thread"))?;
if !vm.is_main_thread() {
return Err(vm.new_value_error("signal only works in main thread"));
}

let sig_handler =
match usize::try_from_borrowed_object(vm, &handler).ok() {
Expand All @@ -245,17 +246,15 @@ pub(crate) mod _signal {
siginterrupt(signalnum, 1);
}

let signal_handlers = vm.signal_handlers.get_or_init(signal::new_signal_handlers);
let old_handler = signal_handlers.borrow_mut()[signalnum as usize].replace(handler);
Ok(old_handler)
}

#[pyfunction]
fn getsignal(signalnum: i32, vm: &VirtualMachine) -> PyResult {
signal::assert_in_range(signalnum, vm)?;
let signal_handlers = vm
.signal_handlers
.as_deref()
.ok_or_else(|| vm.new_value_error("getsignal only works in main thread"))?;
let signal_handlers = vm.signal_handlers.get_or_init(signal::new_signal_handlers);
let handler = signal_handlers.borrow()[signalnum as usize]
.clone()
.unwrap_or_else(|| vm.ctx.none());
Expand Down Expand Up @@ -372,8 +371,8 @@ pub(crate) mod _signal {
#[cfg(not(windows))]
let fd = args.fd;

if vm.signal_handlers.is_none() {
return Err(vm.new_value_error("signal only works in main thread"));
if !vm.is_main_thread() {
return Err(vm.new_value_error("set_wakeup_fd only works in main thread"));
}

#[cfg(windows)]
Expand Down
23 changes: 17 additions & 6 deletions crates/vm/src/vm/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,7 @@ use crate::{
};
use alloc::{borrow::Cow, collections::BTreeMap};
use core::{
cell::{Cell, Ref, RefCell},
cell::{Cell, OnceCell, Ref, RefCell},
sync::atomic::{AtomicBool, Ordering},
};
use crossbeam_utils::atomic::AtomicCell;
Expand Down Expand Up @@ -81,7 +81,7 @@ pub struct VirtualMachine {
pub trace_func: RefCell<PyObjectRef>,
pub use_tracing: Cell<bool>,
pub recursion_limit: Cell<usize>,
pub(crate) signal_handlers: Option<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
pub(crate) signal_handlers: OnceCell<Box<RefCell<[Option<PyObjectRef>; signal::NSIG]>>>,
pub(crate) signal_rx: Option<signal::UserSignalReceiver>,
pub repr_guards: RefCell<HashSet<usize>>,
pub state: PyRc<PyGlobalState>,
Expand Down Expand Up @@ -148,6 +148,20 @@ pub fn process_hash_secret_seed() -> u32 {
}

impl VirtualMachine {
/// Check whether the current thread is the main thread.
/// Mirrors `_Py_ThreadCanHandleSignals`.
#[allow(dead_code)]
pub(crate) fn is_main_thread(&self) -> bool {
#[cfg(feature = "threading")]
{
crate::stdlib::thread::get_ident() == self.state.main_thread_ident.load()
}
#[cfg(not(feature = "threading"))]
{
true
}
}

/// Create a new `VirtualMachine` structure.
pub(crate) fn new(ctx: PyRc<Context>, state: PyRc<PyGlobalState>) -> Self {
flame_guard!("new VirtualMachine");
Expand All @@ -170,10 +184,7 @@ impl VirtualMachine {
let importlib = ctx.none();
let profile_func = RefCell::new(ctx.none());
let trace_func = RefCell::new(ctx.none());
let signal_handlers = Some(Box::new(
// putting it in a const optimizes better, prevents linear initialization of the array
const { RefCell::new([const { None }; signal::NSIG]) },
));
let signal_handlers = OnceCell::from(signal::new_signal_handlers());

let vm = Self {
builtins,
Expand Down
2 changes: 1 addition & 1 deletion crates/vm/src/vm/thread.rs
Original file line number Diff line number Diff line change
Expand Up @@ -275,7 +275,7 @@ impl VirtualMachine {
trace_func: RefCell::new(global_trace.unwrap_or_else(|| self.ctx.none())),
use_tracing: Cell::new(use_tracing),
recursion_limit: self.recursion_limit.clone(),
signal_handlers: None,
signal_handlers: core::cell::OnceCell::new(),
signal_rx: None,
repr_guards: RefCell::default(),
state: self.state.clone(),
Expand Down
Loading