[CBR-243] improve wallet worker start-up and exception handling

k0001 · k0001 · commit 112e345801af · 2018-08-01T16:37:10.000+03:00
diff --git a/pkgs/default.nix b/pkgs/default.nix
@@ -17760,6 +17760,7 @@ license = stdenv.lib.licenses.mit;
 , lens
 , log-warper
 , memory
+, mmorph
 , mtl
 , mwc-random
 , neat-interpolation
@@ -17865,6 +17866,7 @@ json-sop
 lens
 log-warper
 memory
+mmorph
 mtl
 mwc-random
 neat-interpolation
@@ -17887,6 +17889,7 @@ servant-swagger-ui-core
 servant-swagger-ui-redoc
 sqlite-simple
 sqlite-simple-errors
+stm
 swagger2
 text
 time
diff --git a/wallet-new/cardano-sl-wallet-new.cabal b/wallet-new/cardano-sl-wallet-new.cabal
@@ -173,6 +173,7 @@ library
                      , lens
                      , log-warper
                      , memory
+                     , mmorph
                      , mtl
                      , mwc-random
                      , neat-interpolation
@@ -196,6 +197,7 @@ library
                      , servant-swagger-ui-redoc
                      , sqlite-simple
                      , sqlite-simple-errors
+                     , stm
                      , swagger2
                      , text
                      , time
diff --git a/wallet-new/src/Cardano/Wallet/Kernel/Actions.hs b/wallet-new/src/Cardano/Wallet/Kernel/Actions.hs
@@ -2,8 +2,8 @@
 module Cardano.Wallet.Kernel.Actions
     ( WalletAction(..)
     , WalletActionInterp(..)
-    , forkWalletWorker
-    , walletWorker
+    , withWalletWorker
+    , Err_WalletWorkerExpired(..)
     , interp
     , interpList
     , WalletWorkerState
@@ -12,8 +12,10 @@ module Cardano.Wallet.Kernel.Actions
     , isValidState
     ) where
 
-import           Control.Concurrent.Async (async, link)
-import           Control.Concurrent.Chan
+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))
 import           Control.Lens (makeLenses, (%=), (+=), (-=), (.=))
 import           Formatting (bprint, build, shown, (%))
 import qualified Formatting.Buildable
@@ -55,13 +57,11 @@ data WalletWorkerState b = WalletWorkerState
 
 makeLenses ''WalletWorkerState
 
--- A helper function for lifting a `WalletActionInterp` through a monad transformer.
-lifted :: (Monad m, MonadTrans t) => WalletActionInterp m b -> WalletActionInterp (t m) b
-lifted i = WalletActionInterp
-    { applyBlocks  = lift . applyBlocks i
-    , switchToFork = \n bs -> lift (switchToFork i n bs)
-    , emit         = lift . emit i
-    }
+instance MFunctor WalletActionInterp where
+  hoist nat i = WalletActionInterp
+    { applyBlocks = fmap nat (applyBlocks i)
+    , switchToFork = fmap (fmap nat) (switchToFork i)
+    , emit = fmap nat (emit i) }
 
 -- | `interp` is the main interpreter for converting a wallet action to a concrete
 --   transition on the wallet worker's state, perhaps combined with some effects on
@@ -121,20 +121,24 @@ interp walletInterp action = do
       Shutdown -> error "walletWorker: unreacheable dead code, reached!"
 
   where
-    WalletActionInterp{..} = lifted walletInterp
+    WalletActionInterp{..} = hoist lift walletInterp
     prependNewestFirst bs = \nf -> NewestFirst (getNewestFirst bs <> getNewestFirst nf)
 
--- | Connect a wallet action interpreter to a channel of actions.
-walletWorker :: forall b. Chan (WalletAction b) -> WalletActionInterp IO b -> IO ()
-walletWorker chan ops = do
-    emit ops "Starting wallet worker."
-    void $ (`evalStateT` initialWorkerState) tick
-    emit ops "Finishing wallet worker."
-    where
-        tick :: StateT (WalletWorkerState b) IO ()
-        tick = lift (readChan chan) >>= \case
-            Shutdown -> return ()
-            msg      -> interp ops msg >> tick
+-- | Connect a wallet action interpreter to a source actions. This function
+-- returns as soon as the given action returns 'Shutdown'.
+walletWorker
+  :: Ex.MonadMask m
+  => WalletActionInterp m b
+  -> m (WalletAction b)
+  -> m ()
+walletWorker wai getWA = Ex.bracket_
+  (emit wai "Starting wallet worker.")
+  (evalStateT
+     (fix $ \next -> lift getWA >>= \case
+        Shutdown -> pure ()
+        wa -> interp wai wa >> next)
+     initialWorkerState)
+  (emit wai "Stoping wallet worker.")
 
 -- | Connect a wallet action interpreter to a stream of actions.
 interpList :: Monad m => WalletActionInterp m b -> [WalletAction b] -> m (WalletWorkerState b)
@@ -147,13 +151,51 @@ initialWorkerState = WalletWorkerState
     , _lengthPendingBlocks = 0
     }
 
--- | Start up a wallet worker; the worker will respond to actions issued over the
---   returned channel.
-forkWalletWorker :: WalletActionInterp IO b -> IO (WalletAction b -> IO ())
-forkWalletWorker ops = do
-    c <- newChan
-    link =<< async (walletWorker c ops)
-    return (writeChan c)
+-- | 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. Async exceptions from any
+-- source will always be prioritized.
+--
+-- Usage of the obtained 'STM' action after the given continuation has returned
+-- 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
+  -- '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 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
+  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
diff --git a/wallet-new/src/Cardano/Wallet/Kernel/Wallets.hs b/wallet-new/src/Cardano/Wallet/Kernel/Wallets.hs
@@ -59,7 +59,8 @@ instance Exception CreateWalletError
 -------------------------------------------------------------------------------}
 
 -- | Creates a new HD 'Wallet'.
-createHdWallet :: PassiveWallet
+createHdWallet
+             :: PassiveWallet
              -> Mnemonic nat
              -- ^ The set of words (i.e the mnemonic) to generate the initial seed.
              -- See <https://github.com/bitcoin/bips/blob/master/bip-0039.mediawiki#From_mnemonic_to_seed>
diff --git a/wallet-new/src/Cardano/Wallet/WalletLayer/Kernel.hs b/wallet-new/src/Cardano/Wallet/WalletLayer/Kernel.hs
@@ -8,9 +8,9 @@ module Cardano.Wallet.WalletLayer.Kernel
 
 import           Universum
 
+import qualified Control.Concurrent.STM as STM
 import           Control.Lens (to)
 import           Data.Coerce (coerce)
-import           Data.Default (def)
 import           Data.Maybe (fromJust)
 import           Data.Time.Units (Second)
 import           Formatting (build, sformat)
@@ -41,7 +41,6 @@ import           Cardano.Wallet.Kernel.CoinSelection.FromGeneric
                      (CoinSelectionOptions (..), ExpenseRegulation,
                      InputGrouping, newOptions)
 
-import qualified Cardano.Wallet.Kernel.BIP39 as BIP39
 import           Pos.Core (Address, Coin, decodeTextAddress, mkCoin)
 import qualified Pos.Core as Core
 import           Pos.Core.Chrono (OldestFirst (..))
@@ -66,38 +65,25 @@ bracketPassiveWallet
     -> (PassiveWalletLayer n -> Kernel.PassiveWallet -> m a) -> m a
 bracketPassiveWallet logFunction keystore rocksDB f =
     Kernel.bracketPassiveWallet logFunction keystore rocksDB $ \w -> do
-
-      -- Create the wallet worker and its communication endpoint `invoke`.
-      bracket (liftIO $ Actions.forkWalletWorker $ Actions.WalletActionInterp
-                 { Actions.applyBlocks  =  \blunds ->
-                     Kernel.applyBlocks w $
-                         OldestFirst (mapMaybe blundToResolvedBlock (toList (getOldestFirst blunds)))
-                 , Actions.switchToFork = \_ _ -> logFunction Debug "<switchToFork>"
-                 , Actions.emit         = logFunction Debug
-                 }
-              ) (\invoke -> liftIO (invoke Actions.Shutdown))
-              $ \invoke -> do
-                  -- TODO (temporary): build a sample wallet from a backup phrase
-                  _ <- liftIO $ do
-                    Kernel.createHdWallet w
-                                          (def @(BIP39.Mnemonic 12))
-                                          emptyPassphrase
-                                          assuranceLevel
-                                          walletName
-
-                  f (passiveWalletLayer w invoke) w
-
+      let wai = Actions.WalletActionInterp
+                 { Actions.applyBlocks = \blunds ->
+                     Kernel.applyBlocks w
+                        (OldestFirst (mapMaybe blundToResolvedBlock
+                           (toList (getOldestFirst blunds))))
+                 , Actions.switchToFork = \_ _ ->
+                     logFunction Debug "<switchToFork>"
+                 , Actions.emit = logFunction Debug }
+      Actions.withWalletWorker wai $ \invoke -> do
+         f (passiveWalletLayer w invoke) w
   where
-    -- TODO consider defaults
-    walletName       = HD.WalletName "(new wallet)"
-    assuranceLevel   = HD.AssuranceLevelNormal
-
     -- | TODO(ks): Currently not implemented!
     passiveWalletLayer :: Kernel.PassiveWallet
-                       -> (Actions.WalletAction Blund -> IO ())
+                       -> (Actions.WalletAction Blund -> STM ())
                        -> PassiveWalletLayer n
     passiveWalletLayer wallet invoke =
-        PassiveWalletLayer
+        let invokeIO :: forall m'. MonadIO m' => Actions.WalletAction Blund -> m' ()
+            invokeIO = liftIO . STM.atomically . invoke
+        in PassiveWalletLayer
             { _pwlCreateWallet   =
                 \(V1.NewWallet (V1.BackupPhrase mnemonic) mbSpendingPassword v1AssuranceLevel v1WalletName operation) -> do
                     liftIO $ limitExecutionTimeTo (30 :: Second) CreateWalletTimeLimitReached $ do
@@ -166,8 +152,8 @@ bracketPassiveWallet logFunction keystore rocksDB f =
                                      Left  err     -> return (Left $ CreateAddressError err)
             , _pwlGetAddresses   = error "Not implemented!"
 
-            , _pwlApplyBlocks    = liftIO . invoke . Actions.ApplyBlocks
-            , _pwlRollbackBlocks = liftIO . invoke . Actions.RollbackBlocks
+            , _pwlApplyBlocks    = invokeIO . Actions.ApplyBlocks
+            , _pwlRollbackBlocks = invokeIO . Actions.RollbackBlocks
             }
 
     -- The use of the unsafe constructor 'UnsafeRawResolvedBlock' is justified
