Add runPyAsync

cbits/python.c

@@ -210,3 +210,14 @@ void inline_py_Integer_FromPy(
     PyLong_AsNativeBytes(p, buf, size, -1);
 #endif
 }
+
+
+
+static PyObject* AsyncError = 0;
+
+PyObject* inline_py_AsyncError() {
+    if( AsyncError == 0 ) {
+        AsyncError = PyErr_NewException("inline_py.AsyncError", PyExc_BaseException, 0);
+    }
+    return AsyncError;
+}

include/inline-python.h

@@ -81,3 +81,12 @@ void inline_py_Integer_FromPy(
     void*     buf,
     size_t    size
 );
+
+
+
+// ================================================================
+// Async exceptions
+// ================================================================
+
+// Obtain class for async exception
+PyObject* inline_py_AsyncError();

inline-python.cabal

@@ -80,6 +80,7 @@ Library
     Python.Inline.QQ
     Python.Inline.Eval
     Python.Inline.Types
+    Python.Inline.Async
   Other-modules:
     Python.Internal.CAPI
     Python.Internal.Eval
@@ -98,6 +99,7 @@ library test
                   , tasty                >=1.2
                   , tasty-hunit          >=0.10
                   , tasty-quickcheck     >=0.10
+                  , stm
                   , quickcheck-instances >=0.3.33
                   , exceptions
                   , containers

src/Python/Inline/Async.hs

@@ -0,0 +1,23 @@
+-- |
+-- Asynchronous computation using python. Normally library tries to
+-- execute python code in the same thread. Moreover it use global lock
+-- in addition to GIL in order to avoid blocking capability on GIL.
+-- This module provide API for working with concurrent python.
+-- Its API is heavily modelled after @async@ package.
+--
+-- Note it's very experimental and not well tested. Also mixing
+-- concurrency primitives from two languages makes difficult task of
+-- concurrent programming even more complicated.
+module Python.Inline.Async
+  ( PyAsync
+  , PyAsyncCancelled(..)
+  , runPyAsync
+  , withPyAsync
+  , waitPy
+  , waitPyCatch
+  , cancelPy
+  , uninterruptibleCancelPy
+  ) where
+
+import Python.Internal.Eval
+

src/Python/Internal/Eval.hs

@@ -16,6 +16,15 @@ module Python.Internal.Eval
   , runPy
   , runPyInMain
   , unsafeRunPy
+    -- ** Async
+  , PyAsync
+  , PyAsyncCancelled(..)
+  , waitPy
+  , waitPyCatch
+  , cancelPy
+  , uninterruptibleCancelPy
+  , runPyAsync
+  , withPyAsync
     -- * GC-related
   , newPyObject
     -- * C-API wrappers
@@ -54,6 +63,7 @@ import Control.Monad.Trans.Cont
 import Data.Maybe
 import Data.Function
 import Data.ByteString.Unsafe    qualified as BS
+import Data.Word
 import Foreign.Concurrent        qualified as GHC
 import Foreign.Ptr
 import Foreign.ForeignPtr
@@ -274,6 +284,13 @@ releaseLock tid = readTVar globalPyLock >>= \case
         []    -> LockUnlocked
         t':ts -> Locked t' ts
 
+ensureInit :: STM ()
+ensureInit = readTVar globalPyLock >>= \case
+  LockUninialized -> throwSTM PythonNotInitialized
+  LockFinalized   -> throwSTM PythonIsFinalized
+  LockedByGC      -> pure ()
+  LockUnlocked    -> pure ()
+  Locked{}        -> pure ()
 
 
 ----------------------------------------------------------------
@@ -516,13 +533,91 @@ runPyInMain py
                       takeMVar resp `onException` throwTo tid_main InterruptMain
       either throwM pure r
 
-
 -- | Execute python action. This function is unsafe and should be only
 --   called in thread of interpreter.
 unsafeRunPy :: Py a -> IO a
 unsafeRunPy (Py io) = io
 
 
+----------------------------------------------------------------
+-- Async running
+----------------------------------------------------------------
+
+-- | Exception thrown to a thread doing async python computation.
+data PyAsyncCancelled = PyAsyncCancelled
+  deriving (Show, Eq)
+
+instance Exception PyAsyncCancelled
+
+-- | Handle to asynchronous python computation spawned by
+--   'runPyAsync'. It's performed on separate OS thread. Use
+--   'wait'\/'waitCatch' to obtain computation result.
+data PyAsync a = PyAsync
+  { asyncTID   :: !ThreadId
+  , asyncPyTID :: !Word64
+  , asyncWait  :: STM (Either SomeException a)
+  }
+
+-- | Wait for result of asynchronous computation. If it threw an
+--   exception it will be rethrown by @wait@.
+waitPy :: PyAsync a -> STM a
+waitPy a = either throwSTM pure =<< a.asyncWait
+
+-- | Wait for result of asynchronous computation. Exception thrown by
+--   it will be returned as @Left@.
+waitPyCatch :: PyAsync a -> STM (Either SomeException a)
+waitPyCatch = (.asyncWait)
+
+-- | Cancel execution of asynchronous computation. Most likely thread
+--   will be executing some python so first it attempts to raise async
+--   exception in python code. Then it throws 'PyAsyncCancelled' in case
+--   it executes haskell code. This means thread could be terminate
+--   either with 'PyError' or 'PyAsyncCancelled'.
+--
+--   Note that python code generally is not written under assumption
+--   that it could be smitten with exception at an absolutely any
+--   moment.
+cancelPy :: PyAsync a -> IO ()
+cancelPy PyAsync{asyncTID=tid, asyncPyTID=py_tid}
+  | rtsSupportsBoundThreads = runInBoundThread go
+  | otherwise               = go
+  where
+    go = do
+      [C.block| void {
+        int gil = PyGILState_Ensure();
+        int n   = PyThreadState_SetAsyncExc($(uint64_t py_tid), inline_py_AsyncError());
+        printf("My job is done (%d) tid=(%ld)\n", n, $(uint64_t py_tid));
+        PyGILState_Release(gil);
+        }|]
+      throwTo tid PyAsyncCancelled
+
+-- | Variant of 'cancel' which isn't interruptible.
+uninterruptibleCancelPy :: PyAsync a -> IO ()
+uninterruptibleCancelPy = uninterruptibleMask_ . cancelPy
+
+-- | Create new OS thread and execute python code on it.
+runPyAsync :: Py a -> IO (PyAsync a)
+runPyAsync py = do
+  atomically ensureInit
+  result    <- newEmptyTMVarIO
+  py_tid_mv <- newEmptyMVar
+  tid    <- forkOS $ do
+    -- Obtain python thread ID
+    putMVar py_tid_mv =<< [CU.exp| uint64_t { PyThread_get_thread_ident() } |]
+    a <- try $ unsafeRunPy $ ensureGIL py
+    atomically $ putTMVar result a
+  py_tid <- takeMVar py_tid_mv
+  pure PyAsync
+    { asyncTID   = tid
+    , asyncPyTID = py_tid
+    , asyncWait  = takeTMVar result
+    }
+
+-- | Create new OS thread and execute python code on it. Will use
+--   'uninterruptibleCancel' after callback finishes execution.
+withPyAsync :: Py a -> (PyAsync a -> IO b) -> IO b
+withPyAsync py = bracket (runPyAsync py) uninterruptibleCancelPy
+
 
 ----------------------------------------------------------------
 -- GC-related functions

test/TST/Run.hs

@@ -3,6 +3,7 @@
 module TST.Run(tests) where
 
 import Control.Concurrent
+import Control.Concurrent.STM
 import Control.Exception
 import Control.Monad
 import Control.Monad.IO.Class
@@ -12,6 +13,7 @@ import Test.Tasty.HUnit
 import Python.Inline
 import Python.Inline.QQ
 import Python.Inline.Eval
+import Python.Inline.Async
 import TST.Util
 
 tests :: TestTree
@@ -164,8 +166,46 @@ tests = testGroup "Run python"
         assert m_hs.a == 12
         assert m_hs.b == 'asd'
         |]
+  , testGroup "async" $ guardThreaded
+    [ -- We can run async computation at all
+      testCase "runPyAsync" $ do
+        runPy [pymain| dct = {} |]
+        a <- runPyAsync $ [py_| dct[1] = 100 |]
+        _ <- atomically $ waitPy a
+        n <- runPy $ fromPy =<< [pye| dct[1] |]
+        assertEqual "x" (Just (100::Int)) n
+        runPy [pymain| del dct |]
+    , -- Cancellation of python code
+      testCase "cancelPy [python]" $ do
+        a <- runPyAsync [py_|
+          import time
+          while True:
+              time.sleep(1e-3)
+          |]
+        d <- registerDelay 100_000
+        cancelPy a
+        _ <- atomically $ waitPyCatch a `orElse` do readTVar d >>= \case
+                                                      True  -> error "Timeout"
+                                                      False -> retry
+        return ()
+    -- , -- Cancellation of haskell code
+    --   testCase "cancelPy [haskell]" $ do
+    --     a <- runPyAsync $ do
+    --       liftIO $ forever $ threadDelay 1_000_000
+    --     d <- registerDelay 100_000
+    --     cancelPy a
+    --     _ <- atomically $ waitPyCatch a `orElse` do readTVar d >>= \case
+    --                                                   True  -> error "Timeout"
+    --                                                   False -> retry
+    --     return ()
+    ]
   ]
 
 data Stop = Stop
   deriving stock    Show
   deriving anyclass Exception
+
+guardThreaded :: [TestTree] -> [TestTree]
+guardThreaded ts
+  | rtsSupportsBoundThreads = ts
+  | otherwise               = []