Remove duplication between initWorker and registerWorker

lib/core/src/Cardano/Wallet/Api/Server.hs

@@ -147,7 +147,6 @@ import Cardano.Wallet
     , ErrWithdrawalNotWorth (..)
     , ErrWrongPassphrase (..)
     , FeeEstimation (..)
-    , HasLogger
     , HasNetworkLayer
     , TxSubmitLog
     , WalletLog
@@ -673,8 +672,7 @@ postShelleyWallet ctx generateKey body = do
     let state = mkSeqStateFromRootXPrv (rootXPrv, pwd) purposeCIP1852 g
     void $ liftHandler $ initWorker @_ @s @k ctx wid
         (\wrk -> W.createWallet  @(WorkerCtx ctx) @s @k wrk wid wName state)
-        (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @k wrk wid)
-        (\wrk -> W.manageRewardBalance @(WorkerCtx ctx) @s @k (Proxy @n) wrk wid)
+        (\wrk _ -> W.manageRewardBalance @(WorkerCtx ctx) @s @k (Proxy @n) wrk wid)
     withWorkerCtx @_ @s @k ctx wid liftE liftE $ \wrk -> liftHandler $
         W.attachPrivateKeyFromPwd @_ @s @k wrk wid (rootXPrv, pwd)
     fst <$> getWallet ctx (mkShelleyWallet @_ @s @k) (ApiT wid)
@@ -710,8 +708,7 @@ postAccountWallet ctx mkWallet liftKey coworker body = do
     let state = mkSeqStateFromAccountXPub (liftKey accXPub) purposeCIP1852 g
     void $ liftHandler $ initWorker @_ @s @k ctx wid
         (\wrk -> W.createWallet  @(WorkerCtx ctx) @s @k wrk wid wName state)
-        (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @k wrk wid)
-        (`coworker` wid)
+        coworker
     fst <$> getWallet ctx mkWallet (ApiT wid)
   where
     g = maybe defaultAddressPoolGap getApiT (body ^. #addressPoolGap)
@@ -822,8 +819,7 @@ postLegacyWallet
     -> Handler ApiByronWallet
 postLegacyWallet ctx (rootXPrv, pwd) createWallet = do
     void $ liftHandler $ initWorker @_ @s @k ctx wid (`createWallet` wid)
-        (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @k wrk wid)
-        (`idleWorker` wid)
+        idleWorker
     withWorkerCtx ctx wid liftE liftE $ \wrk -> liftHandler $
         W.attachPrivateKeyFromPwd wrk wid (rootXPrv, pwd)
     fst <$> getWallet ctx mkLegacyWallet (ApiT wid)
@@ -928,8 +924,7 @@ postRandomWalletFromXPrv ctx body = do
     s <- liftIO $ mkRndState byronKey <$> getStdRandom random
     void $ liftHandler $ initWorker @_ @s @k ctx wid
         (\wrk -> W.createWallet @(WorkerCtx ctx) @s @k wrk wid wName s)
-        (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @k wrk wid)
-        (`idleWorker` wid)
+        idleWorker
     withWorkerCtx ctx wid liftE liftE $ \wrk -> liftHandler $
         W.attachPrivateKeyFromPwdHash wrk wid (byronKey, pwd)
     fst <$> getWallet ctx mkLegacyWallet (ApiT wid)
@@ -1963,58 +1958,6 @@ postAccountPublicKey ctx (ApiT wid) (ApiT ix) (ApiPostAccountKeyData (ApiT pwd)
                                   Helpers
 -------------------------------------------------------------------------------}
 
--- | see 'Cardano.Wallet#createWallet'
-initWorker
-    :: forall ctx s k.
-        ( HasWorkerRegistry s k ctx
-        , HasDBFactory s k ctx
-        , HasLogger (WorkerLog WalletId WalletLog) ctx
-        )
-    => ctx
-        -- ^ Surrounding API context
-    -> WalletId
-        -- ^ Wallet Id
-    -> (WorkerCtx ctx -> ExceptT ErrWalletAlreadyExists IO WalletId)
-        -- ^ Create action
-    -> (WorkerCtx ctx -> ExceptT ErrNoSuchWallet IO ())
-        -- ^ Restore action
-    -> (WorkerCtx ctx -> IO ())
-        -- ^ Action to run concurrently with restore action
-    -> ExceptT ErrCreateWallet IO WalletId
-initWorker ctx wid createWallet restoreWallet coworker =
-    liftIO (Registry.lookup re wid) >>= \case
-        Just _ ->
-            throwE $ ErrCreateWalletAlreadyExists $ ErrWalletAlreadyExists wid
-        Nothing ->
-            liftIO (Registry.register @_ @ctx re ctx wid config) >>= \case
-                Nothing ->
-                    throwE ErrCreateWalletFailedToCreateWorker
-                Just _ ->
-                    pure wid
-  where
-    config = MkWorker
-        { workerBefore = \ctx' _ -> do
-            void $ unsafeRunExceptT $ createWallet ctx'
-
-        , workerMain = \ctx' _ -> do
-            race_
-                (unsafeRunExceptT $ restoreWallet ctx')
-                (coworker ctx')
-
-        , workerAfter =
-            defaultWorkerAfter
-
-        , workerAcquire =
-            withDatabase df wid
-        }
-    re = ctx ^. workerRegistry @s @k
-    df = ctx ^. dbFactory @s @k
-
--- | Something to pass as the coworker action to 'newApiLayer', which does
--- nothing, and never exits.
-idleWorker :: ctx -> wid -> IO ()
-idleWorker _ _ = forever $ threadDelay maxBound
-
 -- | Handler for fetching the 'ArgGenChange' for the 'RndState' (i.e. the root
 -- XPrv), necessary to derive new change addresses.
 rndStateChange
@@ -2378,35 +2321,82 @@ registerWorker
         , IsOurs s RewardAccount
         , IsOurs s Address
         )
-    => ApiLayer s k
+    => ctx
     -> (WorkerCtx ctx -> WalletId -> IO ())
+        -- ^ Action to run concurrently with restore
     -> WalletId
     -> IO ()
-registerWorker ctx coworker wid =
-    void $ Registry.register @_ @ctx re ctx wid config
+registerWorker ctx coworker = void . startWorker ctx before coworker
   where
+    before ctx' wid =
+        runExceptT (W.checkWalletIntegrity ctx' wid gp)
+        >>= either throwIO pure
     (_, NetworkParameters gp _ _, _) = ctx ^. genesisData
-    re = ctx ^. workerRegistry
+
+-- | Register a wallet creation worker to the registry.
+-- See 'Cardano.Wallet#createWallet'
+initWorker
+    :: forall ctx s k.
+        ( ctx ~ ApiLayer s k
+        , IsOurs s RewardAccount
+        , IsOurs s Address
+        )
+    => ctx
+        -- ^ Surrounding API context
+    -> WalletId
+        -- ^ Wallet Id
+    -> (WorkerCtx ctx -> ExceptT ErrWalletAlreadyExists IO WalletId)
+        -- ^ Create action
+    -> (WorkerCtx ctx -> WalletId -> IO ())
+        -- ^ Action to run concurrently with restore
+    -> ExceptT ErrCreateWallet IO WalletId
+initWorker ctx wid createWallet coworker =
+    liftIO (Registry.lookup re wid) >>= \case
+        Just _ ->
+            throwE $ ErrCreateWalletAlreadyExists $ ErrWalletAlreadyExists wid
+        Nothing ->
+            liftIO (startWorker ctx before coworker wid) >>= \case
+                Nothing -> throwE ErrCreateWalletFailedToCreateWorker
+                Just _ -> pure wid
+  where
+    before ctx' _ = void $ unsafeRunExceptT $ createWallet ctx'
+    re = ctx ^. workerRegistry @s @k
+
+-- | Create a worker for an existing wallet, register it, then start the worker
+-- thread.
+startWorker
+    :: forall ctx s k.
+        ( ctx ~ ApiLayer s k
+        , IsOurs s RewardAccount
+        , IsOurs s Address
+        )
+    => ctx
+    -> (WorkerCtx ctx -> WalletId -> IO ())
+        -- ^ First action to run after starting the worker thread.
+    -> (WorkerCtx ctx -> WalletId -> IO ())
+        -- ^ Action to run concurrently with restore.
+    -> WalletId
+    -> IO (Maybe ctx)
+startWorker ctx before coworker wid =
+    fmap (const ctx) <$> Registry.register @_ @ctx re ctx wid config
+  where
+    re = ctx ^. workerRegistry @s @k
     df = ctx ^. dbFactory
     config = MkWorker
-        { workerBefore = \ctx' _ ->
-            runExceptT (W.checkWalletIntegrity ctx' wid gp)
-                >>= either throwIO pure
-
-        , workerMain = \ctx' _ -> do
-            -- FIXME:
-            -- Review error handling here
-            race_
-                (unsafeRunExceptT $ W.restoreWallet @(WorkerCtx ctx) @s ctx' wid)
-                (coworker ctx' wid)
-
-        , workerAfter =
-            defaultWorkerAfter
-
-        , workerAcquire =
-            withDatabase df wid
+        { workerAcquire = withDatabase df wid
+        , workerBefore = before
+        , workerAfter = defaultWorkerAfter
+        -- fixme: ADP-641 Review error handling here
+        , workerMain = \ctx' _ -> race_
+            (unsafeRunExceptT $ W.restoreWallet ctx' wid)
+            (coworker ctx' wid)
         }
 
+-- | Something to pass as the coworker action to 'newApiLayer', which does
+-- nothing, and never exits.
+idleWorker :: ctx -> wid -> IO a
+idleWorker _ _ = forever $ threadDelay maxBound
+
 -- | Run an action in a particular worker context. Fails if there's no worker
 -- for a given id.
 withWorkerCtx

lib/core/src/Cardano/Wallet/Registry.hs

@@ -245,6 +245,7 @@ register registry ctx k (MkWorker before main after acquire) = do
                                     Logging
 -------------------------------------------------------------------------------}
 
+-- | Log messages relating to a registry worker thread.
 data WorkerLog key msg
     = MsgThreadAfter AfterThreadLog
     | MsgFromWorker key msg
@@ -264,6 +265,7 @@ instance HasSeverityAnnotation msg => HasSeverityAnnotation (WorkerLog key msg)
         MsgThreadAfter msg -> getSeverityAnnotation msg
         MsgFromWorker _ msg -> getSeverityAnnotation msg
 
+-- | Log messages describing how a worker thread exits.
 data AfterThreadLog
     = MsgThreadFinished
     | MsgThreadCancelled