withWalletWorker: use Async and precise exception type

Author: k0001

wallet-new/src/Cardano/Wallet/Kernel/Actions.hs

@@ -3,6 +3,7 @@ module Cardano.Wallet.Kernel.Actions
     ( WalletAction(..)
     , WalletActionInterp(..)
     , withWalletWorker
+    , Err_WalletWorkerExpired(..)
     , interp
     , interpList
     , WalletWorkerState
@@ -11,7 +12,7 @@ module Cardano.Wallet.Kernel.Actions
     , isValidState
     ) where
 
-import           Control.Concurrent (forkFinally)
+import qualified Control.Concurrent.Async as Async
 import qualified Control.Concurrent.STM as STM
 import qualified Control.Exception.Safe as Ex
 import           Control.Monad.Morph (MFunctor(hoist))
@@ -150,51 +151,51 @@ initialWorkerState = WalletWorkerState
     , _lengthPendingBlocks = 0
     }
 
+-- | Thrown by 'withWalletWorker''s continuation in case it's used outside of
+-- its intended scope.
+data Err_WalletWorkerExpired = Err_WalletWorkerExpired deriving (Show)
+instance Ex.Exception Err_WalletWorkerExpired
 
 -- | Start a wallet worker in backround who will react to input provided via the
 -- 'STM' function, in FIFO order.
 --
 -- After the given continuation returns (successfully or due to some exception),
 -- the worker will continue processing any pending input before returning,
--- re-throwing the continuation's exception if any.
+-- re-throwing the continuation's exception if any. Async exceptions from any
+-- source will always be prioritized.
 --
 -- Usage of the obtained 'STM' action after the given continuation has returned
--- will fail with an exception.
+-- is not possible. It will throw 'Err_WalletWorkerExpired'.
 withWalletWorker
   :: (MonadIO m, Ex.MonadMask m)
   => WalletActionInterp IO a
   -> ((WalletAction a -> STM ()) -> m b)
   -> m b
 withWalletWorker wai k = do
-  -- 'mDone' is full if the worker finished.
-  mDone :: MVar (Either Ex.SomeException ()) <- liftIO newEmptyMVar
-  -- 'tqWA' keeps items to be processed by the worker.
+  -- 'tqWA' keeps items to be processed by the worker in FIFO order.
   tqWA :: STM.TQueue (WalletAction a) <- liftIO STM.newTQueueIO
   -- 'tvOpen' is 'True' as long as 'tqWA' can receive new input.
   tvOpen :: STM.TVar Bool <- liftIO (STM.newTVarIO True)
-  -- 'getWA' returns the next action to be processed. This function blocks
+  -- 'getWA' returns the next action to be processed. This function retries
   -- unless 'tvOpen' is 'False', in which case 'Shutdown' is returned.
   let getWA :: STM (WalletAction a)
       getWA = STM.tryReadTQueue tqWA >>= \case
          Just wa -> pure wa
          Nothing -> STM.readTVar tvOpen >>= \case
             False -> pure Shutdown
             True  -> STM.retry
-  -- 'pushWA' adds an action to be executed by the worker, in FIFO order. It
-  -- will throw 'BlockedIndefinitelyOnSTM' if used after `k` returns.
-  let pushWA :: WalletAction a -> STM ()
-      pushWA = \wa -> do STM.check =<< STM.readTVar tvOpen
-                         STM.writeTQueue tqWA wa
-  liftIO $ void $ forkFinally
-     (walletWorker wai (STM.atomically getWA))
-     (putMVar mDone)
-  Ex.finally
-     (k pushWA)
-     (liftIO $ do
-         -- Prevent new input.
-         STM.atomically (STM.writeTVar tvOpen False)
-         -- Wait for the worker to finish.
-         either Ex.throwM pure =<< takeMVar mDone)
+  Ex.bracket
+     (liftIO (Async.async (walletWorker wai (STM.atomically getWA))))
+     (\a1 -> liftIO $ do
+        -- Prevent new input.
+        STM.atomically (STM.writeTVar tvOpen False)
+        -- Wait for the worker to finish, re-throwing any exceptions from it.
+        Async.wait a1)
+     (\_ -> k $ \wa -> do
+        -- Add a WalletAction to the queue, unless it's been closed already.
+        STM.readTVar tvOpen >>= \case
+           True -> STM.writeTQueue tqWA wa
+           False -> Ex.throwM Err_WalletWorkerExpired)
 
 -- | Check if this is the initial worker state.
 isInitialState :: Eq b => WalletWorkerState b -> Bool