Make PyWeakref_NewProxy thread-safe

mpage · mpage · commit c4f845339746 · 2024-03-22T12:29:13.000-07:00
diff --git a/Objects/weakrefobject.c b/Objects/weakrefobject.c
@@ -854,15 +854,26 @@ PyWeakref_NewProxy(PyObject *ob, PyObject *callback)
         return NULL;
     }
     list = GET_WEAKREFS_LISTPTR(ob);
+    /* NB: For free-threaded builds its critical that no code inside the
+     * critical section can release it. We need to recompute ref/proxy after
+     * any code that may release the critical section.
+     */
+    Py_BEGIN_CRITICAL_SECTION(ob);
     get_basic_refs(*list, &ref, &proxy);
     if (callback == Py_None)
         callback = NULL;
     if (callback == NULL)
         /* attempt to return an existing weak reference if it exists */
         result = proxy;
+#ifdef Py_GIL_DISABLED
+    /* Incref will fail if the existing weakref is being destroyed */
+    if (result == NULL ||
+        !_Py_TryIncref((PyObject **)&result, (PyObject *)result)) {
+#else
     if (result != NULL)
         Py_INCREF(result);
     else {
+#endif
         /* We do not need to recompute ref/proxy; new_weakref cannot
            trigger GC.
         */
@@ -888,6 +899,7 @@ PyWeakref_NewProxy(PyObject *ob, PyObject *callback)
                 insert_after(result, prev);
         }
     }
+    Py_END_CRITICAL_SECTION();
     return (PyObject *) result;
 }
 
