Use weakrefs

Author: XuehaiPan

include/pybind11/detail/class.h

@@ -11,10 +11,16 @@
 
 #include <pybind11/attr.h>
 #include <pybind11/options.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
 
 #include "exception_translation.h"
 
 PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
+
+class weakref;
+class cpp_function;
+
 PYBIND11_NAMESPACE_BEGIN(detail)
 
 #if !defined(PYPY_VERSION)
@@ -252,11 +258,6 @@ extern "C" inline void pybind11_meta_dealloc(PyObject *obj) {
     });
 
     PyType_Type.tp_dealloc(obj);
-
-    // Release the references to the internals capsules that were acquired in make_new_python_type.
-    // See the comment there for details on preventing use-after-free during interpreter shutdown.
-    Py_XDECREF(get_internals_capsule());
-    Py_XDECREF(get_local_internals_capsule());
 }
 
 /** This metaclass is assigned by default to all pybind11 types and is required in order
@@ -834,15 +835,6 @@ inline PyObject *make_new_python_type(const type_record &rec) {
 
     PYBIND11_SET_OLDPY_QUALNAME(type, qualname);
 
-    // Prevent use-after-free during interpreter shutdown. GC order is not guaranteed, so the
-    // internals capsule may be destroyed (resetting internals via internals_shutdown) before all
-    // pybind11 types are destroyed. If a type's tp_traverse/tp_clear then calls py::cast, it
-    // would recreate an empty internals and fail because the type registry is gone. By holding
-    // references to the capsules, we ensure they outlive all pybind11 types. The decref happens
-    // in pybind11_meta_dealloc.
-    Py_XINCREF(get_internals_capsule());
-    Py_XINCREF(get_local_internals_capsule());
-
     return reinterpret_cast<PyObject *>(type);
 }
 

include/pybind11/pybind11.h

@@ -1692,6 +1692,30 @@ class generic_type : public object {
 
         m_ptr = make_new_python_type(rec);
 
+        // Prevent use-after-free during interpreter shutdown. GC order is not guaranteed, so the
+        // internals capsule may be destroyed (resetting internals via internals_shutdown) before
+        // all pybind11 types are destroyed. If a type's tp_traverse/tp_clear then calls py::cast,
+        // it would recreate an empty internals and fail because the type registry is gone. By
+        // holding references to the capsules, we ensure they outlive all pybind11 types.
+        // We use weakrefs on the type with a cpp_function callback. When the type is destroyed,
+        // Python will call the callback which releases the capsule reference and the weakref.
+        if (PyObject *capsule = get_internals_capsule()) {
+            Py_INCREF(capsule);
+            (void) weakref(handle(m_ptr), cpp_function([](handle prevent_release) -> void {
+                               Py_XDECREF(get_internals_capsule());
+                               prevent_release.dec_ref();
+                           }))
+                .release();
+        }
+        if (PyObject *capsule = get_local_internals_capsule()) {
+            Py_INCREF(capsule);
+            (void) weakref(handle(m_ptr), cpp_function([](handle prevent_release) -> void {
+                               Py_XDECREF(get_local_internals_capsule());
+                               prevent_release.dec_ref();
+                           }))
+                .release();
+        }
+
         /* Register supplemental type information in C++ dict */
         auto *tinfo = new detail::type_info();
         tinfo->type = reinterpret_cast<PyTypeObject *>(m_ptr);