Skip to content

weakref: reuse proxy objects #8307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 4 commits into from
Jul 18, 2026
Merged

weakref: reuse proxy objects #8307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8a9e505
weakref: reuse callback-less proxy objects like CPython
rlaisqls Jul 17, 2026
a83b08a
weakref: always validate class in basic-proxy slot lookup
rlaisqls Jul 17, 2026
5150cdb
weakref: rustfmt find_generic_proxy_ptr
rlaisqls Jul 17, 2026
f7b36ac
weakref: unify reuse check shape
rlaisqls Jul 18, 2026
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
2 changes: 0 additions & 2 deletions Lib/test/test_weakref.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -289,7 +289,6 @@ def test_ref_reuse(self):
self.assertEqual(weakref.getweakrefcount(o), 1,
"wrong weak ref count for object after deleting proxy")

@unittest.expectedFailure # TODO: RUSTPYTHON
def test_proxy_reuse(self):
o = C()
proxy1 = weakref.proxy(o)
Expand Down Expand Up @@ -383,7 +382,6 @@ def __imatmul__(self, other):
def test_shared_ref_without_callback(self):
self.check_shared_without_callback(weakref.ref)

@unittest.expectedFailure # TODO: RUSTPYTHON
def test_shared_proxy_without_callback(self):
self.check_shared_without_callback(weakref.proxy)

Expand Down
2 changes: 0 additions & 2 deletions crates/capi/src/weakrefobject.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,8 +20,6 @@ pub unsafe extern "C" fn PyWeakref_GetRef(
let reference = unsafe { &*reference };
let upgraded = if let Some(weak) = reference.downcast_ref::<PyWeak>() {
weak.upgrade()
} else if let Some(proxy) = reference.downcast_ref::<PyWeakProxy>() {
proxy.get_weak().upgrade()
} else {
return Err(vm.new_type_error("expected a weakref"));
};
Expand Down
71 changes: 31 additions & 40 deletions crates/vm/src/builtins/weakproxy.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ use crate::{
Context, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, VirtualMachine, atomic_func,
class::PyClassImpl,
common::hash::PyHash,
function::{OptionalArg, PyArithmeticValue, PyComparisonValue, PySetterValue},
function::{FuncArgs, OptionalArg, PyArithmeticValue, PyComparisonValue, PySetterValue},
protocol::{PyIter, PyIterReturn, PyMappingMethods, PyNumberMethods, PySequenceMethods},
stdlib::builtins::reversed,
types::{
Expand All @@ -13,13 +13,22 @@ use crate::{
},
};

#[pyclass(module = false, name = "weakproxy", unhashable = true, traverse)]
#[pyclass(module = false, name = "weakproxy", unhashable = true)]
#[derive(Debug)]
pub struct PyWeakProxy {
weak: PyRef<PyWeak>,
}
#[repr(transparent)]
pub struct PyWeakProxy(PyWeak);

impl PyPayload for PyWeakProxy {
const PAYLOAD_TYPE_ID: core::any::TypeId = <PyWeak as PyPayload>::PAYLOAD_TYPE_ID;

#[inline]
unsafe fn validate_downcastable_from(obj: &PyObject) -> bool {
<Self as ::rustpython_vm::class::PyClassDef>::BASICSIZE <= obj.class().slots.basicsize
&& obj
.class()
.fast_issubclass(<Self as ::rustpython_vm::class::StaticType>::static_type())
}

#[inline]
fn class(ctx: &Context) -> &'static Py<PyType> {
ctx.types.weakproxy_type
Expand All @@ -37,52 +46,34 @@ pub struct WeakProxyNewArgs {
impl Constructor for PyWeakProxy {
type Args = WeakProxyNewArgs;

fn py_new(
_cls: &Py<PyType>,
Self::Args { referent, callback }: Self::Args,
vm: &VirtualMachine,
) -> PyResult<Self> {
let weak = Self::new_weak(referent.as_ref(), callback.into_option(), vm)?;
// TODO: PyWeakProxy should use the same payload as PyWeak
Ok(Self { weak })
fn slot_new(cls: PyTypeRef, args: FuncArgs, vm: &VirtualMachine) -> PyResult {
let _ = cls;
let Self::Args { referent, callback } = args.bind(vm)?;
let callback = callback
.into_option()
.filter(|callback| !vm.is_none(callback));
let proxy = Self::new_weakproxy(referent.as_ref(), callback, vm)?;
Ok(proxy.into())
}
}

crate::common::static_cell! {
static WEAK_SUBCLASS: PyTypeRef;
fn py_new(_cls: &Py<PyType>, _args: Self::Args, _vm: &VirtualMachine) -> PyResult<Self> {
unimplemented!("use slot_new")
}
}

impl PyWeakProxy {
fn new_weak(
referent: &PyObject,
callback: Option<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<PyRef<PyWeak>> {
// using an internal subclass as the class prevents us from getting the generic weakref,
// which would mess up the weakref count
let weak_cls = WEAK_SUBCLASS.get_or_init(|| {
vm.ctx.new_class(
None,
"__weakproxy",
vm.ctx.types.weakref_type.to_owned(),
super::PyWeak::make_slots(),
)
});
referent.downgrade_with_typ(callback, weak_cls.clone(), vm)
}

pub fn new_weakproxy(
referent: &PyObject,
callback: Option<PyObjectRef>,
vm: &VirtualMachine,
) -> PyResult<PyRef<Self>> {
let weak = Self::new_weak(referent, callback, vm)?;
Ok(Self { weak }.into_ref(&vm.ctx))
) -> PyResult<PyRef<PyWeak>> {
let typ = vm.ctx.types.weakproxy_type.to_owned();
referent.downgrade_with_typ(callback, typ, vm)
}

#[must_use]
pub fn get_weak(&self) -> &PyRef<PyWeak> {
&self.weak
pub fn get_weak(&self) -> &PyWeak {
&self.0
}
}

Expand All @@ -99,7 +90,7 @@ impl PyWeakProxy {
))]
impl PyWeakProxy {
fn try_upgrade(&self, vm: &VirtualMachine) -> PyResult {
self.weak.upgrade().ok_or_else(|| new_reference_error(vm))
self.0.upgrade().ok_or_else(|| new_reference_error(vm))
}

#[pymethod]
Expand Down
166 changes: 108 additions & 58 deletions crates/vm/src/object/core.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -582,6 +582,18 @@ unsafe fn unlink_weakref(wrl: &WeakRefList, node: NonNull<Py<PyWeak>>) {
}
}

// try_reuse_basic_ref
unsafe fn try_reuse_weakref(ptr: *mut Py<PyWeak>) -> Option<PyRef<PyWeak>> {
if ptr.is_null() {
return None;
}
let node = unsafe { &*ptr };
node.0
.ref_count
.safe_inc()
.then(|| unsafe { PyRef::from_raw(ptr) })
}

impl WeakRefList {
pub(super) fn new() -> Self {
Self {
Expand All @@ -596,22 +608,25 @@ impl WeakRefList {
obj: &PyObject,
cls: PyTypeRef,
cls_is_weakref: bool,
cls_is_weakproxy: bool,
callback: Option<PyObjectRef>,
dict: Option<PyDictRef>,
) -> PyRef<PyWeak> {
let is_generic = cls_is_weakref && callback.is_none();
let is_generic_proxy = cls_is_weakproxy && callback.is_none();

// Try reuse under lock first (fast path, no allocation)
{
let _lock = weakref_lock::lock(obj as *const PyObject as usize);
if is_generic {
let generic_ptr = self.generic.load(Ordering::Relaxed);
if !generic_ptr.is_null() {
let generic = unsafe { &*generic_ptr };
if generic.0.ref_count.safe_inc() {
return unsafe { PyRef::from_raw(generic_ptr) };
}
}
let existing = if is_generic {
unsafe { try_reuse_weakref(self.generic.load(Ordering::Relaxed)) }
} else if is_generic_proxy {
unsafe { try_reuse_weakref(self.find_generic_proxy_ptr()) }
} else {
None
};
if let Some(existing) = existing {
return existing;
}
}

Expand All @@ -629,64 +644,98 @@ impl WeakRefList {
// Re-acquire lock for linked list insertion
let _lock = weakref_lock::lock(obj as *const PyObject as usize);

// Re-check: another thread may have inserted a generic ref while we
// were allocating outside the lock. If so, reuse it and drop ours.
if is_generic {
let generic_ptr = self.generic.load(Ordering::Relaxed);
if !generic_ptr.is_null() {
let generic = unsafe { &*generic_ptr };
if generic.0.ref_count.safe_inc() {
// Nullify wr_object so drop_inner won't unlink an
// un-inserted node (which would corrupt the list head).
weak.wr_object.store(ptr::null_mut(), Ordering::Relaxed);
return unsafe { PyRef::from_raw(generic_ptr) };
}
}
// Re-check: another thread may have inserted a generic ref/proxy
// while we were allocating outside the lock. If so, reuse it and
// drop ours.
let existing = if is_generic {
unsafe { try_reuse_weakref(self.generic.load(Ordering::Relaxed)) }
} else if is_generic_proxy {
unsafe { try_reuse_weakref(self.find_generic_proxy_ptr()) }
} else {
None
};
if let Some(existing) = existing {
// Nullify wr_object so drop_inner won't unlink an
// un-inserted node (which would corrupt the list head).
weak.wr_object.store(ptr::null_mut(), Ordering::Relaxed);
return existing;
}

// Insert into linked list under stripe lock
// (insert_weakref: generic ref at head, generic proxy right after it)
let node_ptr = NonNull::from(&*weak);
let after = if is_generic {
None
} else if is_generic_proxy {
NonNull::new(self.generic.load(Ordering::Relaxed))
} else {
NonNull::new(self.find_generic_proxy_ptr())
.or_else(|| NonNull::new(self.generic.load(Ordering::Relaxed)))
};
Comment thread
coderabbitai[bot] marked this conversation as resolved.
match after {
Some(after) => unsafe { self.insert_after(after, node_ptr) },
None => unsafe { self.insert_at_head(node_ptr) },
}
if is_generic {
self.generic.store(node_ptr.as_ptr(), Ordering::Relaxed);
}

weak
}

unsafe fn insert_at_head(&self, node_ptr: NonNull<Py<PyWeak>>) {
unsafe {
let mut ptrs = WeakLink::pointers(node_ptr);
if is_generic {
// Generic ref goes to head (insert_head for basic ref)
let old_head = self.head.load(Ordering::Relaxed);
ptrs.as_mut().set_next(NonNull::new(old_head));
ptrs.as_mut().set_prev(None);
if let Some(old_head) = NonNull::new(old_head) {
WeakLink::pointers(old_head)
.as_mut()
.set_prev(Some(node_ptr));
}
self.head.store(node_ptr.as_ptr(), Ordering::Relaxed);
self.generic.store(node_ptr.as_ptr(), Ordering::Relaxed);
} else {
// Non-generic refs go after generic ref (insert_after)
let generic_ptr = self.generic.load(Ordering::Relaxed);
if let Some(after) = NonNull::new(generic_ptr) {
let after_next = WeakLink::pointers(after).as_ref().get_next();
ptrs.as_mut().set_prev(Some(after));
ptrs.as_mut().set_next(after_next);
WeakLink::pointers(after).as_mut().set_next(Some(node_ptr));
if let Some(next) = after_next {
WeakLink::pointers(next).as_mut().set_prev(Some(node_ptr));
}
let old_head = self.head.load(Ordering::Relaxed);
ptrs.as_mut().set_next(NonNull::new(old_head));
ptrs.as_mut().set_prev(None);
if let Some(old_head) = NonNull::new(old_head) {
WeakLink::pointers(old_head)
.as_mut()
.set_prev(Some(node_ptr));
}
self.head.store(node_ptr.as_ptr(), Ordering::Relaxed);
}
}

unsafe fn insert_after(&self, after: NonNull<Py<PyWeak>>, node_ptr: NonNull<Py<PyWeak>>) {
unsafe {
let mut ptrs = WeakLink::pointers(node_ptr);
let after_next = WeakLink::pointers(after).as_ref().get_next();
ptrs.as_mut().set_prev(Some(after));
ptrs.as_mut().set_next(after_next);
WeakLink::pointers(after).as_mut().set_next(Some(node_ptr));
if let Some(next) = after_next {
WeakLink::pointers(next).as_mut().set_prev(Some(node_ptr));
}
}
}

// get_basic_refs
fn find_generic_proxy_ptr(&self) -> *mut Py<PyWeak> {
let generic_ptr = self.generic.load(Ordering::Relaxed);
let candidate_ptr = if let Some(generic_node) = NonNull::new(generic_ptr) {
unsafe { WeakLink::pointers(generic_node).as_ref().get_next() }
.map_or(ptr::null_mut(), |n| n.as_ptr())
} else {
self.head.load(Ordering::Relaxed)
};
match NonNull::new(candidate_ptr) {
Some(candidate) => {
let node = unsafe { candidate.as_ref() };
let has_callback = unsafe { (&*node.0.payload.callback.get()).is_some() };
// PyWeakref_CheckProxy: the basic-proxy slot is reserved for
// the canonical proxy type; subclasses and callback-less ref
// subclasses must not be mistaken for it.
let is_proxy = node.class().is(crate::builtins::PyWeakProxy::static_type());
if has_callback || !is_proxy {
ptr::null_mut()
} else {
// No generic ref; insert at head
let old_head = self.head.load(Ordering::Relaxed);
ptrs.as_mut().set_next(NonNull::new(old_head));
ptrs.as_mut().set_prev(None);
if let Some(old_head) = NonNull::new(old_head) {
WeakLink::pointers(old_head)
.as_mut()
.set_prev(Some(node_ptr));
}
self.head.store(node_ptr.as_ptr(), Ordering::Relaxed);
candidate_ptr
}
}
None => ptr::null_mut(),
}

weak
}

/// Clear all weakrefs and call their callbacks.
Expand Down Expand Up @@ -1445,7 +1494,7 @@ impl PyObject {
typ: PyTypeRef,
) -> Option<PyRef<PyWeak>> {
self.weak_ref_list()
.map(|wrl| wrl.add(self, typ, true, callback, None))
.map(|wrl| wrl.add(self, typ, true, false, callback, None))
}

pub(crate) fn downgrade_with_typ(
Expand Down Expand Up @@ -1476,13 +1525,14 @@ impl PyObject {
None
};
let cls_is_weakref = typ.is(vm.ctx.types.weakref_type);
let cls_is_weakproxy = typ.is(vm.ctx.types.weakproxy_type);
let wrl = self.weak_ref_list().ok_or_else(|| {
vm.new_type_error(format!(
"cannot create weak reference to '{}' object",
self.class().name()
))
})?;
Ok(wrl.add(self, typ, cls_is_weakref, callback, dict))
Ok(wrl.add(self, typ, cls_is_weakref, cls_is_weakproxy, callback, dict))
}

pub fn downgrade(
Expand Down
Loading