fix JIT invalidation mechanism for FT

Author: Fidget-Spinner

Include/cpython/code.h

@@ -16,6 +16,7 @@ typedef struct {
 } _PyCoCached;
 
 typedef struct {
+    uint8_t is_finalizing;
     int size;
     int capacity;
     struct _PyExecutorObject *executors[1];

Include/internal/pycore_optimizer.h

@@ -77,7 +77,9 @@ PyAPI_FUNC(void) _Py_Executor_DependsOn(_PyExecutorObject *executor, void *obj);
 #define _Py_MAX_ALLOWED_GLOBALS_MODIFICATIONS 6
 
 #ifdef _Py_TIER2
+PyAPI_FUNC(void) _Py_Executors_InvalidateDependencyWorldStopped(PyInterpreterState *interp, void *obj, int is_invalidation);
 PyAPI_FUNC(void) _Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is_invalidation);
+PyAPI_FUNC(void) _Py_Executors_InvalidateAllWorldStopped(PyInterpreterState *interp, int is_invalidation);
 PyAPI_FUNC(void) _Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation);
 PyAPI_FUNC(void) _Py_Executors_InvalidateCold(PyThreadState *tstate);
 

Objects/codeobject.c

@@ -2185,6 +2185,13 @@ _PyCode_ReturnsOnlyNone(PyCodeObject *co)
 static void
 clear_executors(PyCodeObject *co)
 {
+#ifdef Py_GIL_DISABLED
+    uint8_t expected = 0;
+    if (!_Py_atomic_compare_exchange_uint8(&co->co_executors->is_finalizing, &expected, 1)) {
+        // Another thread is already finalizing this.
+        return;
+    }
+#endif
     assert(co->co_executors);
     for (int i = 0; i < co->co_executors->size; i++) {
         if (co->co_executors->executors[i]) {

Python/instrumentation.c

@@ -1785,7 +1785,7 @@ force_instrument_lock_held(PyCodeObject *code, PyInterpreterState *interp)
     if (code->co_executors != NULL) {
         _PyCode_Clear_Executors(code);
     }
-    _Py_Executors_InvalidateDependency(_PyInterpreterState_GET(), code, 1);
+    _Py_Executors_InvalidateDependencyWorldStopped(_PyInterpreterState_GET(), code, 1);
 #endif
     int code_len = (int)Py_SIZE(code);
     /* Exit early to avoid creating instrumentation
@@ -2028,7 +2028,7 @@ _PyMonitoring_SetEvents(int tool_id, _PyMonitoringEventSet events)
     }
     set_global_version(tstate, new_version);
 #ifdef _Py_TIER2
-    _Py_Executors_InvalidateAll(interp, 1);
+    _Py_Executors_InvalidateAllWorldStopped(interp, 1);
 #endif
     return instrument_all_executing_code_objects(interp);
 }

Python/optimizer.c

@@ -81,6 +81,7 @@ get_index_for_executor(PyCodeObject *code, _Py_CODEUNIT *instr)
         if (new == NULL) {
             return -1;
         }
+        new->is_finalizing = 0;
         new->capacity = new_capacity;
         new->size = size;
         code->co_executors = new;
@@ -1698,7 +1699,7 @@ _PyJit_Tracer_InvalidateDependency(PyThreadState *tstate, void *obj)
  * May cause other executors to be invalidated as well
  */
 void
-_Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is_invalidation)
+_Py_Executors_InvalidateDependencyWorldStopped(PyInterpreterState *interp, void *obj, int is_invalidation)
 {
     _PyBloomFilter obj_filter;
     _Py_BloomFilter_Init(&obj_filter);
@@ -1711,7 +1712,6 @@ _Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is
     }
     /* Clearing an executor can deallocate others, so we need to make a list of
      * executors to invalidate first */
-    _PyEval_StopTheWorldAll(&_PyRuntime);
     _Py_FOR_EACH_TSTATE_UNLOCKED(interp, p) {
         _PyJit_Tracer_InvalidateDependency(p, obj);
         for (_PyExecutorObject *exec = ((_PyThreadStateImpl *)p)->jit_executor_state.executor_list_head; exec != NULL;) {
@@ -1728,7 +1728,6 @@ _Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is
             exec = next;
         }
     }
-    _PyEval_StartTheWorldAll(&_PyRuntime);
     for (Py_ssize_t i = 0; i < PyList_GET_SIZE(invalidate); i++) {
         PyObject *exec = PyList_GET_ITEM(invalidate, i);
         executor_clear(exec);
@@ -1741,13 +1740,19 @@ _Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is
 error:
     PyErr_Clear();
     Py_XDECREF(invalidate);
-    // If we're truly out of memory, wiping out everything is a fine fallback:
-    _Py_Executors_InvalidateAll(interp, is_invalidation);
+    // If we're truly out of memory, DO NOT wipe everything as the world is stopped, and we might deadlock.
 }
 
-/* Invalidate all executors */
 void
-_Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation)
+_Py_Executors_InvalidateDependency(PyInterpreterState *interp, void *obj, int is_invalidation)
+{
+    _PyEval_StopTheWorld(interp);
+    _Py_Executors_InvalidateDependencyWorldStopped(interp, obj, is_invalidation);
+    _PyEval_StartTheWorld(interp);
+}
+
+void
+_Py_Executors_InvalidateAllWorldStopped(PyInterpreterState *interp, int is_invalidation)
 {
     PyObject *invalidate = PyList_New(0);
     if (invalidate == NULL) {
@@ -1756,7 +1761,6 @@ _Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation)
     }
     /* Clearing an executor can deallocate others, so we need to make a list of
      * executors to invalidate first */
-    _PyEval_StopTheWorldAll(&_PyRuntime);
     _Py_FOR_EACH_TSTATE_UNLOCKED(interp, p) {
         for (_PyExecutorObject *exec = ((_PyThreadStateImpl *)p)->jit_executor_state.executor_list_head; exec != NULL;) {
             assert(exec->vm_data.valid);
@@ -1771,7 +1775,6 @@ _Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation)
             exec = next;
         }
     }
-    _PyEval_StartTheWorldAll(&_PyRuntime);
     Py_ssize_t list_len = PyList_GET_SIZE(invalidate);
     for (Py_ssize_t i = 0; i < list_len; i++) {
         _PyExecutorObject *executor = (_PyExecutorObject *)PyList_GET_ITEM(invalidate, i);
@@ -1787,6 +1790,16 @@ _Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation)
         }
     }
     Py_DECREF(invalidate);
+    // If we're truly out of memory, DO NOT wipe everything as the world is stopped, and we might deadlock.
+}
+
+/* Invalidate all executors */
+void
+_Py_Executors_InvalidateAll(PyInterpreterState *interp, int is_invalidation)
+{
+    _PyEval_StopTheWorld(interp);
+    _Py_Executors_InvalidateAllWorldStopped(interp, is_invalidation);
+    _PyEval_StartTheWorld(interp);
 }