diff --git a/crates/stdlib/src/_asyncio.rs b/crates/stdlib/src/_asyncio.rs index 7e8ae5292..6e7a8c6e0 100644 --- a/crates/stdlib/src/_asyncio.rs +++ b/crates/stdlib/src/_asyncio.rs @@ -17,6 +17,7 @@ pub(crate) mod _asyncio { extend_module, function::{FuncArgs, KwArgs, OptionalArg, OptionalOption, PySetterValue}, protocol::PyIterReturn, + recursion::ReprGuard, types::{ Callable, Constructor, Destructor, Initializer, IterNext, Iterable, Representable, SelfIter, @@ -45,6 +46,15 @@ pub(crate) mod _asyncio { "_current_tasks" => current_tasks, }); + // Register fork handler to clear task state in child process + #[cfg(unix)] + { + let on_fork = vm + .get_attribute_opt(module.to_owned().into(), vm.ctx.intern_str("_on_fork"))? + .expect("_on_fork not found in _asyncio module"); + vm.state.after_forkers_child.lock().push(on_fork); + } + Ok(()) } @@ -830,9 +840,13 @@ pub(crate) mod _asyncio { impl Representable for PyFuture { fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let info = get_future_repr_info(zelf.as_object(), vm)?; let class_name = zelf.class().name().to_string(); - Ok(format!("<{} {}>", class_name, info)) + if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) { + let info = get_future_repr_info(zelf.as_object(), vm)?; + Ok(format!("<{} {}>", class_name, info)) + } else { + Ok(format!("<{} ...>", class_name)) + } } } @@ -1915,34 +1929,38 @@ pub(crate) mod _asyncio { fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult { let class_name = zelf.class().name().to_string(); - // Try to use _task_repr_info if available - if let Ok(info) = get_task_repr_info(zelf.as_object(), vm) - && info != "state=unknown" - { - return Ok(format!("<{} {}>", class_name, info)); + if let Some(_guard) = ReprGuard::enter(vm, zelf.as_object()) { + // Try to use _task_repr_info if available + if let Ok(info) = get_task_repr_info(zelf.as_object(), vm) + && info != "state=unknown" + { + return Ok(format!("<{} {}>", class_name, info)); + } + + // Fallback: build repr from task properties directly + let state = zelf.base.fut_state.load().as_str().to_lowercase(); + let name = zelf + .task_name + .read() + .as_ref() + .and_then(|n| n.str(vm).ok()) + .map(|s| s.as_str().to_string()) + .unwrap_or_else(|| "?".to_string()); + let coro_repr = zelf + .task_coro + .read() + .as_ref() + .and_then(|c| c.repr(vm).ok()) + .map(|s| s.as_str().to_string()) + .unwrap_or_else(|| "?".to_string()); + + Ok(format!( + "<{} {} name='{}' coro={}>", + class_name, state, name, coro_repr + )) + } else { + Ok(format!("<{} ...>", class_name)) } - - // Fallback: build repr from task properties directly - let state = zelf.base.fut_state.load().as_str().to_lowercase(); - let name = zelf - .task_name - .read() - .as_ref() - .and_then(|n| n.str(vm).ok()) - .map(|s| s.as_str().to_string()) - .unwrap_or_else(|| "?".to_string()); - let coro_repr = zelf - .task_coro - .read() - .as_ref() - .and_then(|c| c.repr(vm).ok()) - .map(|s| s.as_str().to_string()) - .unwrap_or_else(|| "?".to_string()); - - Ok(format!( - "<{} {} name='{}' coro={}>", - class_name, state, name, coro_repr - )) } } @@ -2203,9 +2221,12 @@ pub(crate) mod _asyncio { *task.task_fut_waiter.write() = Some(result.clone()); let task_obj: PyObjectRef = task.clone().into(); - let wakeup_wrapper = TaskWakeupMethWrapper::new(task_obj).into_ref(&vm.ctx); + let wakeup_wrapper = TaskWakeupMethWrapper::new(task_obj.clone()).into_ref(&vm.ctx); vm.call_method(&result, "add_done_callback", (wakeup_wrapper,))?; + // Track awaited_by relationship for introspection + future_add_to_awaited_by(result.clone(), task_obj, vm)?; + // If task_must_cancel is set, cancel the awaited future immediately // This propagates the cancellation through the future chain if task.task_must_cancel.load(Ordering::Relaxed) { @@ -2295,6 +2316,9 @@ pub(crate) mod _asyncio { .downcast() .map_err(|_| vm.new_type_error("task_wakeup called with non-Task object"))?; + // Remove awaited_by relationship before resuming + future_discard_from_awaited_by(fut.clone(), task.clone(), vm)?; + *task_ref.task_fut_waiter.write() = None; // Call result() on the awaited future to get either result or exception @@ -2376,7 +2400,6 @@ pub(crate) mod _asyncio { Some(l) if !vm.is_none(&l) => l, _ => { // When loop is None or not provided, use the running loop - // and raise an error if there's no running loop match vm.asyncio_running_loop.borrow().clone() { Some(l) => l, None => return Err(vm.new_runtime_error("no running event loop")), @@ -2384,6 +2407,23 @@ pub(crate) mod _asyncio { } }; + // Fast path: if the loop is the current thread's running loop, + // return the per-thread running task directly + let is_current_loop = vm + .asyncio_running_loop + .borrow() + .as_ref() + .is_some_and(|rl| rl.is(&loop_obj)); + + if is_current_loop { + return Ok(vm + .asyncio_running_task + .borrow() + .clone() + .unwrap_or_else(|| vm.ctx.none())); + } + + // Slow path: look up in the module-level dict for cross-thread queries let current_tasks = get_current_tasks_dict(vm)?; let dict: PyDictRef = current_tasks.downcast().unwrap(); @@ -2464,89 +2504,116 @@ pub(crate) mod _asyncio { #[pyfunction] fn _enter_task(loop_: PyObjectRef, task: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - let current_tasks = get_current_tasks_dict(vm)?; - let dict: PyDictRef = current_tasks.downcast().unwrap(); - - if dict.contains_key(&*loop_, vm) { - return Err( - vm.new_runtime_error("Cannot enter task while another task is already running") - ); + // Per-thread check, matching CPython's ts->asyncio_running_task + { + let running_task = vm.asyncio_running_task.borrow(); + if running_task.is_some() { + return Err(vm.new_runtime_error(format!( + "Cannot enter into task {:?} while another task {:?} is being executed.", + task, + running_task.as_ref().unwrap() + ))); + } } - dict.set_item(&*loop_, task, vm)?; + *vm.asyncio_running_task.borrow_mut() = Some(task.clone()); + + // Also update the module-level dict for cross-thread queries + if let Ok(current_tasks) = get_current_tasks_dict(vm) + && let Ok(dict) = current_tasks.downcast::() + { + let _ = dict.set_item(&*loop_, task, vm); + } Ok(()) } #[pyfunction] fn _leave_task(loop_: PyObjectRef, task: PyObjectRef, vm: &VirtualMachine) -> PyResult<()> { - let current_tasks = get_current_tasks_dict(vm)?; - let dict: PyDictRef = current_tasks.downcast().unwrap(); - - let current = dict.get_item(&*loop_, vm).ok(); - match current { - None => { - return Err( - vm.new_runtime_error("_leave_task: task is not the current task".to_owned()) - ); + // Per-thread check, matching CPython's ts->asyncio_running_task + { + let running_task = vm.asyncio_running_task.borrow(); + match running_task.as_ref() { + None => { + return Err(vm.new_runtime_error( + "_leave_task: task is not the current task".to_owned(), + )); + } + Some(current) if !current.is(&task) => { + return Err(vm.new_runtime_error( + "_leave_task: task is not the current task".to_owned(), + )); + } + _ => {} } - Some(current) if !current.is(&task) => { - return Err( - vm.new_runtime_error("_leave_task: task is not the current task".to_owned()) - ); - } - _ => {} } - dict.del_item(&*loop_, vm)?; + *vm.asyncio_running_task.borrow_mut() = None; + + // Also update the module-level dict + if let Ok(current_tasks) = get_current_tasks_dict(vm) + && let Ok(dict) = current_tasks.downcast::() + { + let _ = dict.del_item(&*loop_, vm); + } Ok(()) } #[pyfunction] fn _swap_current_task(loop_: PyObjectRef, task: PyObjectRef, vm: &VirtualMachine) -> PyResult { - let current_tasks = get_current_tasks_dict(vm)?; - let dict: PyDictRef = current_tasks.downcast().unwrap(); - - let prev = match dict.get_item(&*loop_, vm) { - Ok(t) => t, - Err(_) => vm.ctx.none(), - }; + // Per-thread swap, matching CPython's swap_current_task + let prev = vm + .asyncio_running_task + .borrow() + .clone() + .unwrap_or_else(|| vm.ctx.none()); if vm.is_none(&task) { - let _ = dict.del_item(&*loop_, vm); + *vm.asyncio_running_task.borrow_mut() = None; } else { - dict.set_item(&*loop_, task, vm)?; + *vm.asyncio_running_task.borrow_mut() = Some(task.clone()); + } + + // Also update the module-level dict for cross-thread queries + if let Ok(current_tasks) = get_current_tasks_dict(vm) + && let Ok(dict) = current_tasks.downcast::() + { + if vm.is_none(&task) { + let _ = dict.del_item(&*loop_, vm); + } else { + let _ = dict.set_item(&*loop_, task, vm); + } } Ok(prev) } + /// Reset task state after fork in child process. + #[pyfunction] + fn _on_fork(vm: &VirtualMachine) -> PyResult<()> { + // Clear current_tasks dict so child process doesn't inherit parent's tasks + if let Ok(current_tasks) = get_current_tasks_dict(vm) { + vm.call_method(¤t_tasks, "clear", ())?; + } + // Clear the running loop and task + *vm.asyncio_running_loop.borrow_mut() = None; + *vm.asyncio_running_task.borrow_mut() = None; + Ok(()) + } + #[pyfunction] fn future_add_to_awaited_by( fut: PyObjectRef, waiter: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - // Use optimized method for native Future/Task + // Only operate on native Future/Task objects (including subclasses). + // Non-native objects are silently ignored. if let Some(task) = fut.downcast_ref::() { return task.base.awaited_by_add(waiter, vm); } if let Some(future) = fut.downcast_ref::() { return future.awaited_by_add(waiter, vm); } - - // Fallback for non-native futures (Python subclasses) - if let Ok(Some(awaited_by)) = - vm.get_attribute_opt(fut.clone(), vm.ctx.intern_str("_asyncio_awaited_by")) - && !vm.is_none(&awaited_by) - { - vm.call_method(&awaited_by, "add", (waiter,))?; - return Ok(()); - } - - let new_set = PySet::default().into_ref(&vm.ctx); - new_set.add(waiter, vm)?; - fut.set_attr(vm.ctx.intern_str("_asyncio_awaited_by"), new_set, vm)?; - Ok(()) } @@ -2556,21 +2623,14 @@ pub(crate) mod _asyncio { waiter: PyObjectRef, vm: &VirtualMachine, ) -> PyResult<()> { - // Use optimized method for native Future/Task + // Only operate on native Future/Task objects (including subclasses). + // Non-native objects are silently ignored. if let Some(task) = fut.downcast_ref::() { return task.base.awaited_by_discard(&waiter, vm); } if let Some(future) = fut.downcast_ref::() { return future.awaited_by_discard(&waiter, vm); } - - // Fallback for non-native futures - if let Ok(Some(awaited_by)) = - vm.get_attribute_opt(fut.clone(), vm.ctx.intern_str("_asyncio_awaited_by")) - && !vm.is_none(&awaited_by) - { - vm.call_method(&awaited_by, "discard", (waiter,))?; - } Ok(()) } @@ -2596,6 +2656,14 @@ pub(crate) mod _asyncio { fn __self__(&self, vm: &VirtualMachine) -> PyObjectRef { self.task.read().clone().unwrap_or_else(|| vm.ctx.none()) } + + #[pygetset] + fn __qualname__(&self, vm: &VirtualMachine) -> PyResult> { + match self.task.read().as_ref() { + Some(t) => vm.get_attribute_opt(t.clone(), vm.ctx.intern_str("__qualname__")), + None => Ok(None), + } + } } impl Callable for TaskStepMethWrapper { @@ -2610,12 +2678,12 @@ pub(crate) mod _asyncio { } impl Representable for TaskStepMethWrapper { - fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let task = zelf.task.read().clone(); - match task { - Some(t) => Ok(t.repr(vm)?.as_str().to_string()), - None => Ok("".to_string()), - } + fn repr_str(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + Ok(format!( + "<{} object at {:#x}>", + zelf.class().name(), + zelf.get_id() + )) } } @@ -2634,6 +2702,14 @@ pub(crate) mod _asyncio { task: PyRwLock::new(Some(task)), } } + + #[pygetset] + fn __qualname__(&self, vm: &VirtualMachine) -> PyResult> { + match self.task.read().as_ref() { + Some(t) => vm.get_attribute_opt(t.clone(), vm.ctx.intern_str("__qualname__")), + None => Ok(None), + } + } } impl Callable for TaskWakeupMethWrapper { @@ -2648,12 +2724,12 @@ pub(crate) mod _asyncio { } impl Representable for TaskWakeupMethWrapper { - fn repr_str(zelf: &Py, vm: &VirtualMachine) -> PyResult { - let task = zelf.task.read().clone(); - match task { - Some(t) => Ok(t.repr(vm)?.as_str().to_string()), - None => Ok("".to_string()), - } + fn repr_str(zelf: &Py, _vm: &VirtualMachine) -> PyResult { + Ok(format!( + "<{} object at {:#x}>", + zelf.class().name(), + zelf.get_id() + )) } } diff --git a/crates/stdlib/src/_remote_debugging.rs b/crates/stdlib/src/_remote_debugging.rs new file mode 100644 index 000000000..57aa9876a --- /dev/null +++ b/crates/stdlib/src/_remote_debugging.rs @@ -0,0 +1,107 @@ +pub(crate) use _remote_debugging::module_def; + +#[pymodule] +mod _remote_debugging { + use crate::vm::{ + Py, PyObjectRef, PyResult, VirtualMachine, + builtins::PyType, + function::FuncArgs, + types::{Constructor, PyStructSequence}, + }; + + #[pystruct_sequence_data] + struct FrameInfoData { + filename: String, + lineno: i64, + funcname: String, + } + + #[pyattr] + #[pystruct_sequence( + name = "FrameInfo", + module = "_remote_debugging", + data = "FrameInfoData" + )] + struct FrameInfo; + + #[pyclass(with(PyStructSequence))] + impl FrameInfo {} + + #[pystruct_sequence_data] + struct TaskInfoData { + task_id: PyObjectRef, + task_name: PyObjectRef, + coroutine_stack: PyObjectRef, + awaited_by: PyObjectRef, + } + + #[pyattr] + #[pystruct_sequence(name = "TaskInfo", module = "_remote_debugging", data = "TaskInfoData")] + struct TaskInfo; + + #[pyclass(with(PyStructSequence))] + impl TaskInfo {} + + #[pystruct_sequence_data] + struct CoroInfoData { + call_stack: PyObjectRef, + task_name: PyObjectRef, + } + + #[pyattr] + #[pystruct_sequence(name = "CoroInfo", module = "_remote_debugging", data = "CoroInfoData")] + struct CoroInfo; + + #[pyclass(with(PyStructSequence))] + impl CoroInfo {} + + #[pystruct_sequence_data] + struct ThreadInfoData { + thread_id: PyObjectRef, + frame_info: PyObjectRef, + } + + #[pyattr] + #[pystruct_sequence( + name = "ThreadInfo", + module = "_remote_debugging", + data = "ThreadInfoData" + )] + struct ThreadInfo; + + #[pyclass(with(PyStructSequence))] + impl ThreadInfo {} + + #[pystruct_sequence_data] + struct AwaitedInfoData { + thread_id: PyObjectRef, + awaited_by: PyObjectRef, + } + + #[pyattr] + #[pystruct_sequence( + name = "AwaitedInfo", + module = "_remote_debugging", + data = "AwaitedInfoData" + )] + struct AwaitedInfo; + + #[pyclass(with(PyStructSequence))] + impl AwaitedInfo {} + + #[pyattr] + #[pyclass(name = "RemoteUnwinder", module = "_remote_debugging")] + #[derive(Debug, PyPayload)] + struct RemoteUnwinder {} + + impl Constructor for RemoteUnwinder { + type Args = FuncArgs; + + fn py_new(_cls: &Py, _args: Self::Args, vm: &VirtualMachine) -> PyResult { + Err(vm.new_not_implemented_error("_remote_debugging is not available".to_owned())) + } + } + + #[pyclass(with(Constructor))] + impl RemoteUnwinder {} +} diff --git a/crates/stdlib/src/lib.rs b/crates/stdlib/src/lib.rs index ae69f6f86..7fb79f6f4 100644 --- a/crates/stdlib/src/lib.rs +++ b/crates/stdlib/src/lib.rs @@ -12,6 +12,7 @@ extern crate alloc; pub(crate) mod macros; mod _asyncio; +mod _remote_debugging; pub mod array; mod binascii; mod bisect; @@ -113,6 +114,7 @@ pub fn stdlib_module_defs(ctx: &Context) -> Vec<&'static builtins::PyModuleDef> vec![ _asyncio::module_def(ctx), _opcode::module_def(ctx), + _remote_debugging::module_def(ctx), array::module_def(ctx), binascii::module_def(ctx), bisect::module_def(ctx), diff --git a/crates/vm/src/builtins/asyncgenerator.rs b/crates/vm/src/builtins/asyncgenerator.rs index 59f94ffb3..9903c9280 100644 --- a/crates/vm/src/builtins/asyncgenerator.rs +++ b/crates/vm/src/builtins/asyncgenerator.rs @@ -614,7 +614,7 @@ impl PyAnextAwaitable { await_method?.call((), vm)? } else { return Err(vm.new_type_error(format!( - "object {} can't be used in 'await' expression", + "'{}' object can't be awaited", wrapped.class().name() ))); } @@ -624,7 +624,7 @@ impl PyAnextAwaitable { await_method?.call((), vm)? } else { return Err(vm.new_type_error(format!( - "object {} can't be used in 'await' expression", + "'{}' object can't be awaited", wrapped.class().name() ))); } diff --git a/crates/vm/src/builtins/frame.rs b/crates/vm/src/builtins/frame.rs index 743b6e58a..94eaab544 100644 --- a/crates/vm/src/builtins/frame.rs +++ b/crates/vm/src/builtins/frame.rs @@ -151,6 +151,11 @@ impl Frame { #[pyclass] impl Py { + #[pygetset] + fn f_generator(&self) -> Option { + self.generator.lock().clone() + } + #[pygetset] pub fn f_back(&self, vm: &VirtualMachine) -> Option> { // TODO: actually store f_back inside Frame struct diff --git a/crates/vm/src/builtins/function.rs b/crates/vm/src/builtins/function.rs index 7eb74dec4..14c393f0b 100644 --- a/crates/vm/src/builtins/function.rs +++ b/crates/vm/src/builtins/function.rs @@ -529,13 +529,22 @@ impl Py { let is_coro = code.flags.contains(bytecode::CodeFlags::COROUTINE); match (is_gen, is_coro) { (true, false) => { - Ok(PyGenerator::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm)) + let obj = PyGenerator::new(frame.clone(), self.__name__(), self.__qualname__()) + .into_pyobject(vm); + frame.set_generator(obj.clone()); + Ok(obj) } (false, true) => { - Ok(PyCoroutine::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm)) + let obj = PyCoroutine::new(frame.clone(), self.__name__(), self.__qualname__()) + .into_pyobject(vm); + frame.set_generator(obj.clone()); + Ok(obj) } (true, true) => { - Ok(PyAsyncGen::new(frame, self.__name__(), self.__qualname__()).into_pyobject(vm)) + let obj = PyAsyncGen::new(frame.clone(), self.__name__(), self.__qualname__()) + .into_pyobject(vm); + frame.set_generator(obj.clone()); + Ok(obj) } (false, false) => vm.run_frame(frame), } diff --git a/crates/vm/src/frame.rs b/crates/vm/src/frame.rs index 4fc3cb260..1dacd61f6 100644 --- a/crates/vm/src/frame.rs +++ b/crates/vm/src/frame.rs @@ -85,6 +85,8 @@ pub struct Frame { pub trace_lines: PyMutex, pub trace_opcodes: PyMutex, pub temporary_refs: PyMutex>, + /// Back-reference to owning generator/coroutine/async generator + pub generator: PyMutex>, } impl PyPayload for Frame { @@ -112,6 +114,7 @@ unsafe impl Traverse for Frame { self.trace.traverse(tracer_fn); self.state.traverse(tracer_fn); self.temporary_refs.traverse(tracer_fn); + self.generator.traverse(tracer_fn); } } @@ -172,9 +175,14 @@ impl Frame { trace_lines: PyMutex::new(true), trace_opcodes: PyMutex::new(false), temporary_refs: PyMutex::new(vec![]), + generator: PyMutex::new(None), } } + pub fn set_generator(&self, generator: PyObjectRef) { + *self.generator.lock() = Some(generator); + } + pub fn current_location(&self) -> SourceLocation { self.code.locations[self.lasti() as usize - 1].0 } @@ -276,7 +284,21 @@ impl Py { } pub fn yield_from_target(&self) -> Option { - self.with_exec(|exec| exec.yield_from_target().map(PyObject::to_owned)) + // Use try_lock to avoid deadlock when the frame is currently executing. + // A running coroutine has no yield-from target. + let mut state = self.state.try_lock()?; + let exec = ExecutingFrame { + code: &self.code, + fastlocals: &self.fastlocals, + cells_frees: &self.cells_frees, + locals: &self.locals, + globals: &self.globals, + builtins: &self.builtins, + lasti: &self.lasti, + object: self, + state: &mut state, + }; + exec.yield_from_target().map(PyObject::to_owned) } pub fn is_internal_frame(&self) -> bool {