Fix calling python from callbacks in threaded runtime

cbits/python.c

@@ -1,22 +1,93 @@
 #include <inline-python.h>
 #include <stdlib.h>
 
-PyObject *inline_py_function_wrapper(PyCFunction fun, int flags) {
-    PyMethodDef *meth = malloc(sizeof(PyMethodDef));
-    meth->ml_name  = "[inline_python]";
-    meth->ml_meth  = fun;
-    meth->ml_flags = flags;
-    meth->ml_doc   = "Wrapper constructed by inline-python";
-    // Python wrapper which carries PyMethodDef
-    PyObject* meth_obj = PyCapsule_New(meth, NULL, &inline_py_free_capsule);
+// ================================================================
+// Callbacks
+//
+// General idea: we store function pointer (haskell's FunPtr) in
+// PyCapsule and use to call function. Most importantly we must
+// release GIL before calling into haskell. Haskell callback will
+// happen on different thread (on threaded RTS). So it'll have to
+// reacquire GIL there.
+// ================================================================
+
+int inline_py_callback_depth = 0;
+
+static PyObject* callback_METH_O(PyObject* self, PyObject* arg) {
+    PyObject    *res;
+    PyCFunction *fun = PyCapsule_GetPointer(self, NULL);
+    //--
+    inline_py_callback_depth++;
+Py_BEGIN_ALLOW_THREADS
+    res = (*fun)(self, arg);
+Py_END_ALLOW_THREADS
+    inline_py_callback_depth--;
+    return res;
+}
+
+static PyObject* callback_METH_FASTCALL(PyObject* self, PyObject** args, Py_ssize_t nargs) {
+    PyObject        *res;
+    PyCFunctionFast *fun = PyCapsule_GetPointer(self, NULL);
+    //--
+    inline_py_callback_depth++;
+Py_BEGIN_ALLOW_THREADS
+    res = (*fun)(self, args, nargs);
+Py_END_ALLOW_THREADS
+    inline_py_callback_depth--;
+    return res;
+}
+
+static void capsule_free_FunPtr(PyObject* capsule) {
+    PyCFunction *fun = PyCapsule_GetPointer(capsule, NULL);
+    // We call directly to haskell RTS to free FunPtr. Only question
+    // is how stable is this API.
+    freeHaskellFunctionPtr(*fun);
+    free(fun);
+}
+
+static PyMethodDef method_METH_O = {
+    .ml_name  = "[inline_python]",
+    .ml_meth  = callback_METH_O,
+    .ml_flags = METH_O,
+    .ml_doc   = "Wrapper for haskell callback"
+};
+
+static PyMethodDef method_METH_FASTCALL = {
+    .ml_name  = "[inline_python]",
+    .ml_meth  = (PyCFunction)callback_METH_FASTCALL,
+    .ml_flags = METH_FASTCALL,
+    .ml_doc   = "Wrapper for haskell callback"
+};
+
+PyObject *inline_py_callback_METH_O(PyCFunction fun) {
+    PyCFunction *buf = malloc(sizeof(PyCFunction));
+    *buf = fun;
+    PyObject* self = PyCapsule_New(buf, NULL, &capsule_free_FunPtr);
     if( PyErr_Occurred() )
         return NULL;
     // Python function
-    PyObject* f = PyCFunction_New(meth, meth_obj);
-    Py_DECREF(meth_obj);
-    return f;    
+    PyObject* f = PyCFunction_New(&method_METH_O, self);
+    Py_DECREF(self);
+    return f;
 }
 
+PyObject *inline_py_callback_METH_FASTCALL(PyCFunctionFast fun) {
+    PyCFunctionFast *buf = malloc(sizeof(PyCFunctionFast));
+    *buf = fun;
+    PyObject* self = PyCapsule_New(buf, NULL, &capsule_free_FunPtr);
+    if( PyErr_Occurred() )
+        return NULL;
+    // Python function
+    PyObject* f = PyCFunction_New(&method_METH_FASTCALL, self);
+    Py_DECREF(self);
+    return f;
+}
+
+
+// ================================================================
+// Marshalling
+// ================================================================
+
 int inline_py_unpack_iterable(PyObject *iterable, int n, PyObject **out) {
     // Initialize iterator. If object is not an iterable we treat this
     // as not an exception but as a conversion failure
@@ -57,11 +128,3 @@ int inline_py_unpack_iterable(PyObject *iterable, int n, PyObject **out) {
     return -1;
 }
 
-void inline_py_free_capsule(PyObject* py) {
-    PyMethodDef *meth = PyCapsule_GetPointer(py, NULL);
-    // HACK: We want to release wrappers created by wrapper. It
-    //       doesn't seems to be nice and stable C API
-    freeHaskellFunctionPtr(meth->ml_meth);
-    free(meth);
-}
-

include/inline-python.h

@@ -5,14 +5,39 @@
 #include <Rts.h>
 
 
+// Use new stable API from 
+#ifndef PyCFunctionFast
+typedef _PyCFunctionFast PyCFunctionFast;
+#endif
+
 // ----------------------------------------------------------------
 // Standard status codes
 
 #define IPY_OK          0
 #define IPY_ERR_PYTHON  1
 #define IPY_ERR_COMPILE 2
 
-// ----------------------------------------------------------------
+
+
+// ================================================================
+// Callbacks
+// ================================================================
+
+// Callback depth. It's used to decide whether we want to just
+// continue in bound thread. Should only be modified while GIL is held
+extern int inline_py_callback_depth;
+
+// Wrap haskell callback using METH_O calling convention
+PyObject *inline_py_callback_METH_O(PyCFunction fun);
+
+// Wrap haskell callback using METH_FASTCALL calling convention
+PyObject *inline_py_callback_METH_FASTCALL(PyCFunctionFast fun);
+
+
+
+// ================================================================
+// Callbacks
+// ================================================================
 
 // Unpack iterable into array of PyObjects. Iterable must contain
 // exactly N elements.
@@ -27,13 +52,3 @@ int inline_py_unpack_iterable(
     int        n,
     PyObject **out
     );
-
-// Allocate python function object which carrries its own PyMethodDef.
-// Returns function object or NULL with error raised.
-//
-// See NOTE: [Creation of python functions]
-PyObject *inline_py_function_wrapper(PyCFunction fun, int flags);
-
-// Free malloc'd buffer inside PyCapsule
-void inline_py_free_capsule(PyObject*);
-

inline-python.cabal

@@ -79,6 +79,8 @@ Library
 ----------------------------------------------------------------
 library test
   import:           language
+  Default-Extensions:
+    QuasiQuotes
   build-depends:    base
                   , inline-python
                   , tasty            >=1.2
@@ -88,6 +90,8 @@ library test
     TST.Run
     TST.ToPy
     TST.FromPy
+    TST.Callbacks
+    TST.Util
 
 test-suite inline-python-tests
   import:           language

src/Python/Inline/Literal.hs

@@ -12,7 +12,6 @@ module Python.Inline.Literal
   , fromPy'
   ) where
 
-import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class
 import Control.Monad.Trans.Class
@@ -23,7 +22,6 @@ import Data.Word
 import Data.Foldable
 import Foreign.Ptr
 import Foreign.C.Types
-import Foreign.Marshal.Alloc
 import Foreign.Storable
 
 import Language.C.Inline         qualified as C
@@ -285,7 +283,7 @@ instance (FromPy a) => FromPy [a] where
 --        with async exception out of the blue
 
 
-instance (FromPy a, ToPy b) => ToPy (a -> IO b) where
+instance (FromPy a, Show a, ToPy b) => ToPy (a -> IO b) where
   basicToPy f = Py $ do
     -- C function pointer for callback
     f_ptr <- wrapO $ \_ p_a -> pyCallback $ do
@@ -295,10 +293,9 @@ instance (FromPy a, ToPy b) => ToPy (a -> IO b) where
         Right a            -> pure a
       liftIO $ unPy . basicToPy =<< f a
     --
-    [C.exp| PyObject* {
-      inline_py_function_wrapper(
-          $(PyObject* (*f_ptr)(PyObject*, PyObject*)),
-          METH_O)
+    [CU.block| PyObject* {
+      inline_py_callback_METH_O(
+          $(PyObject* (*f_ptr)(PyObject*, PyObject*)));
       }|]
 
 instance (FromPy a1, FromPy a2, ToPy b) => ToPy (a1 -> a2 -> IO b) where
@@ -311,10 +308,8 @@ instance (FromPy a1, FromPy a2, ToPy b) => ToPy (a1 -> a2 -> IO b) where
       liftIO $ unPy . basicToPy =<< f a b
     -- Create python function
     [C.block| PyObject* {
-      _PyCFunctionFast impl = $(PyObject* (*f_ptr)(PyObject*, PyObject*const*, int64_t));
-      return inline_py_function_wrapper(
-          (PyCFunction)impl,
-          METH_FASTCALL);
+      PyCFunctionFast impl = $(PyObject* (*f_ptr)(PyObject*, PyObject*const*, int64_t));
+      return inline_py_callback_METH_FASTCALL(impl);
       }|]
 
 loadArgFastcall :: FromPy a => Ptr (Ptr PyObject) -> Int -> Int64 -> Program (Ptr PyObject) a
@@ -331,7 +326,7 @@ loadArgFastcall p_arr i tot = do
 ----------------------------------------------------------------
 
 pyCallback :: Program (Ptr PyObject) (Ptr PyObject) -> IO (Ptr PyObject)
-pyCallback io = unPy $ evalContT io `catchPy` convertHaskell2Py
+pyCallback io = unPy $ ensureGIL $ evalContT io `catchPy` convertHaskell2Py
 
 raiseUndecodedArg :: CInt -> CInt -> Py (Ptr PyObject)
 raiseUndecodedArg n tot = Py [CU.block| PyObject* {

src/Python/Internal/Eval.hs

@@ -15,6 +15,7 @@ module Python.Internal.Eval
     -- * PyObject wrapper
   , newPyObject
   , decref
+  , ensureGIL
     -- * Exceptions
   , convertHaskell2Py
   , convertPy2Haskell
@@ -186,26 +187,30 @@ runPy :: Py a -> IO a
 -- See NOTE: [Threading and exceptions]
 runPy py
   -- Multithreaded RTS
-  | rtsSupportsBoundThreads = do
-      result <- newEmptyMVar
-      status <- newMVar Pending
-      let onExc :: SomeException -> IO b
-          onExc e = do
-            modifyMVar_ status $ \case
-              Pending   -> pure Cancelled
-              Cancelled -> pure Cancelled
-              Done      -> pure Done
-              Running   -> Cancelled <$ [CU.exp| void { PyErr_SetInterrupt() } |]
-            throwIO e
-      (do putMVar toPythonThread $ PyEvalReq{ prog=py, ..}
-          takeMVar result >>= \case
-            Left  e -> throwIO e
-            Right a -> pure a
-        ) `catch` onExc
+  --
+  -- Here we check whether we're in callback or creating a new call
+  | rtsSupportsBoundThreads = [CU.exp| int { inline_py_callback_depth } |] >>= \case
+      0 -> do
+        result <- newEmptyMVar
+        status <- newMVar Pending
+        let onExc :: SomeException -> IO b
+            onExc e = do
+              modifyMVar_ status $ \case
+                Pending   -> pure Cancelled
+                Cancelled -> pure Cancelled
+                Done      -> pure Done
+                Running   -> Cancelled <$ [CU.exp| void { PyErr_SetInterrupt() } |]
+              throwIO e
+        (do putMVar toPythonThread $ PyEvalReq{ prog=py, ..}
+            takeMVar result >>= \case
+              Left  e -> throwIO e
+              Right a -> pure a
+          ) `catch` onExc
+      _ -> unPy $ ensureGIL py
   -- Single-threaded RTS
   --
   -- See NOTE: [Async exceptions]
-  | otherwise = mask_ $ unPy py
+  | otherwise = mask_ $ unPy $ ensureGIL py
 
 
 -- | Execute python action. This function is unsafe and should be only
@@ -311,7 +316,7 @@ evalReq :: IO ()
 -- See NOTE: [Python and threading]
 -- See NOTE: [Threading and exceptions]
 evalReq = do
-  PyEvalReq{prog=Py io, result, status} <- takeMVar toPythonThread
+  PyEvalReq{prog, result, status} <- takeMVar toPythonThread
   -- GC
   let decrefList Nil = pure ()
       decrefList (p `Cons` ps) = do [CU.exp| void { Py_XDECREF($(PyObject* p)) } |]
@@ -324,7 +329,7 @@ evalReq = do
     Cancelled -> return (Cancelled,False)
     Pending   -> return (Running,  True)
   when do_eval $ do
-    a <- (Right <$> mask_ io) `catches`
+    a <- (Right <$> mask_ (unPy $ ensureGIL prog)) `catches`
          [ Handler $ \(e :: AsyncException)     -> throwIO e
          , Handler $ \(e :: SomeAsyncException) -> throwIO e
          , Handler $ \(e :: SomeException)      -> pure (Left e)
@@ -347,6 +352,16 @@ evalReq = do
 decref :: Ptr PyObject -> Py ()
 decref p = Py [CU.exp| void { Py_DECREF($(PyObject* p)) } |]
 
+-- | Ensure that we hold GIL for duration of action
+ensureGIL :: Py a -> Py a
+ensureGIL action = do
+  -- NOTE: We're cheating here and looking behind the veil.
+  --       PyGILState_STATE is defined as enum. Let hope it will stay
+  --       this way.
+  gil_state <- Py [CU.exp| int { PyGILState_Ensure() } |]
+  action `finallyPy` Py [CU.exp| void { PyGILState_Release($(int gil_state)) } |]
+
+
 -- | Wrap raw python object into
 newPyObject :: Ptr PyObject -> Py PyObject
 -- We need to use different implementation for different RTS

src/Python/Internal/EvalQQ.hs

@@ -14,14 +14,11 @@ module Python.Internal.EvalQQ
   , unindent
   ) where
 
-import Control.Exception
 import Control.Monad.IO.Class
 import Control.Monad.Trans.Class
 import Control.Monad.Trans.Cont
 import Data.Char
 import Foreign.C.Types
-import Foreign.C.String
-import Foreign.Marshal.Alloc
 import Foreign.Ptr
 import Foreign.Storable
 import System.Exit

src/Python/Internal/Program.hs

@@ -14,7 +14,6 @@ module Python.Internal.Program
   , withPyWCString
   ) where
 
-import Control.Exception
 import Control.Monad.Trans.Cont
 import Data.Coerce
 import Foreign.Ptr

src/Python/Types.hs

@@ -10,13 +10,8 @@ module Python.Types
   ) where
 
 import Data.Coerce
-
 import Foreign.Ptr
-import Foreign.ForeignPtr
-import Language.C.Inline         qualified as C
-
 import GHC.ForeignPtr
-
 import Python.Internal.Types
 
 unsafeWithPyObject :: forall a. PyObject -> (Ptr PyObject -> Py a) -> Py a