sys.set_asyncgen_hook #6439
sys.set_asyncgen_hook #6439
Conversation
WalkthroughThe changes introduce lifecycle and state management patterns to async generators and coroutines, adding hook initialization, finalizer invocation, and closed-state tracking to prevent reuse after completion. Additionally, await() result validation is added to enforce iterator protocol compliance. Changes
Sequence DiagramsequenceDiagram
participant App as Application
participant VM as VirtualMachine
participant AG as AsyncGenerator
participant Hooks as Thread-local Hooks
participant Awaitable as PyAnextAwaitable
App->>AG: Call __anext__()
activate AG
AG->>VM: init_hooks(self, vm)
activate VM
VM->>Hooks: Read finalizer & firstiter
Hooks-->>VM: Return hooks
VM->>AG: Store finalizer in ag_finalizer
VM->>VM: Call firstiter (if present)
VM-->>AG: Hooks initialized
deactivate VM
AG->>Awaitable: Create new PyAnextAwaitable(state=Init)
AG-->>App: Return Awaitable
deactivate AG
App->>Awaitable: Call send(val)
activate Awaitable
Awaitable->>Awaitable: check_closed()?
alt Not Closed
Awaitable->>Awaitable: state = Iter
Awaitable-->>App: Yield value
else Closed
Awaitable-->>App: RuntimeError
end
deactivate Awaitable
App->>Awaitable: Call close()
activate Awaitable
Awaitable->>Awaitable: state = Closed
Awaitable-->>App: Closed
deactivate Awaitable
Note over AG,Awaitable: On AsyncGen finalization
AG->>AG: call_finalizer(self, vm)
AG->>VM: Invoke stored finalizer
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes
Possibly related PRs
Poem
Pre-merge checks and finishing touches✅ Passed checks (3 passed)
✨ Finishing touches
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (3)
crates/vm/src/frame.rs (1)
921-922: Remove redundant local import.
PyIteris already imported at line 15 in the file's global imports block (protocol::{PyIter, PyIterReturn}). This local import is unnecessary.Apply this diff:
bytecode::Instruction::GetAwaitable => { - use crate::protocol::PyIter; - let awaited_obj = self.pop_value();crates/vm/src/builtins/asyncgenerator.rs (2)
588-595: Inconsistency:PyAnextAwaitabledoesn't close on error unlikePyAsyncGenASend.In
PyAsyncGenASend::send(lines 299-301), when the result is an error,self.close()is called. However,PyAnextAwaitable::senddoesn't close on error, only on explicitclose()call.This inconsistency could lead to different behavior when errors occur. Consider aligning the behavior:
fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult { self.check_closed(vm)?; self.state.store(AwaitableState::Iter); let awaitable = self.get_awaitable_iter(vm)?; let result = vm.call_method(&awaitable, "send", (val,)); - self.handle_result(result, vm) + let res = self.handle_result(result, vm); + if res.is_err() { + self.state.store(AwaitableState::Closed); + } + res }
597-618: Same inconsistency inthrow- should close on error.Similar to
send, thethrowmethod doesn't mark the state as closed when an error occurs, unlikePyAsyncGenASend::throw.fn throw( &self, exc_type: PyObjectRef, exc_val: OptionalArg, exc_tb: OptionalArg, vm: &VirtualMachine, ) -> PyResult { self.check_closed(vm)?; self.state.store(AwaitableState::Iter); let awaitable = self.get_awaitable_iter(vm)?; let result = vm.call_method( &awaitable, "throw", ( exc_type, exc_val.unwrap_or_none(vm), exc_tb.unwrap_or_none(vm), ), ); - self.handle_result(result, vm) + let res = self.handle_result(result, vm); + if res.is_err() { + self.state.store(AwaitableState::Closed); + } + res }
📜 Review details
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro
⛔ Files ignored due to path filters (1)
Lib/test/test_asyncgen.pyis excluded by!Lib/**
📒 Files selected for processing (3)
crates/vm/src/builtins/asyncgenerator.rs(13 hunks)crates/vm/src/builtins/coroutine.rs(5 hunks)crates/vm/src/frame.rs(2 hunks)
🧰 Additional context used
📓 Path-based instructions (1)
**/*.rs
📄 CodeRabbit inference engine (.github/copilot-instructions.md)
**/*.rs: Follow the default rustfmt code style by runningcargo fmtto format Rust code
Always run clippy to lint Rust code (cargo clippy) before completing tasks and fix any warnings or lints introduced by changes
Follow Rust best practices for error handling and memory management
Use the macro system (pyclass,pymodule,pyfunction, etc.) when implementing Python functionality in Rust
Files:
crates/vm/src/frame.rscrates/vm/src/builtins/coroutine.rscrates/vm/src/builtins/asyncgenerator.rs
🧬 Code graph analysis (1)
crates/vm/src/builtins/coroutine.rs (1)
crates/vm/src/coroutine.rs (3)
new(51-60)send(92-132)close(150-170)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Run snippets and cpython tests (macos-latest)
🔇 Additional comments (8)
crates/vm/src/frame.rs (1)
937-945: LGTM - iterator validation for__await__()result.The validation correctly enforces that
__await__()must return an iterator, raising aTypeErrorwith an appropriate message if not. This aligns with Python's async protocol requirements.crates/vm/src/builtins/coroutine.rs (3)
60-65: LGTM - closed-state initialization for coroutine wrapper.The wrapper now correctly initializes with
closed: falseand tracks whether the coroutine has been awaited to prevent reuse.
159-164: LGTM - reuse prevention check.The
check_closedmethod correctly prevents reusing an already-awaited coroutine with an appropriate error message.
195-198: LGTM - close marks wrapper as closed first.Setting
closedbefore delegating to the underlying coroutine ensures the wrapper cannot be reused even if the close operation fails.crates/vm/src/builtins/asyncgenerator.rs (4)
50-74: LGTM - async generator hook initialization.The
init_hooksmethod correctly implements one-time initialization of async generator hooks by:
- Checking if already initialized (idempotent)
- Storing the finalizer from thread-local storage
- Calling the firstiter hook if set
This aligns with CPython's
async_gen_init_hooksbehavior.
141-148: LGTM - hook initialization in__anext__.Correctly initializes hooks before creating the async send object, ensuring firstiter callback is invoked on first iteration.
524-529: Correct error message forPyAnextAwaitable::check_closed.The error message correctly identifies the reuse scenario for
__anext__()/asend()context.
76-88: Dead code added for future deallocation support.The
call_finalizermethod is marked as dead code but appears to be scaffolding for implementing proper async generator finalization during deallocation (per the "= gen_dealloc" comment). This is acceptable for incremental development.Note: Lines 80-81 use
let_chainssyntax (if let Some(x) = y && cond), which requires Rust 1.76+. The project's MSRV is 1.89.0, so this is fully supported.
| #[pymethod] | ||
| fn send(&self, val: PyObjectRef, vm: &VirtualMachine) -> PyResult<PyIterReturn> { | ||
| self.coro.send(val, vm) | ||
| self.check_closed(vm)?; | ||
| let result = self.coro.send(val, vm); | ||
| // Mark as closed if exhausted | ||
| if let Ok(PyIterReturn::StopIteration(_)) = &result { | ||
| self.closed.store(true); | ||
| } | ||
| result | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Check if there's a pattern in the codebase for closing on errors
rg -n "closed\.store\(true\)" --type rust -C3Repository: RustPython/RustPython
Length of output: 4621
🏁 Script executed:
#!/bin/bash
# Get the full coroutine.rs file to understand the structure
wc -l crates/vm/src/builtins/coroutine.rsRepository: RustPython/RustPython
Length of output: 104
🏁 Script executed:
#!/bin/bash
# Find the Coro struct and understand its send method
rg -n "impl.*Coro" --type rust -A 5 crates/vm/src/builtins/Repository: RustPython/RustPython
Length of output: 3678
🏁 Script executed:
#!/bin/bash
# Check if there are other methods that handle closing and errors
rg -n "fn.*(&self.*vm.*PyResult" crates/vm/src/builtins/coroutine.rs -A 10Repository: RustPython/RustPython
Length of output: 144
🏁 Script executed:
#!/bin/bash
# Get the Coro struct and send method implementation
sed -n '40,120p' crates/vm/src/coroutine.rs | head -80Repository: RustPython/RustPython
Length of output: 2533
🏁 Script executed:
#!/bin/bash
# Find the send method in the Coro struct
rg -n "fn send" crates/vm/src/coroutine.rs -A 15Repository: RustPython/RustPython
Length of output: 647
🏁 Script executed:
#!/bin/bash
# See the full context of the PyCoroutineWrapper send and throw methods
sed -n '166,200p' crates/vm/src/builtins/coroutine.rsRepository: RustPython/RustPython
Length of output: 1028
🏁 Script executed:
#!/bin/bash
# Look at the rest of Coro.send to see full error handling
sed -n '92,125p' crates/vm/src/coroutine.rsRepository: RustPython/RustPython
Length of output: 1455
🏁 Script executed:
#!/bin/bash
# Verify throw method in Coro as well to see if it uses maybe_close
rg -n "pub fn throw" crates/vm/src/coroutine.rs -A 20Repository: RustPython/RustPython
Length of output: 866
Mark wrapper as closed on errors, not just StopIteration.
The wrapper's closed flag only closes when StopIteration is returned. When send() returns an error, the wrapper remains "open", allowing another call that will unnecessarily pass through check_closed() before the underlying Coro (which already closed itself via maybe_close) returns StopIteration. Close the wrapper on errors to match the underlying Coro's error-handling pattern and align with CPython behavior where a coroutine cannot resume after raising an exception.
🤖 Prompt for AI Agents
In crates/vm/src/builtins/coroutine.rs around lines 166 to 175, the wrapper's
closed flag is only set when send() returns Ok(PyIterReturn::StopIteration(_)),
but it must also be marked closed when send() returns an Err so the wrapper
cannot be resumed after the underlying coroutine has raised; update the logic to
set self.closed.store(true) both for Ok(PyIterReturn::StopIteration(_)) and for
Err(_) (i.e., treat any error as terminal), preserving the original return value
so the error or StopIteration still propagates to the caller.
Summary by CodeRabbit
__await__returns invalid types✏️ Tip: You can customize this high-level summary in your review settings.