Make weakrefs to immutable objects thread-safe

Include/cpython/weakrefobject.h

@@ -16,6 +16,11 @@ struct _PyWeakReference {
 
     /* A callable to invoke when wr_object dies, or NULL if none. */
     PyObject *wr_callback;
+    /* ID of the interpreter where the callback resides.
+     * Used for immutable objects to know which interpreter to call back into.
+     * This is -1 if the callback is NULL.
+     */
+    int64_t callback_ipid;
 
     /* A cache for wr_object's hash code.  As usual for hashes, this is -1
      * if the hash code isn't known yet.

Include/internal/pycore_object.h

@@ -813,6 +813,7 @@ _Py_TryIncref(PyObject *op)
 #ifdef Py_GIL_DISABLED
     return _Py_TryIncrefFast(op) || _Py_TryIncRefShared(op);
 #else
+    assert(!_Py_IsImmutable(op) && "Use _Py_TryIncref_Immutable for immutable objects");
     if (Py_REFCNT(op) > 0) {
         Py_INCREF(op);
         return 1;
@@ -821,6 +822,9 @@ _Py_TryIncref(PyObject *op)
 #endif
 }
 
+int _Py_TryIncref_Immutable(PyObject *op);
+int _Py_IsDead_Immutable(PyObject *op);
+
 // Enqueue an object to be freed possibly after some delay
 #ifdef Py_GIL_DISABLED
 PyAPI_FUNC(void) _PyObject_XDecRefDelayed(PyObject *obj);

Include/internal/pycore_weakref.h

@@ -37,11 +37,14 @@ extern "C" {
 
 #else
 
-#define LOCK_WEAKREFS(obj)
-#define UNLOCK_WEAKREFS(obj)
+// Lock used for weakrefs to immutable objects
+extern PyMutex _PyWeakref_Lock;
 
-#define LOCK_WEAKREFS_FOR_WR(wr)
-#define UNLOCK_WEAKREFS_FOR_WR(wr)
+#define LOCK_WEAKREFS(obj) PyMutex_LockFlags(&_PyWeakref_Lock, _Py_LOCK_DONT_DETACH)
+#define UNLOCK_WEAKREFS(obj) PyMutex_Unlock(&_PyWeakref_Lock)
+
+#define LOCK_WEAKREFS_FOR_WR(wr) PyMutex_LockFlags(&_PyWeakref_Lock, _Py_LOCK_DONT_DETACH)
+#define UNLOCK_WEAKREFS_FOR_WR(wr) PyMutex_Unlock(&_PyWeakref_Lock)
 
 #define FT_CLEAR_WEAKREFS(obj, weakref_list)        \
     do {                                            \
@@ -65,35 +68,59 @@ static inline int _is_dead(PyObject *obj)
     Py_ssize_t shared = _Py_atomic_load_ssize_relaxed(&obj->ob_ref_shared);
     return shared == _Py_REF_SHARED(0, _Py_REF_MERGED);
 #else
-    return (Py_REFCNT(obj) == 0);
+    if (_Py_IsImmutable(obj)) {
+        return _Py_IsDead_Immutable(obj);
+    }
+    else {
+        return (Py_REFCNT(obj) == 0);
+    }
 #endif
 }
 
+static inline PyObject* get_ref_lock_held(PyWeakReference *ref, PyObject *obj)
+{
+#ifdef Py_GIL_DISABLED
+    // Need to check again because the object could have been deallocated
+    if (ref->wr_object == Py_None) {
+        // clear_weakref() was called
+        return NULL;
+    }
+    if (_Py_TryIncref(obj)) {
+        return obj;
+    }
+#else
+    if (_Py_IsImmutable(obj)) {
+        // Need to check again because the object could have been deallocated
+        if (ref->wr_object == Py_None) {
+            // clear_weakref() was called
+            return NULL;
+        }
+        if (_Py_TryIncref_Immutable(obj)) {
+            return obj;
+        }
+    }
+    else if (_Py_TryIncref(obj)) {
+        return obj;
+    }
+#endif
+    return NULL;
+}
+
 static inline PyObject* _PyWeakref_GET_REF(PyObject *ref_obj)
 {
     assert(PyWeakref_Check(ref_obj));
     PyWeakReference *ref = _Py_CAST(PyWeakReference*, ref_obj);
 
-    PyObject *obj = FT_ATOMIC_LOAD_PTR(ref->wr_object);
+    PyObject *obj = _Py_atomic_load_ptr(&ref->wr_object);
     if (obj == Py_None) {
         // clear_weakref() was called
         return NULL;
     }
 
     LOCK_WEAKREFS(obj);
-#ifdef Py_GIL_DISABLED
-    if (ref->wr_object == Py_None) {
-        // clear_weakref() was called
-        UNLOCK_WEAKREFS(obj);
-        return NULL;
-    }
-#endif
-    if (_Py_TryIncref(obj)) {
-        UNLOCK_WEAKREFS(obj);
-        return obj;
-    }
+    PyObject* result = get_ref_lock_held(ref, obj);
     UNLOCK_WEAKREFS(obj);
-    return NULL;
+    return result;
 }
 
 static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
@@ -108,12 +135,9 @@ static inline int _PyWeakref_IS_DEAD(PyObject *ref_obj)
     }
     else {
         LOCK_WEAKREFS(obj);
+        // Immutable objects and free-threaded builds require a new check
         // See _PyWeakref_GET_REF() for the rationale of this test
-#ifdef Py_GIL_DISABLED
         ret = (ref->wr_object == Py_None) || _is_dead(obj);
-#else
-        ret = _is_dead(obj);
-#endif
         UNLOCK_WEAKREFS(obj);
     }
     return ret;
@@ -125,6 +149,8 @@ extern Py_ssize_t _PyWeakref_GetWeakrefCount(PyObject *obj);
 // intact.
 extern void _PyWeakref_ClearWeakRefsNoCallbacks(PyObject *obj);
 
+PyAPI_FUNC(void) _PyWeakref_OnObjectFreeze(PyObject *object);
+PyAPI_FUNC(void) _PyImmutability_ClearWeakRefs(PyObject *object, PyWeakReference **callbacks);
 PyAPI_FUNC(int) _PyWeakref_IsDead(PyObject *weakref);
 
 #ifdef __cplusplus

Lib/test/test_freeze/test_core.py

@@ -1,5 +1,6 @@
 import unittest
 from immutable import freeze, NotFreezable, isfrozen
+import weakref
 
 from .test_common import BaseNotFreezableTest, BaseObjectTest
 
@@ -445,12 +446,11 @@ class B:
     class C:
         # Function that takes a object, and stores it in a weakref field.
         def __init__(self, obj):
-            import weakref
             self.obj = weakref.ref(obj)
         def val(self):
             return self.obj()
 
-    def test_weakref(self):
+    def test_freeze_object_with_weakref(self):
         obj = TestWeakRef.B()
         c = TestWeakRef.C(obj)
         freeze(c)
@@ -461,6 +461,30 @@ def test_weakref(self):
         # The reference should remain as it was reachable through a frozen weakref.
         self.assertTrue(c.val() is not None)
 
+    def test_weakref_to_frozen_object(self):
+        obj = TestWeakRef.B()
+        freeze(obj)
+        c = TestWeakRef.C(obj)
+        self.assertTrue(isfrozen(c.val()))
+        # TODO(Immutable): Do we want to freeze the weakref?
+        # self.assertTrue(isfrozen(c.obj))
+
+    def test_remove_weakref(self):
+        obj = TestWeakRef.B()
+        c = TestWeakRef.C(obj)
+        freeze(c)
+        obj = None
+        # The reference should be removed
+        self.assertTrue(weakref.getweakrefcount(obj) == 0)
+
+    def test_reuse_weakref(self):
+        obj = TestWeakRef.B()
+        freeze(obj)
+        c1 = TestWeakRef.C(obj)
+        c2 = TestWeakRef.C(obj)
+        # The weakrefs should be the same, as they refer to the same object
+        self.assertTrue(c1.obj is c2.obj)
+
 class TestStackCapture(unittest.TestCase):
      def test_stack_capture(self):
          import sys

Modules/_weakref.c

@@ -88,7 +88,9 @@ _weakref_getweakrefs(PyObject *module, PyObject *object)
     PyWeakReference *current = *GET_WEAKREFS_LISTPTR(object);
     while (current != NULL) {
         PyObject *curobj = (PyObject *) current;
-        if (_Py_TryIncref(curobj)) {
+        int incref_res = _Py_IsImmutable(curobj) ?
+            _Py_TryIncref_Immutable(curobj) : _Py_TryIncref(curobj);
+        if (incref_res) {
             if (PyList_Append(result, curobj)) {
                 UNLOCK_WEAKREFS(object);
                 Py_DECREF(curobj);