Create weakrefs for subclasses with sentinel callback to properly handle them in GC

Author: sergey-miryanov

Include/cpython/weakrefobject.h

@@ -38,7 +38,6 @@ struct _PyWeakReference {
      */
     PyMutex *weakrefs_lock;
 #endif
-    uint8_t is_subclass;
 };
 
 PyAPI_FUNC(void) _PyWeakref_ClearRef(PyWeakReference *self);

Include/internal/pycore_interp_structs.h

@@ -701,6 +701,7 @@ struct _Py_interp_static_objects {
         _PyGC_Head_UNUSED _hamt_empty_gc_not_used;
         PyHamtObject hamt_empty;
         PyBaseExceptionObject last_resort_memory_error;
+        PyObject *subclasses_weakref_sentinel;
     } singletons;
 };
 

Include/internal/pycore_weakref.h

@@ -9,6 +9,7 @@ extern "C" {
 #endif
 
 #include "pycore_critical_section.h" // Py_BEGIN_CRITICAL_SECTION()
+#include "pycore_interp_structs.h"   // PyInterpreterState
 #include "pycore_lock.h"             // PyMutex_LockFlags()
 #include "pycore_object.h"           // _Py_REF_IS_MERGED()
 #include "pycore_pyatomic_ft_wrappers.h"
@@ -127,6 +128,10 @@ extern void _PyWeakref_ClearWeakRefsNoCallbacks(PyObject *obj);
 
 PyAPI_FUNC(int) _PyWeakref_IsDead(PyObject *weakref);
 
+int _PyWeakref_InitSubclassSentinel(PyInterpreterState *interp);
+PyObject * _PyWeakref_NewSubclassRef(PyObject *ob);
+int _PyWeakref_IsSubclassRef(PyWeakReference *weakref);
+
 #ifdef __cplusplus
 }
 #endif

Objects/typeobject.c

@@ -9259,15 +9259,12 @@ add_subclass(PyTypeObject *base, PyTypeObject *type)
     if (key == NULL)
         return -1;
 
-    PyObject *ref = PyWeakref_NewRef((PyObject *)type, NULL);
+    PyObject *ref = _PyWeakref_NewSubclassRef((PyObject *)type);
     if (ref == NULL) {
         Py_DECREF(key);
         return -1;
     }
 
-    ((PyWeakReference*)ref)->is_subclass = 1;
-
-
     // Only get tp_subclasses after creating the key and value.
     // PyWeakref_NewRef() can trigger a garbage collection which can execute
     // arbitrary Python code and so modify base->tp_subclasses.

Objects/weakrefobject.c

@@ -1,5 +1,7 @@
 #include "Python.h"
 #include "pycore_critical_section.h"
+#include "pycore_global_objects.h" // _Py_INTERP_STATIC_OBJECT()
+#include "pycore_interp_structs.h" // _PyInterpreterState_GET()
 #include "pycore_lock.h"
 #include "pycore_modsupport.h"    // _PyArg_NoKwnames()
 #include "pycore_object.h"        // _PyObject_GET_WEAKREFS_LISTPTR()
@@ -72,7 +74,6 @@ init_weakref(PyWeakReference *self, PyObject *ob, PyObject *callback)
     _PyObject_SetMaybeWeakref(ob);
     _PyObject_SetMaybeWeakref((PyObject *)self);
 #endif
-    self->is_subclass = 0;
 }
 
 // Clear the weakref and steal its callback into `callback`, if provided.
@@ -1133,3 +1134,55 @@ _PyWeakref_IsDead(PyObject *weakref)
 {
     return _PyWeakref_IS_DEAD(weakref);
 }
+
+int
+_PyWeakref_InitSubclassSentinel(PyInterpreterState *interp)
+{
+    PyObject *code =  (PyObject *)PyCode_NewEmpty(
+        "<generated>",
+        "<subclasses_weakref_sentinel>",
+        0);
+    if (code == NULL) {
+        return -1;
+    }
+
+    PyObject *globals = PyDict_New();
+    if (globals == NULL) {
+        Py_DECREF(code);
+        return -1;
+    }
+
+    PyObject *func = PyFunction_New(code, globals);
+    Py_DECREF(globals);
+    Py_DECREF(code);
+    if (func == NULL) {
+        return -1;
+    }
+
+    _Py_INTERP_SINGLETON(interp, subclasses_weakref_sentinel) = func;
+    return 0;
+}
+
+PyObject *
+_PyWeakref_NewSubclassRef(PyObject *ob)
+{
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    if (interp != interp->runtime->interpreters.main) {
+        interp = interp->runtime->interpreters.main;
+    }
+    PyObject *func = _Py_INTERP_SINGLETON(interp, subclasses_weakref_sentinel);
+    return PyWeakref_NewRef(ob, func);
+}
+
+int
+_PyWeakref_IsSubclassRef(PyWeakReference *weakref)
+{
+    assert(weakref != NULL);
+
+    PyInterpreterState *interp = _PyInterpreterState_GET();
+    if (interp != interp->runtime->interpreters.main) {
+        interp = interp->runtime->interpreters.main;
+    }
+    PyObject *func = _Py_INTERP_SINGLETON(interp, subclasses_weakref_sentinel);
+    return weakref->wr_callback == func;
+}

Python/gc.c

@@ -907,7 +907,7 @@ handle_weakrefs(PyGC_Head *unreachable, PyGC_Head *old)
              * will run and potentially cause a crash.  See bpo-38006 for
              * one example.
              */
-            if (!((PyWeakReference *)op)->is_subclass) {
+            if (!_PyWeakref_IsSubclassRef((PyWeakReference *)op)) {
                 _PyWeakref_ClearRef((PyWeakReference *)op);
             }
             else {
@@ -935,7 +935,7 @@ handle_weakrefs(PyGC_Head *unreachable, PyGC_Head *old)
             PyGC_Head *wrasgc;                  /* AS_GC(wr) */
 
             wr_next = wr->wr_next;
-            if (wr->is_subclass == 1) {
+            if (_PyWeakref_IsSubclassRef(wr) == 1) {
                 continue;
             }
 

Python/pylifecycle.c

@@ -1376,6 +1376,12 @@ init_interp_main(PyThreadState *tstate)
         return _PyStatus_ERR("failed to set builtin dict watcher");
     }
 
+    if (is_main_interp) {
+        if (_PyWeakref_InitSubclassSentinel(interp) < 0) {
+            return _PyStatus_ERR("failed to create subclasses weakref sentinel");
+        }
+    }
+
     assert(!_PyErr_Occurred(tstate));
 
     return _PyStatus_OK();

Python/pystate.c

@@ -760,6 +760,8 @@ interpreter_clear(PyInterpreterState *interp, PyThreadState *tstate)
 
     Py_CLEAR(interp->audit_hooks);
 
+    Py_CLEAR(_Py_INTERP_SINGLETON(interp, subclasses_weakref_sentinel));
+
     // At this time, all the threads should be cleared so we don't need atomic
     // operations for instrumentation_version or eval_breaker.
     interp->ceval.instrumentation_version = 0;