Fix a few bugs and add a test.

Author: ZeroIntensity

Include/internal/pycore_interp_structs.h

@@ -797,6 +797,7 @@ struct _is {
             Py_ssize_t countdown;
             PyEvent finished;
             PyMutex mutex;
+            int shutting_down;
         } finalizing;
     } threads;
 

Modules/_testcapimodule.c

@@ -2571,6 +2571,26 @@ test_interp_refcount(PyObject *self, PyObject *unused)
     Py_RETURN_NONE;
 }
 
+static PyObject *
+test_interp_lookup(PyObject *self, PyObject *unused)
+{
+    PyInterpreterState *interp = PyInterpreterState_Get();
+    assert(_PyInterpreterState_Refcount(interp) == 0);
+    int64_t interp_id = PyInterpreterState_GetID(interp);
+    PyInterpreterState *ref = PyInterpreterState_Lookup(interp_id);
+    assert(ref == interp);
+    assert(_PyInterpreterState_Refcount(interp) == 1);
+    PyInterpreterState_Release(ref);
+    assert(PyInterpreterState_Lookup(10000) == NULL);
+    Py_BEGIN_ALLOW_THREADS;
+    ref = PyInterpreterState_Lookup(interp_id);
+    assert(ref == interp);
+    PyInterpreterState_Release(ref);
+    Py_END_ALLOW_THREADS;
+
+    Py_RETURN_NONE;
+}
+
 static PyMethodDef TestMethods[] = {
     {"set_errno",               set_errno,                       METH_VARARGS},
     {"test_config",             test_config,                     METH_NOARGS},
@@ -2666,6 +2686,7 @@ static PyMethodDef TestMethods[] = {
     {"code_offset_to_line", _PyCFunction_CAST(code_offset_to_line), METH_FASTCALL},
     {"toggle_reftrace_printer", toggle_reftrace_printer, METH_O},
     {"test_interp_refcount", test_interp_refcount, METH_NOARGS},
+    {"test_interp_lookup", test_interp_lookup, METH_NOARGS},
     {NULL, NULL} /* sentinel */
 };
 

Python/pylifecycle.c

@@ -3469,6 +3469,7 @@ wait_for_native_shutdown(PyInterpreterState *interp)
 {
     assert(interp != NULL);
     struct _Py_finalizing_threads *finalizing = &interp->threads.finalizing;
+    _Py_atomic_store_int_release(&finalizing->shutting_down, 1);
     PyMutex_Lock(&finalizing->mutex);
     if (_Py_atomic_load_ssize_relaxed(&finalizing->countdown) == 0) {
         // Nothing to do.

Python/pystate.c

@@ -1373,12 +1373,8 @@ interp_look_up_id(_PyRuntimeState *runtime, int64_t requested_id)
     return NULL;
 }
 
-/* Return the interpreter state with the given ID.
-
-   Fail with RuntimeError if the interpreter is not found. */
-
 PyInterpreterState *
-_PyInterpreterState_LookUpID(int64_t requested_id)
+_PyInterpreterState_LookUpIDNoErr(int64_t requested_id)
 {
     PyInterpreterState *interp = NULL;
     if (requested_id >= 0) {
@@ -1387,6 +1383,18 @@ _PyInterpreterState_LookUpID(int64_t requested_id)
         interp = interp_look_up_id(runtime, requested_id);
         HEAD_UNLOCK(runtime);
     }
+    return interp;
+}
+
+/* Return the interpreter state with the given ID.
+
+   Fail with RuntimeError if the interpreter is not found. */
+
+PyInterpreterState *
+_PyInterpreterState_LookUpID(int64_t requested_id)
+{
+    assert(_PyThreadState_GET() != NULL);
+    PyInterpreterState *interp = _PyInterpreterState_LookUpIDNoErr(requested_id);
     if (interp == NULL && !PyErr_Occurred()) {
         PyErr_Format(PyExc_InterpreterNotFoundError,
                      "unrecognized interpreter ID %lld", requested_id);
@@ -1807,13 +1815,23 @@ PyThreadState_Clear(PyThreadState *tstate)
 static void
 decrement_stoptheworld_countdown(struct _stoptheworld_state *stw);
 
+static int
+shutting_down_natives(PyInterpreterState *interp)
+{
+    assert(interp != NULL);
+    return _Py_atomic_load_int_relaxed(&interp->threads.finalizing.shutting_down);
+}
+
 static void
 decrement_daemon_count(PyInterpreterState *interp)
 {
     assert(interp != NULL);
     struct _Py_finalizing_threads *finalizing = &interp->threads.finalizing;
-    if (_Py_atomic_add_ssize(&finalizing->countdown, -1) == 1) {
+    Py_ssize_t old = _Py_atomic_add_ssize(&finalizing->countdown, -1);
+    if (old == 1 && shutting_down_natives(interp)) {
         _PyEvent_Notify(&finalizing->finished);
+    } else if (old <= 0) {
+        Py_FatalError("interpreter has negative reference count");
     }
 }
 
@@ -3074,6 +3092,9 @@ _PyThreadState_CheckConsistency(PyThreadState *tstate)
 int
 _PyThreadState_MustExit(PyThreadState *tstate)
 {
+    if (!tstate->daemon) {
+        return 0;
+    }
     int state = _Py_atomic_load_int_relaxed(&tstate->state);
     return state == _Py_THREAD_SHUTTING_DOWN;
 }
@@ -3221,19 +3242,19 @@ PyInterpreterState_Hold(void)
 PyInterpreterState *
 PyInterpreterState_Lookup(int64_t interp_id)
 {
-    PyInterpreterState *interp = _PyInterpreterState_LookUpID(interp_id);
+    PyInterpreterState *interp = _PyInterpreterState_LookUpIDNoErr(interp_id);
     if (interp == NULL) {
         return NULL;
     }
     HEAD_LOCK(&_PyRuntime); // Prevent deletion
-    struct _Py_finalizing_threads finalizing = interp->threads.finalizing;
-    PyMutex *mutex = &finalizing.mutex;
+    struct _Py_finalizing_threads *finalizing = &interp->threads.finalizing;
+    PyMutex *mutex = &finalizing->mutex;
     PyMutex_Lock(mutex); // Synchronize TOCTOU with the event flag
-    if (_PyEvent_IsSet(&finalizing.finished)) {
+    if (_PyEvent_IsSet(&finalizing->finished)) {
         /* Interpreter has already finished threads */
         interp = NULL;
     } else {
-        _Py_atomic_add_ssize(&finalizing.countdown, 1);
+        _Py_atomic_add_ssize(&finalizing->countdown, 1);
     }
     PyMutex_Unlock(mutex);
     HEAD_UNLOCK(&_PyRuntime);