check env var when re-enabling

Author: Fidget-Spinner

Include/internal/pycore_pylifecycle.h

@@ -129,6 +129,7 @@ PyAPI_FUNC(int) _PyInterpreterConfig_UpdateFromDict(
     PyInterpreterConfig *,
     PyObject *);
 
+extern void _PyInterpreter_SetJitWithEnvVar(PyConfig *config, PyInterpreterState *interp);
 
 #ifdef __cplusplus
 }

Python/pylifecycle.c

@@ -1194,6 +1194,34 @@ run_presite(PyThreadState *tstate)
 }
 #endif
 
+void
+_PyInterpreter_SetJitWithEnvVar(PyConfig *config, PyInterpreterState *interp)
+{
+    int enabled = 1;
+#if _Py_TIER2 & 2
+    enabled = 0;
+#endif
+    char *env = Py_GETENV("PYTHON_JIT");
+    if (env && *env != '\0') {
+        // PYTHON_JIT=0|1 overrides the default
+        enabled = *env != '0';
+    }
+    if (enabled) {
+#ifdef _Py_JIT
+        // perf profiler works fine with tier 2 interpreter, so
+        // only checking for a "real JIT".
+        if (config->perf_profiling > 0) {
+            (void)PyErr_WarnEx(
+                PyExc_RuntimeWarning,
+                "JIT deactivated as perf profiling support is active",
+                0);
+        } else
+#endif
+        {
+            FT_ATOMIC_STORE_CHAR_RELAXED(interp->jit, 1);
+        }
+    }
+}
 
 static PyStatus
 init_interp_main(PyThreadState *tstate)
@@ -1345,30 +1373,7 @@ init_interp_main(PyThreadState *tstate)
     // This is also needed when the JIT is enabled
 #ifdef _Py_TIER2
     if (is_main_interp) {
-        int enabled = 1;
-#if _Py_TIER2 & 2
-        enabled = 0;
-#endif
-        char *env = Py_GETENV("PYTHON_JIT");
-        if (env && *env != '\0') {
-            // PYTHON_JIT=0|1 overrides the default
-            enabled = *env != '0';
-        }
-        if (enabled) {
-#ifdef _Py_JIT
-            // perf profiler works fine with tier 2 interpreter, so
-            // only checking for a "real JIT".
-            if (config->perf_profiling > 0) {
-                (void)PyErr_WarnEx(
-                    PyExc_RuntimeWarning,
-                    "JIT deactivated as perf profiling support is active",
-                    0);
-            } else
-#endif
-            {
-                interp->jit = 1;
-            }
-        }
+        _PyInterpreter_SetJitWithEnvVar(config, interp);
     }
 #endif
 

Python/pystate.c

@@ -1880,9 +1880,9 @@ tstate_delete_common(PyThreadState *tstate, int release_gil)
     // There's only one thread. Re-enable JIT.
     PyThreadState *curr = interp->threads.head;
     if (curr != NULL && curr->prev == NULL && curr->next == NULL) {
-       FT_ATOMIC_STORE_UINT8(interp->jit, 1);
+        _PyInterpreter_SetJitWithEnvVar(_PyInterpreterState_GetConfig(interp), interp);
     }
-#   endif    
+#   endif
 #endif
 
     HEAD_UNLOCK(runtime);