Tighten native thread state fallback

Author: fangerer

graalpython/com.oracle.graal.python.cext/src/pystate.c

@@ -83,6 +83,16 @@ extern "C" {
 static inline PyThreadState *
 _get_thread_state() {
     PyThreadState *ts = tstate_current;
+    if (UNLIKELY(ts == NULL)) {
+        /*
+         * Very unlikely fallback: this can happen if another thread initializes the C API while
+         * the current thread is attached to Python but blocked and therefore misses eager
+         * initialization of its native 'tstate_current' TLS slot.
+         */
+        ts = GraalPyPrivate_ThreadState_Get(&tstate_current);
+        assert(ts != NULL);
+        tstate_current = ts;
+    }
     assert(ts != NULL);
     return ts;
 }

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/modules/cext/PythonCextPyStateBuiltins.java

@@ -43,6 +43,7 @@
 import static com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiCallPath.Direct;
 import static com.oracle.graal.python.builtins.modules.cext.PythonCextBuiltins.CApiCallPath.Ignored;
 import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.Int;
+import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.Pointer;
 import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyFrameObjectTransfer;
 import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyObject;
 import static com.oracle.graal.python.builtins.objects.cext.capi.transitions.ArgDescriptor.PyObjectBorrowed;
@@ -112,6 +113,42 @@ static Object restore(
         }
     }
 
+    /**
+     * Very unlikely fallback for threads that were already attached when another thread initialized
+     * the C API, but were blocked at that time and therefore could not process the thread-local
+     * action that eagerly initializes their native 'tstate_current' TLS slot.
+     */
+    @CApiBuiltin(ret = PyThreadState, args = {Pointer}, acquireGil = false, call = Ignored)
+    abstract static class GraalPyPrivate_ThreadState_Get extends CApiUnaryBuiltinNode {
+        private static final TruffleLogger LOGGER = CApiContext.getLogger(GraalPyPrivate_ThreadState_Get.class);
+
+        @Specialization
+        @TruffleBoundary
+        static Object get(Object tstateCurrentPtr) {
+            PythonContext context = PythonContext.get(null);
+            PythonThreadState threadState = context.getThreadState(context.getLanguage());
+
+            /*
+             * The C caller may have observed 'tstate_current == NULL' before entering this upcall.
+             * While entering this builtin, the same thread may process a queued thread-local action
+             * from C API initialization and initialize its native thread state eagerly. So the
+             * fallback decision made in C can be stale by the time we get here.
+             */
+            if (threadState.isNativeThreadStateInitialized()) {
+                LOGGER.fine(() -> String.format("Lazy initialization attempt of native thread state for thread %s aborted. Was initialized in the meantime.", Thread.currentThread()));
+                Object nativeThreadState = PThreadState.getNativeThreadState(threadState);
+                assert nativeThreadState != null;
+                return nativeThreadState;
+            }
+
+            LOGGER.fine(() -> "Lazy (fallback) initialization of native thread state for thread " + Thread.currentThread());
+            assert PThreadState.getNativeThreadState(threadState) == null;
+            Object nativeThreadState = PThreadState.getOrCreateNativeThreadState(threadState);
+            threadState.setNativeThreadLocalVarPointer(tstateCurrentPtr);
+            return nativeThreadState;
+        }
+    }
+
     @CApiBuiltin(ret = Void, args = {}, call = Ignored)
     abstract static class GraalPyPrivate_BeforeThreadDetach extends CApiNullaryBuiltinNode {
         @Specialization

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/CApiContext.java

@@ -859,7 +859,7 @@ protected void perform(ThreadLocalAction.Access access) {
                 context.initializeNativeThreadState();
             }
         };
-        submitThreadLocalActions(context, threads, action);
+        context.getEnv().submitThreadLocal(threads, action);
     }
 
     private static Thread[] getOtherAliveAttachedThreads(PythonContext context) {
@@ -873,12 +873,6 @@ private static Thread[] getOtherAliveAttachedThreads(PythonContext context) {
         return threads.toArray(Thread[]::new);
     }
 
-    private static void submitThreadLocalActions(PythonContext context, Thread[] threads, ThreadLocalAction action) {
-        for (Thread thread : threads) {
-            context.getEnv().submitThreadLocal(new Thread[]{thread}, action);
-        }
-    }
-
     private static CApiContext loadCApi(Node node, PythonContext context, TruffleString name, TruffleString path, String reason) throws IOException, ImportException, ApiInitException {
         Env env = context.getEnv();
         InteropLibrary U = InteropLibrary.getUncached();

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/CExtNodes.java

@@ -853,15 +853,6 @@ static Object doWithoutContext(NativeCAPISymbol symbol, Object[] args,
                 if (capiState != PythonContext.CApiState.INITIALIZING && capiState != PythonContext.CApiState.INITIALIZED) {
                     CompilerDirectives.transferToInterpreterAndInvalidate();
                     CApiContext.ensureCapiWasLoaded("call internal native GraalPy function");
-                    capiState = pythonContext.getCApiState();
-                }
-                if (capiState == PythonContext.CApiState.INITIALIZED) {
-                    PythonContext.PythonThreadState threadState = pythonContext.getThreadState(pythonContext.getLanguage(inliningTarget));
-                    // Native thread-state bootstrap may itself call internal C API helpers before
-                    // 'tstate_current' has been published for this thread.
-                    if (!threadState.isNativeThreadStateInitializationInProgress()) {
-                        pythonContext.ensureNativeThreadStateInitialized(threadState);
-                    }
                 }
                 // TODO review EnsureTruffleStringNode with GR-37896
                 Object callable = CApiContext.getNativeSymbol(inliningTarget, symbol);

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/builtins/objects/cext/capi/ExternalFunctionNodes.java

@@ -804,7 +804,6 @@ private static Object invoke(VirtualFrame frame, PythonContext ctx, CApiTiming t
                         CheckFunctionResultNode checkResultNode, CExtToJavaNode convertReturnValue, PForeignToPTypeNode fromForeign, GetThreadStateNode getThreadStateNode,
                         ExternalFunctionInvokeNode invokeNode) {
             PythonThreadState threadState = getThreadStateNode.execute(inliningTarget, ctx);
-            ctx.ensureNativeThreadStateInitialized(threadState);
             Object result = invokeNode.execute(frame, inliningTarget, threadState, timing, name, callable, cArguments);
             result = checkResultNode.execute(threadState, name, result);
             if (convertReturnValue != null) {

graalpython/com.oracle.graal.python/src/com/oracle/graal/python/runtime/PythonContext.java

@@ -593,6 +593,10 @@ public void setNativeThreadLocalVarPointer(Object ptr) {
             this.nativeThreadLocalVarPointer = ptr;
         }
 
+        public Object getNativeThreadLocalVarPointer() {
+            return nativeThreadLocalVarPointer;
+        }
+
         public boolean isNativeThreadStateInitialized() {
             return nativeThreadLocalVarPointer != null;
         }
@@ -2604,15 +2608,10 @@ public synchronized void attachThread(Thread thread, ContextThreadLocal<PythonTh
 
     @TruffleBoundary
     public void initializeNativeThreadState() {
+        LOGGER.fine(() -> "Initializing native thread state for thread " + Thread.currentThread());
         initializeNativeThreadState(getThreadState(getLanguage()));
     }
 
-    public void ensureNativeThreadStateInitialized(PythonThreadState pythonThreadState) {
-        if (getCApiState() == CApiState.INITIALIZED && !pythonThreadState.isNativeThreadStateInitialized()) {
-            initializeNativeThreadState(pythonThreadState);
-        }
-    }
-
     @SuppressWarnings("try")
     public void initializeNativeThreadState(PythonThreadState pythonThreadState) {
         CompilerAsserts.neverPartOfCompilation();