Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads by youknowone · Pull Request #7317 · RustPython/RustPython · GitHub

Change the code field from PyMutex<PyRef> to PyAtomicRef, eliminating mutex lock/unlock on every function call. The setter uses swap_to_temporary_refs for safe deferred drop of the old code object.

Summary by CodeRabbit

Refactor
- Optimized internal code handling mechanisms in the function execution system.

Change the code field from PyMutex<PyRef<PyCode>> to PyAtomicRef<PyCode>, eliminating mutex lock/unlock on every function call. The setter uses swap_to_temporary_refs for safe deferred drop of the old code object.

📝 Walkthrough

Walkthrough

The pull request refactors PyFunction's code field from a mutex-protected reference to an atomic reference, eliminating lock-based access patterns across getters, setters, and JIT paths while maintaining thread-safe semantics.

Changes

Cohort / File(s)	Summary
Core PyFunction Implementation crates/vm/src/builtins/function.rs	Replaced PyMutex<PyRef<PyCode>> field with PyAtomicRef<PyCode>, updated initialization to use PyAtomicRef::from(), modified __code__ getter to return (*self.code).to_owned(), added VirtualMachine parameter to set___code__ method, and converted all code access from lock-based retrieval to direct atomic dereferencing across can_specialize_call, invoke_with_locals, invoke_exact_args, and error messaging paths.
JIT Code Access Patterns crates/vm/src/builtins/function/jit.rs	Replaced func.code.lock() calls with direct reference let code: &Py<PyCode> = &func.code in get_jit_arg_types and get_jit_args, removed explicit drop(code), and updated all code introspection operations to use direct atomic reference semantics.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related PRs

Fix PyCode constructor/replace #6193: Directly modifies PyFunction.code field storage mechanism with complementary synchronization strategy changes.
SetFunctionAttribute #5968: Related through overlapping PyFunction implementation modifications including field handling and initialization paths.

Suggested reviewers

ShaharNaveh

Poem

🐰✨ From mutex locks to atoms so bright,
The code now flows without holding tight,
No lock() calls, just references free,
Swift and lean, as refs should be!
A hoppy refactor, atomic and light! 🎉

🚥 Pre-merge checks | ✅ 3 ✅ Passed checks (3 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The pull request title 'Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads' accurately summarizes the main change: replacing a mutex-backed reference type with an atomic reference type to enable lock-free reads.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches 🧪 Generate unit tests (beta)

Create PR with unit tests
Post copyable unit tests in a comment

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@crates/vm/src/builtins/function.rs`:
- Around line 682-684: When assigning a new code object in set___code__, also
invalidate any cached JIT artifacts so the old machine code cannot be used;
specifically, after self.code.swap_to_temporary_refs(...) and before/after
self.func_version.store(0, Relaxed), clear the cached jitted_code (and any
associated owner/version fields used by the JIT fast path) so invoke_with_locals
cannot find/execute stale compiled code — update the fields that hold the
compiled entry (jitted_code) to a neutral/empty state consistent with how they
are checked in invoke_with_locals.

ℹ️ Review info

Configuration used: Path: .coderabbit.yml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 0a6a6f8 and 5defc42.

📒 Files selected for processing (2)

crates/vm/src/builtins/function.rs
crates/vm/src/builtins/function/jit.rs

⚠️ Potential issue | 🟠 Major

Invalidate cached JIT code when __code__ is reassigned.

Line 683 swaps the code object, but the cached jitted_code is not reset. The JIT fast path in invoke_with_locals (Line 533 onward) still executes cached compiled code when present, so f.__code__ = ... can leave the function running stale machine code.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@crates/vm/src/builtins/function.rs` around lines 682 - 684, When assigning a
new code object in set___code__, also invalidate any cached JIT artifacts so the
old machine code cannot be used; specifically, after
self.code.swap_to_temporary_refs(...) and before/after
self.func_version.store(0, Relaxed), clear the cached jitted_code (and any
associated owner/version fields used by the JIT fast path) so invoke_with_locals
cannot find/execute stale compiled code — update the fields that hold the
compiled entry (jitted_code) to a neutral/empty state consistent with how they
are checked in invoke_with_locals.

…ustPython#7317) Change the code field from PyMutex<PyRef<PyCode>> to PyAtomicRef<PyCode>, eliminating mutex lock/unlock on every function call. The setter uses swap_to_temporary_refs for safe deferred drop of the old code object.

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads

5defc42

Change the code field from PyMutex<PyRef<PyCode>> to PyAtomicRef<PyCode>, eliminating mutex lock/unlock on every function call. The setter uses swap_to_temporary_refs for safe deferred drop of the old code object.

youknowone marked this pull request as ready for review March 3, 2026 00:05

coderabbitai Bot reviewed Mar 3, 2026

View reviewed changes

youknowone merged commit 43e2df3 into RustPython:main Mar 3, 2026
13 checks passed

youknowone deleted the func-code-lockfree branch March 3, 2026 00:18

coderabbitai Bot mentioned this pull request Mar 4, 2026

Align specialization guards and caching with CPython #7341

Merged

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads by youknowone · Pull Request #7317 · RustPython/RustPython · GitHub

Uh oh!

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads - #7317

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads#7317
youknowone merged 1 commit into
RustPython:mainfrom
youknowone:func-code-lockfree

youknowone commented Mar 2, 2026 •
edited by coderabbitai Bot
Loading

Uh oh!

Uh oh!

coderabbitai Bot commented Mar 2, 2026 •
edited
Loading

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Uh oh!

coderabbitai Bot Mar 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

+                  fn set___code__(&self, code: PyRef<PyCode>, vm: &VirtualMachine) {
+                      self.code.swap_to_temporary_refs(code, vm);
                       self.func_version.store(0, Relaxed);

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads by youknowone · Pull Request #7317 · RustPython/RustPython · GitHub

Uh oh!

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads - #7317

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads#7317youknowone merged 1 commit intoRustPython:mainRustPython/RustPython:mainfrom youknowone:func-code-lockfreeyouknowone/RustPython:func-code-lockfreeCopy head branch name to clipboard

Conversation

youknowone commented Mar 2, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Mar 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested reviewers

Poem

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

coderabbitai Bot Mar 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Replace PyFunction.code PyMutex with PyAtomicRef for lock-free reads#7317
youknowone merged 1 commit into
RustPython:mainfrom
youknowone:func-code-lockfree

youknowone commented Mar 2, 2026 •
edited by coderabbitai Bot
Loading

coderabbitai Bot commented Mar 2, 2026 •
edited
Loading