Updates for the new proposal.

ZeroIntensity · ZeroIntensity · commit d66157883d85 · 2025-05-27T21:02:39.000-04:00
diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h
@@ -270,32 +270,39 @@ PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameFunc(
 
 typedef uintptr_t PyInterpreterRef;
 
-PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Get(void);
+PyAPI_FUNC(int) PyInterpreterRef_Get(PyInterpreterRef *ptr);
 PyAPI_FUNC(PyInterpreterRef) PyInterpreterRef_Dup(PyInterpreterRef ref);
-PyAPI_FUNC(int) PyInterpreterState_AsStrong(PyInterpreterState *interp, PyInterpreterRef *strong_ptr);
+PyAPI_FUNC(int) PyInterpreterRef_Main(PyInterpreterRef *strong_ptr);
 PyAPI_FUNC(void) PyInterpreterRef_Close(PyInterpreterRef ref);
 PyAPI_FUNC(PyInterpreterState *) PyInterpreterRef_AsInterpreter(PyInterpreterRef ref);
 
 #define PyInterpreterRef_Close(ref) do {    \
     PyInterpreterRef_Close(ref);            \
     ref = 0;                                \
-} while (0);                                \
+} while (0);
 
 /* Weak interpreter references */
 
 typedef struct _interpreter_weakref {
     int64_t id;
     Py_ssize_t refcount;
-} PyInterpreterWeakRef;
+} _PyInterpreterWeakRef;
 
-PyAPI_FUNC(PyInterpreterWeakRef) PyInterpreterWeakRef_Get(void);
+typedef _PyInterpreterWeakRef *PyInterpreterWeakRef;
+
+PyAPI_FUNC(int) PyInterpreterWeakRef_Get(PyInterpreterWeakRef *ptr);
 PyAPI_FUNC(PyInterpreterWeakRef) PyInterpreterWeakRef_Dup(PyInterpreterWeakRef wref);
 PyAPI_FUNC(int) PyInterpreterWeakRef_AsStrong(PyInterpreterWeakRef wref, PyInterpreterRef *strong_ptr);
 PyAPI_FUNC(void) PyInterpreterWeakRef_Close(PyInterpreterWeakRef wref);
 
+#define PyInterpreterWeakRef_Close(ref) do {    \
+    PyInterpreterWeakRef_Close(ref);            \
+    ref = 0;                                    \
+} while (0);
+
 // Exports for '_testcapi' shared extension
 PyAPI_FUNC(Py_ssize_t) _PyInterpreterState_Refcount(PyInterpreterState *interp);
-PyAPI_FUNC(void) _PyInterpreterState_Incref(PyInterpreterState *interp);
+PyAPI_FUNC(int) _PyInterpreterState_Incref(PyInterpreterState *interp);
 
 PyAPI_FUNC(int) PyThreadState_Ensure(PyInterpreterRef interp_ref);
 
diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -2546,6 +2546,16 @@ toggle_reftrace_printer(PyObject *ob, PyObject *arg)
     Py_RETURN_NONE;
 }
 
+static PyInterpreterRef
+get_strong_ref(void)
+{
+    PyInterpreterRef ref;
+    if (PyInterpreterRef_Get(&ref) < 0) {
+        Py_FatalError("strong reference should not have failed");
+    }
+    return ref;
+}
+
 static PyObject *
 test_interp_refcount(PyObject *self, PyObject *unused)
 {
@@ -2555,16 +2565,16 @@ test_interp_refcount(PyObject *self, PyObject *unused)
 
     // Reference counts are technically 0 by default
     assert(_PyInterpreterState_Refcount(interp) == 0);
-    ref1 = PyInterpreterRef_Get();
+    ref1 = get_strong_ref();
     assert(_PyInterpreterState_Refcount(interp) == 1);
-    ref2 = PyInterpreterRef_Get();
+    ref2 = get_strong_ref();
     assert(_PyInterpreterState_Refcount(interp) == 2);
     PyInterpreterRef_Close(ref1);
     assert(_PyInterpreterState_Refcount(interp) == 1);
     PyInterpreterRef_Close(ref2);
     assert(_PyInterpreterState_Refcount(interp) == 0);
 
-    ref1 = PyInterpreterRef_Get();
+    ref1 = get_strong_ref();
     ref2 = PyInterpreterRef_Dup(ref1);
     assert(_PyInterpreterState_Refcount(interp) == 2);
     assert(PyInterpreterRef_AsInterpreter(ref1) == interp);
@@ -2580,7 +2590,10 @@ static PyObject *
 test_interp_weak_ref(PyObject *self, PyObject *unused)
 {
     PyInterpreterState *interp = PyInterpreterState_Get();
-    PyInterpreterWeakRef wref = PyInterpreterWeakRef_Get();
+    PyInterpreterWeakRef wref;
+    if (PyInterpreterWeakRef_Get(&wref) < 0) {
+        return NULL;
+    }
     assert(_PyInterpreterState_Refcount(interp) == 0);
 
     PyInterpreterRef ref;
@@ -2598,10 +2611,10 @@ static PyObject *
 test_interp_ensure(PyObject *self, PyObject *unused)
 {
     PyInterpreterState *interp = PyInterpreterState_Get();
-    PyInterpreterRef ref = PyInterpreterRef_Get();
+    PyInterpreterRef ref = get_strong_ref();
     PyThreadState *save_tstate = PyThreadState_Swap(NULL);
     PyThreadState *tstate = Py_NewInterpreter();
-    PyInterpreterRef sub_ref = PyInterpreterRef_Get();
+    PyInterpreterRef sub_ref = get_strong_ref();
     PyInterpreterState *subinterp = PyThreadState_GetInterpreter(tstate);
 
     for (int i = 0; i < 10; ++i) {
diff --git a/Programs/_testembed.c b/Programs/_testembed.c
@@ -2354,7 +2354,10 @@ test_thread_state_ensure(void)
     _testembed_initialize();
     PyThread_handle_t handle;
     PyThread_ident_t ident;
-    PyInterpreterRef ref = PyInterpreterRef_Get();
+    PyInterpreterRef ref;
+    if (PyInterpreterRef_Get(&ref) < 0) {
+        return -1;
+    };
     ThreadData data = { ref };
     if (PyThread_start_joinable_thread(do_tstate_ensure, &data,
                                        &ident, &handle) < 0) {
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -3145,12 +3145,22 @@ _PyInterpreterState_Refcount(PyInterpreterState *interp)
     return refcount;
 }
 
-void
+int
 _PyInterpreterState_Incref(PyInterpreterState *interp)
 {
     assert(interp != NULL);
-    assert(_Py_atomic_load_ssize_relaxed(&interp->threads.finalizing.countdown) >= 0);
+    struct _Py_finalizing_threads *finalizing = &interp->threads.finalizing;
+    assert(_Py_atomic_load_ssize_relaxed(&finalizing->countdown) >= 0);
+    PyMutex *mutex = &finalizing->mutex;
+    PyMutex_Lock(mutex);
+    if (_PyEvent_IsSet(&finalizing->finished)) {
+        PyMutex_Unlock(mutex);
+        return -1;
+    }
+
     _Py_atomic_add_ssize(&interp->threads.finalizing.countdown, 1);
+    PyMutex_Unlock(mutex);
+    return 0;
 }
 
 static PyInterpreterState *
@@ -3164,19 +3174,29 @@ ref_as_interp(PyInterpreterRef ref)
     return interp;
 }
 
-PyInterpreterRef
-PyInterpreterRef_Get(void)
+int
+PyInterpreterRef_Get(PyInterpreterRef *ref_ptr)
 {
+    assert(ref_ptr != NULL);
     PyInterpreterState *interp = PyInterpreterState_Get();
-    _PyInterpreterState_Incref(interp);
-    return (PyInterpreterRef)interp;
+    if (_PyInterpreterState_Incref(interp) < 0) {
+        PyErr_SetString(PyExc_PythonFinalizationError,
+                        "Cannot acquire strong interpreter references anymore");
+        return -1;
+    }
+    *ref_ptr = (PyInterpreterRef)interp;
+    return 0;
 }
 
 PyInterpreterRef
 PyInterpreterRef_Dup(PyInterpreterRef ref)
 {
     PyInterpreterState *interp = ref_as_interp(ref);
-    _PyInterpreterState_Incref(interp);
+    int res = _PyInterpreterState_Incref(interp);
+    (void)res;
+    // We already hold a strong reference, so it shouldn't be possible
+    // for the interpreter to be at a point where references don't work anymore
+    assert(res == 0);
     return (PyInterpreterRef)interp;
 }
 
@@ -3195,24 +3215,48 @@ PyInterpreterRef_AsInterpreter(PyInterpreterRef ref)
     return interp;
 }
 
-PyInterpreterWeakRef
-PyInterpreterWeakRef_Get(void)
+int
+PyInterpreterWeakRef_Get(PyInterpreterWeakRef *wref_ptr)
 {
     PyInterpreterState *interp = PyInterpreterState_Get();
-    PyInterpreterWeakRef wref = { interp->id };
+    _PyInterpreterWeakRef *wref = PyMem_RawMalloc(sizeof(_PyInterpreterWeakRef));
+    if (wref == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    wref->refcount = 1;
+    wref->id = interp->id;
+    *wref_ptr = (PyInterpreterWeakRef)wref;
+    return 0;
+}
+
+static _PyInterpreterWeakRef *
+wref_handle_as_ptr(PyInterpreterWeakRef wref_handle)
+{
+    _PyInterpreterWeakRef *wref = (_PyInterpreterWeakRef *)wref_handle;
+    if (wref == NULL) {
+        Py_FatalError("Got a null weak interpreter reference, likely due to use after close.");
+    }
+
     return wref;
 }
 
 PyInterpreterWeakRef
-PyInterpreterWeakRef_Dup(PyInterpreterWeakRef wref)
+PyInterpreterWeakRef_Dup(PyInterpreterWeakRef wref_handle)
 {
+    _PyInterpreterWeakRef *wref = wref_handle_as_ptr(wref_handle);
+    ++wref->refcount;
     return wref;
 }
 
+#undef PyInterpreterWeakRef_Close
 void
-PyInterpreterWeakRef_Close(PyInterpreterWeakRef wref)
+PyInterpreterWeakRef_Close(PyInterpreterWeakRef wref_handle)
 {
-    return;
+    _PyInterpreterWeakRef *wref = wref_handle_as_ptr(wref_handle);
+    if (--wref->refcount == 0) {
+        PyMem_RawFree(wref);
+    }
 }
 
 static int
@@ -3234,10 +3278,11 @@ try_acquire_strong_ref(PyInterpreterState *interp, PyInterpreterRef *strong_ptr)
 }
 
 int
-PyInterpreterWeakRef_AsStrong(PyInterpreterWeakRef wref, PyInterpreterRef *strong_ptr)
+PyInterpreterWeakRef_AsStrong(PyInterpreterWeakRef wref_handle, PyInterpreterRef *strong_ptr)
 {
     assert(strong_ptr != NULL);
-    int64_t interp_id = wref.id;
+    _PyInterpreterWeakRef *wref = wref_handle_as_ptr(wref_handle);
+    int64_t interp_id = wref->id;
     /* Interpreters cannot be deleted while we hold the runtime lock. */
     _PyRuntimeState *runtime = &_PyRuntime;
     HEAD_LOCK(runtime);
@@ -3254,13 +3299,12 @@ PyInterpreterWeakRef_AsStrong(PyInterpreterWeakRef wref, PyInterpreterRef *stron
 }
 
 int
-PyInterpreterState_AsStrong(PyInterpreterState *interp, PyInterpreterRef *strong_ptr)
+PyInterpreterRef_Main(PyInterpreterRef *strong_ptr)
 {
-    assert(interp != NULL);
     assert(strong_ptr != NULL);
     _PyRuntimeState *runtime = &_PyRuntime;
     HEAD_LOCK(runtime);
-    int res = try_acquire_strong_ref(interp, strong_ptr);
+    int res = try_acquire_strong_ref(&runtime->_main_interpreter, strong_ptr);
     HEAD_UNLOCK(runtime);
 
     return res;
