[GR-46577] Support calling python from natively created threads

PullRequest: graalpython/2835

Author: msimacek

graalpython/com.oracle.graal.python.cext/include/pystate.h

@@ -1,4 +1,4 @@
-/* Copyright (c) 2018, 2022, Oracle and/or its affiliates.
+/* Copyright (c) 2018, 2023, Oracle and/or its affiliates.
  * Copyright (C) 1996-2020 Python Software Foundation
  *
  * Licensed under the PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2
@@ -97,9 +97,8 @@ PyAPI_FUNC(PyFrameObject*) PyThreadState_GetFrame(PyThreadState *tstate);
 PyAPI_FUNC(uint64_t) PyThreadState_GetID(PyThreadState *tstate);
 #endif
 
-typedef
-    enum {PyGILState_LOCKED, PyGILState_UNLOCKED}
-        PyGILState_STATE;
+/* GraalVM change: we need more state bits */
+typedef int PyGILState_STATE;
 
 
 /* Ensure that the current thread is ready to call the Python

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

@@ -1729,14 +1729,6 @@ PyAPI_FUNC(PyObject*) PyFrozenSet_New(PyObject* a) {
 PyAPI_FUNC(int) PyGILState_Check() {
     return GraalPyGILState_Check();
 }
-#undef PyGILState_Ensure
-PyAPI_FUNC(PyGILState_STATE) PyGILState_Ensure() {
-    return GraalPyGILState_Ensure();
-}
-#undef PyGILState_Release
-PyAPI_FUNC(void) PyGILState_Release(PyGILState_STATE a) {
-    GraalPyGILState_Release(a);
-}
 #undef PyImport_GetModuleDict
 PyAPI_FUNC(PyObject*) PyImport_GetModuleDict() {
     return GraalPyImport_GetModuleDict();
@@ -2213,6 +2205,14 @@ PyAPI_FUNC(int) PyTraceMalloc_Track(unsigned int a, uintptr_t b, size_t c) {
 PyAPI_FUNC(int) PyTraceMalloc_Untrack(unsigned int a, uintptr_t b) {
     return GraalPyTraceMalloc_Untrack(a, b);
 }
+#undef PyTruffleGILState_Ensure
+PyAPI_FUNC(int) PyTruffleGILState_Ensure() {
+    return GraalPyTruffleGILState_Ensure();
+}
+#undef PyTruffleGILState_Release
+PyAPI_FUNC(void) PyTruffleGILState_Release() {
+    GraalPyTruffleGILState_Release();
+}
 #undef PyTruffle_Debug
 PyAPI_FUNC(int) PyTruffle_Debug(void* a) {
     return GraalPyTruffle_Debug(a);

graalpython/com.oracle.graal.python.cext/src/capi.h

@@ -170,8 +170,6 @@ typedef struct {
     BUILTIN(PyFrame_New, PyFrameObject*, PyThreadState*, PyCodeObject*, PyObject*, PyObject*) \
     BUILTIN(PyFrozenSet_New, PyObject*, PyObject*) \
     BUILTIN(PyGILState_Check, int) \
-    BUILTIN(PyGILState_Ensure, PyGILState_STATE) \
-    BUILTIN(PyGILState_Release, void, PyGILState_STATE) \
     BUILTIN(PyImport_GetModuleDict, PyObject*) \
     BUILTIN(PyImport_Import, PyObject*, PyObject*) \
     BUILTIN(PyImport_ImportModule, PyObject*, const char*) \
@@ -314,6 +312,8 @@ typedef struct {
     BUILTIN(PyTruffleErr_GetExcInfo, PyObject*) \
     BUILTIN(PyTruffleErr_WarnExplicit, PyObject*, PyObject*, PyObject*, PyObject*, int, PyObject*, PyObject*) \
     BUILTIN(PyTruffleFloat_AsDouble, double, PyObject*) \
+    BUILTIN(PyTruffleGILState_Ensure, int) \
+    BUILTIN(PyTruffleGILState_Release, void) \
     BUILTIN(PyTruffleHash_InitSecret, void, void*) \
     BUILTIN(PyTruffleLong_AsPrimitive, long, PyObject*, int, long) \
     BUILTIN(PyTruffleLong_FromString, PyObject*, const char*, int, int) \

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

@@ -115,3 +115,15 @@ int PyState_RemoveModule(struct PyModuleDef* def) {
     // TODO(fa): implement
     return 0;
 }
+
+// This function has a different implementation on NFI in capi_native.c
+PyAPI_FUNC(PyGILState_STATE) PyGILState_Ensure() {
+    return GraalPyTruffleGILState_Ensure();
+}
+
+// This function has a different implementation on NFI in capi_native.c
+PyAPI_FUNC(void) PyGILState_Release(PyGILState_STATE state) {
+    if (state) {
+        GraalPyTruffleGILState_Release();
+    }
+}

graalpython/com.oracle.graal.python.jni/src/capi_forwards.h

@@ -737,6 +737,8 @@ void unimplemented(const char* name) {
 #undef PyTruffleErr_WarnExplicit
 #undef PyTruffleFloat_AsDouble
 #undef PyTruffleFrame_New
+#undef PyTruffleGILState_Ensure
+#undef PyTruffleGILState_Release
 #undef PyTruffleHash_InitSecret
 #undef PyTruffleLong_AsPrimitive
 #undef PyTruffleLong_FromString
@@ -2564,18 +2566,11 @@ PyAPI_FUNC(int) PyGILState_Check() {
     int result = (int) GraalPyGILState_Check();
     return result;
 }
-PyAPI_FUNC(PyGILState_STATE) PyGILState_Ensure() {
-    PyGILState_STATE result = (PyGILState_STATE) GraalPyGILState_Ensure();
-    return result;
-}
 PyThreadState* (*__target__PyGILState_GetThisThreadState)() = NULL;
 PyAPI_FUNC(PyThreadState*) PyGILState_GetThisThreadState() {
     PyThreadState* result = (PyThreadState*) __target__PyGILState_GetThisThreadState();
     return result;
 }
-PyAPI_FUNC(void) PyGILState_Release(PyGILState_STATE a) {
-    GraalPyGILState_Release(a);
-}
 PyObject* (*__target__PyGen_New)(PyFrameObject*) = NULL;
 PyAPI_FUNC(PyObject*) PyGen_New(PyFrameObject* a) {
     PyObject* result = (PyObject*) __target__PyGen_New(a);
@@ -3955,6 +3950,13 @@ PyAPI_FUNC(int) PyTraceMalloc_Untrack(unsigned int a, uintptr_t b) {
 PyAPI_FUNC(PyFrameObject*) PyTruffleFrame_New(PyThreadState* a, PyCodeObject* b, PyObject* c, PyObject* d) {
     unimplemented("PyTruffleFrame_New"); exit(-1);
 }
+PyAPI_FUNC(int) PyTruffleGILState_Ensure() {
+    int result = (int) GraalPyTruffleGILState_Ensure();
+    return result;
+}
+PyAPI_FUNC(void) PyTruffleGILState_Release() {
+    GraalPyTruffleGILState_Release();
+}
 PyAPI_FUNC(int) PyTruffle_Debug(void* a) {
     int result = (int) GraalPyTruffle_Debug(a);
     return result;

graalpython/com.oracle.graal.python.jni/src/capi_native.c

@@ -50,12 +50,16 @@
 #include <pycore_pymem.h>
 #include <pycore_moduleobject.h>
 
+#include <trufflenfi.h>
+
 #include <stdio.h>
 #include <stdint.h>
 #include <time.h>
 
 #define MUST_INLINE __attribute__((always_inline)) inline
 
+TruffleContext* TRUFFLE_CONTEXT;
+
 #define PY_TYPE_OBJECTS(OBJECT) \
 OBJECT(PyAsyncGen_Type, async_generator) \
 OBJECT(PyBaseObject_Type, object) \
@@ -299,14 +303,16 @@ int initNativeForwardCalled = 0;
 /**
  * Returns 1 on success, 0 on error (if it was already initialized).
  */
-PyAPI_FUNC(int) initNativeForward(void* (*getBuiltin)(int), void* (*getAPI)(const char*), void* (*getType)(const char*), void (*setTypeStore)(const char*, void*), void (*initialize_native_locations)(void*,void*,void*)) {
+PyAPI_FUNC(int) initNativeForward(TruffleEnv* env, void* (*getBuiltin)(int), void* (*getAPI)(const char*), void* (*getType)(const char*), void (*setTypeStore)(const char*, void*), void (*initialize_native_locations)(void*,void*,void*)) {
     if (initNativeForwardCalled) {
     	return 0;
     }
     initNativeForwardCalled = 1;
     clock_t t;
     t = clock();
 
+    TRUFFLE_CONTEXT = (*env)->getTruffleContext(env);
+
 #define SET_TYPE_OBJECT_STORE(NAME, TYPENAME) setTypeStore(#TYPENAME, (void*) &NAME);
     PY_TYPE_OBJECTS(SET_TYPE_OBJECT_STORE)
 #undef SET_TYPE_OBJECT_STORE
@@ -459,6 +465,31 @@ void finalizeCAPI() {
 	GraalPy_set_PyObject_ob_refcnt = nop_GraalPy_set_PyObject_ob_refcnt;
 }
 
+#define _PYGILSTATE_LOCKED   0x1
+#define _PYGILSTATE_ATTACHED 0x2
+
+PyAPI_FUNC(PyGILState_STATE) PyGILState_Ensure() {
+    int result = 0;
+    if ((*TRUFFLE_CONTEXT)->getTruffleEnv(TRUFFLE_CONTEXT) == NULL) {
+        (*TRUFFLE_CONTEXT)->attachCurrentThread(TRUFFLE_CONTEXT);
+        result |= _PYGILSTATE_ATTACHED;
+    }
+    int locked = PyTruffleGILState_Ensure();
+    if (locked) {
+        result |= _PYGILSTATE_LOCKED;
+    }
+    return result;
+}
+
+PyAPI_FUNC(void) PyGILState_Release(PyGILState_STATE state) {
+    if (state & _PYGILSTATE_LOCKED) {
+        PyTruffleGILState_Release();
+    }
+    if (state & _PYGILSTATE_ATTACHED) {
+        (*TRUFFLE_CONTEXT)->detachCurrentThread(TRUFFLE_CONTEXT);
+    }
+}
+
 PyObject* PyTuple_Pack(Py_ssize_t n, ...) {
     va_list vargs;
     va_start(vargs, n);

graalpython/com.oracle.graal.python.test/src/tests/cpyext/test_thread.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2018, 2021, Oracle and/or its affiliates. All rights reserved.
+# Copyright (c) 2018, 2023, Oracle and/or its affiliates. All rights reserved.
 # DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 #
 # The Universal Permissive License (UPL), Version 1.0
@@ -36,8 +36,9 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
+import threading
 
-from . import CPyExtTestCase, CPyExtFunction, unhandled_error_compare
+from . import CPyExtTestCase, CPyExtFunction, unhandled_error_compare, CPyExtType
 
 __dir__ = __file__.rpartition("/")[0]
 
@@ -83,3 +84,43 @@ def compile_module(self, name):
         callfunction="test_PyThread_tss_functions",
         cmpfunc=unhandled_error_compare
     )
+
+
+class TestNativeThread:
+    def test_register_new_thread(self):
+        TestThread = CPyExtType(
+            name="TestThread",
+            includes="#include <pthread.h>",
+            code=r'''
+            void* thread_entrypoint(void* arg) {
+                PyObject* callable = (PyObject*)arg;
+                PyGILState_STATE gstate;
+                gstate = PyGILState_Ensure();
+                if (!PyObject_CallNoArgs(callable)) {
+                    PyErr_WriteUnraisable(callable);
+                }
+                PyGILState_Release(gstate);
+                return NULL;
+            }
+            PyObject* run_in_thread(PyObject* self, PyObject* callable) {
+                Py_BEGIN_ALLOW_THREADS;
+                pthread_t thread;
+                pthread_create(&thread, NULL, thread_entrypoint, callable);
+                pthread_join(thread, NULL);
+                Py_END_ALLOW_THREADS;
+                Py_RETURN_NONE;
+            }
+            ''',
+            tp_methods='{"run_in_thread", (PyCFunction)run_in_thread, METH_O | METH_STATIC, ""}'
+        )
+
+        thread = None
+
+        def callable():
+            nonlocal thread
+            thread = threading.current_thread()
+
+        TestThread.run_in_thread(callable)
+
+        assert thread
+        assert thread is not threading.current_thread()