Initialize autoTSSkey earlier.

ericsnowcurrently · ericsnowcurrently · commit 7bdf5a41afee · 2023-01-19T08:38:58.000-07:00
diff --git a/Include/internal/pycore_pylifecycle.h b/Include/internal/pycore_pylifecycle.h
@@ -69,7 +69,7 @@ extern void _PyThread_FiniType(PyInterpreterState *interp);
 extern void _Py_Deepfreeze_Fini(void);
 extern void _PyArg_Fini(void);
 
-extern PyStatus _PyGILState_Init(_PyRuntimeState *runtime);
+extern PyStatus _PyGILState_Init(PyInterpreterState *interp);
 extern PyStatus _PyGILState_SetTstate(PyThreadState *tstate);
 extern void _PyGILState_Fini(PyInterpreterState *interp);
 
diff --git a/Include/internal/pycore_pystate.h b/Include/internal/pycore_pystate.h
@@ -149,7 +149,6 @@ PyAPI_FUNC(PyStatus) _PyInterpreterState_Enable(_PyRuntimeState *runtime);
 
 #ifdef HAVE_FORK
 extern PyStatus _PyInterpreterState_DeleteExceptMain(_PyRuntimeState *runtime);
-extern PyStatus _PyGILState_Reinit(_PyRuntimeState *runtime);
 extern void _PySignal_AfterFork(void);
 #endif
 
diff --git a/Modules/posixmodule.c b/Modules/posixmodule.c
@@ -587,7 +587,7 @@ PyOS_AfterFork_Child(void)
     PyStatus status;
     _PyRuntimeState *runtime = &_PyRuntime;
 
-    status = _PyGILState_Reinit(runtime);
+    status = _PyRuntimeState_ReInitThreads(runtime);
     if (_PyStatus_EXCEPTION(status)) {
         goto fatal_error;
     }
@@ -611,11 +611,6 @@ PyOS_AfterFork_Child(void)
 
     _PySignal_AfterFork();
 
-    status = _PyRuntimeState_ReInitThreads(runtime);
-    if (_PyStatus_EXCEPTION(status)) {
-        goto fatal_error;
-    }
-
     status = _PyInterpreterState_DeleteExceptMain(runtime);
     if (_PyStatus_EXCEPTION(status)) {
         goto fatal_error;
diff --git a/Python/pylifecycle.c b/Python/pylifecycle.c
@@ -675,12 +675,7 @@ pycore_create_interpreter(_PyRuntimeState *runtime,
                           const PyConfig *src_config,
                           PyThreadState **tstate_p)
 {
-    /* Auto-thread-state API */
-    PyStatus status = _PyGILState_Init(runtime);
-    if (_PyStatus_EXCEPTION(status)) {
-        return status;
-    }
-
+    PyStatus status;
     PyInterpreterState *interp = PyInterpreterState_New();
     if (interp == NULL) {
         return _PyStatus_ERR("can't make main interpreter");
@@ -692,6 +687,12 @@ pycore_create_interpreter(_PyRuntimeState *runtime,
         return status;
     }
 
+    /* Auto-thread-state API */
+    status = _PyGILState_Init(interp);
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
+    }
+
     const _PyInterpreterConfig config = _PyInterpreterConfig_LEGACY_INIT;
     init_interp_settings(interp, &config);
 
@@ -1795,10 +1796,8 @@ finalize_interp_clear(PyThreadState *tstate)
 static void
 finalize_interp_delete(PyInterpreterState *interp)
 {
-    if (_Py_IsMainInterpreter(interp)) {
-        /* Cleanup auto-thread-state */
-        _PyGILState_Fini(interp);
-    }
+    /* Cleanup auto-thread-state */
+    _PyGILState_Fini(interp);
 
     /* We can't call _PyEval_FiniGIL() here because destroying the GIL lock can
        fail when it is being awaited by another running daemon thread (see
diff --git a/Python/pystate.c b/Python/pystate.c
@@ -121,6 +121,34 @@ current_tss_clear(_PyRuntimeState *runtime)
     }
 }
 
+#ifdef HAVE_FORK
+/* Reset the TSS key - called by PyOS_AfterFork_Child().
+ * This should not be necessary, but some - buggy - pthread implementations
+ * don't reset TSS upon fork(), see issue #10517.
+ */
+static PyStatus
+current_tss_reinit(_PyRuntimeState *runtime)
+{
+    if (!current_tss_initialized(runtime)) {
+        return _PyStatus_OK();
+    }
+    PyThreadState *tstate = current_tss_get(runtime);
+
+    current_tss_fini(runtime);
+    PyStatus status = current_tss_init(runtime);
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
+    }
+
+    /* If the thread had an associated auto thread state, reassociate it with
+     * the new key. */
+    if (tstate && _current_tss_set(runtime, tstate) != 0) {
+        return _PyStatus_ERR("failed to set autoTSSkey");
+    }
+    return _PyStatus_OK();
+}
+#endif
+
 
 /* the global runtime */
 
@@ -243,6 +271,13 @@ _PyRuntimeState_Init(_PyRuntimeState *runtime)
         // Reset to _PyRuntimeState_INIT.
         memcpy(runtime, &initial, sizeof(*runtime));
     }
+
+    PyStatus status = current_tss_init(runtime);
+    if (_PyStatus_EXCEPTION(status)) {
+        _PyRuntimeState_Fini(runtime);
+        return status;
+    }
+
     init_runtime(runtime, open_code_hook, open_code_userdata, audit_hook_head,
                  unicode_next_index, lock1, lock2, lock3, lock4);
 
@@ -252,6 +287,10 @@ _PyRuntimeState_Init(_PyRuntimeState *runtime)
 void
 _PyRuntimeState_Fini(_PyRuntimeState *runtime)
 {
+    if (current_tss_initialized(runtime)) {
+        current_tss_fini(runtime);
+    }
+
     /* Force the allocator used by _PyRuntimeState_Init(). */
     PyMemAllocatorEx old_alloc;
     _PyMem_SetDefaultAllocator(PYMEM_DOMAIN_RAW, &old_alloc);
@@ -304,6 +343,12 @@ _PyRuntimeState_ReInitThreads(_PyRuntimeState *runtime)
         return _PyStatus_ERR("Failed to reinitialize runtime locks");
 
     }
+
+    PyStatus status = current_tss_reinit(runtime);
+    if (_PyStatus_EXCEPTION(status)) {
+        return status;
+    }
+
     return _PyStatus_OK();
 }
 #endif
@@ -978,6 +1023,7 @@ _PyThreadState_Init(PyThreadState *tstate)
 void
 _PyThreadState_SetCurrent(PyThreadState *tstate)
 {
+    assert(tstate != NULL);
     _PyGILState_NoteThreadState(tstate);
 }
 
@@ -1614,19 +1660,33 @@ PyThreadState_IsCurrent(PyThreadState *tstate)
    Py_Initialize/Py_FinalizeEx
 */
 PyStatus
-_PyGILState_Init(_PyRuntimeState *runtime)
+_PyGILState_Init(PyInterpreterState *interp)
 {
-    PyStatus status = current_tss_init(runtime);
-    if (_PyStatus_EXCEPTION(status)) {
-        return status;
+    if (!_Py_IsMainInterpreter(interp)) {
+        /* Currently, PyGILState is shared by all interpreters. The main
+         * interpreter is responsible to initialize it. */
+        return _PyStatus_OK();
     }
-    // PyThreadState_New() calls _PyGILState_NoteThreadState() which does
-    // nothing before autoInterpreterState is set.
+    _PyRuntimeState *runtime = interp->runtime;
+    assert(current_tss_get(runtime) == NULL);
     assert(runtime->gilstate.autoInterpreterState == NULL);
+    runtime->gilstate.autoInterpreterState = interp;
     return _PyStatus_OK();
 }
 
+void
+_PyGILState_Fini(PyInterpreterState *interp)
+{
+    if (!_Py_IsMainInterpreter(interp)) {
+        /* Currently, PyGILState is shared by all interpreters. The main
+         * interpreter is responsible to initialize it. */
+        return;
+    }
+    interp->runtime->gilstate.autoInterpreterState = NULL;
+}
+
 
+// XXX Drop this.
 PyStatus
 _PyGILState_SetTstate(PyThreadState *tstate)
 {
@@ -1641,11 +1701,10 @@ _PyGILState_SetTstate(PyThreadState *tstate)
     }
     _PyRuntimeState *runtime = tstate->interp->runtime;
 
-    runtime->gilstate.autoInterpreterState = tstate->interp;
-    assert(current_tss_get(runtime) == NULL);
-    assert(tstate->gilstate_counter == 0);
+    assert(runtime->gilstate.autoInterpreterState == tstate->interp);
+    assert(current_tss_get(runtime) == tstate);
+    assert(tstate->gilstate_counter == 1);
 
-    _PyGILState_NoteThreadState(tstate);
     return _PyStatus_OK();
 }
 
@@ -1655,41 +1714,6 @@ _PyGILState_GetInterpreterStateUnsafe(void)
     return _PyRuntime.gilstate.autoInterpreterState;
 }
 
-void
-_PyGILState_Fini(PyInterpreterState *interp)
-{
-    current_tss_fini(interp->runtime);
-    interp->runtime->gilstate.autoInterpreterState = NULL;
-}
-
-#ifdef HAVE_FORK
-/* Reset the TSS key - called by PyOS_AfterFork_Child().
- * This should not be necessary, but some - buggy - pthread implementations
- * don't reset TSS upon fork(), see issue #10517.
- */
-PyStatus
-_PyGILState_Reinit(_PyRuntimeState *runtime)
-{
-    if (!current_tss_initialized(runtime)) {
-        return _PyStatus_OK();
-    }
-    PyThreadState *tstate = current_tss_get(runtime);
-
-    current_tss_fini(runtime);
-    PyStatus status = current_tss_init(runtime);
-    if (_PyStatus_EXCEPTION(status)) {
-        return status;
-    }
-
-    /* If the thread had an associated auto thread state, reassociate it with
-     * the new key. */
-    if (tstate && _current_tss_set(runtime, tstate) != 0) {
-        return _PyStatus_ERR("failed to set autoTSSkey");
-    }
-    return _PyStatus_OK();
-}
-#endif
-
 /* When a thread state is created for a thread by some mechanism other than
    PyGILState_Ensure, it's important that the GILState machinery knows about
    it so it doesn't try to create another thread state for the thread (this is
@@ -1700,16 +1724,7 @@ _PyGILState_NoteThreadState(PyThreadState* tstate)
 {
     assert(tstate != NULL);
     _PyRuntimeState *runtime = tstate->interp->runtime;
-    /* If autoTSSkey isn't initialized, this must be the very first
-       threadstate created in Py_Initialize().  Don't do anything for now
-       (we'll be back here when _PyGILState_Init is called). */
-    if (!current_tss_initialized(runtime)) {
-        return;
-    }
-    // XXX Drop this check?
-    if (!runtime->gilstate.autoInterpreterState) {
-        return;
-    }
+    assert(current_tss_initialized(runtime));
 
     /* Stick the thread state for this thread in thread specific storage.
 
