Store delegation reward account balances in the database

lib/byron/src/Cardano/Wallet/Byron.hs

@@ -274,6 +274,7 @@ serveWallet
             (DefaultFieldValues $ getActiveSlotCoefficient gp)
             databaseDir
         Server.newApiLayer walletEngineTracer params nl' tl db
+            (\_ _ -> pure ())
       where
         gp@GenesisParameters
             { getGenesisBlockHash

lib/byron/src/Cardano/Wallet/Byron/Network.hs

@@ -241,6 +241,7 @@ withNetworkLayer tr np addrInfo versionData action = do
             , postTx = _postTx localTxSubmissionQ
             , stakeDistribution = _stakeDistribution
             , getAccountBalance = _getAccountBalance
+            , watchNodeTip = _watchNodeTip
             }
   where
     gp@W.GenesisParameters
@@ -295,9 +296,15 @@ withNetworkLayer tr np addrInfo versionData action = do
         case result of
             SubmitSuccess -> pure ()
             SubmitFail err -> throwE $ ErrPostTxBadRequest $ T.pack (show err)
+
     _stakeDistribution =
         notImplemented "stakeDistribution"
 
+    -- At the moment it's not necessasy to implement this method unless
+    -- monitoring a reward account.
+    _watchNodeTip =
+        notImplemented "stakeDistribution"
+
 -- | Type representing a network client running two mini-protocols to sync
 -- from the chain and, submit transactions.
 type NetworkClient m = OuroborosApplication

lib/core/src/Cardano/Wallet.hs

@@ -74,6 +74,7 @@ module Cardano.Wallet
     , updateWalletPassphrase
     , walletSyncProgress
     , fetchRewardBalance
+    , manageRewardBalance
     , rollbackBlocks
     , checkWalletIntegrity
     , ErrWalletAlreadyExists (..)
@@ -209,8 +210,6 @@ import Cardano.Wallet.Primitive.AddressDerivation.Byron
     ( ByronKey )
 import Cardano.Wallet.Primitive.AddressDerivation.Icarus
     ( IcarusKey )
-import Cardano.Wallet.Primitive.AddressDerivation.Shelley
-    ( ShelleyKey )
 import Cardano.Wallet.Primitive.AddressDiscovery
     ( CompareDiscovery (..)
     , GenChange (..)
@@ -355,7 +354,7 @@ import Data.Generics.Product.Typed
 import Data.List.NonEmpty
     ( NonEmpty )
 import Data.Maybe
-    ( isJust, mapMaybe )
+    ( mapMaybe )
 import Data.Quantity
     ( Quantity (..) )
 import Data.Set
@@ -364,8 +363,6 @@ import Data.Text.Class
     ( ToText (..) )
 import Data.Time.Clock
     ( UTCTime, getCurrentTime )
-import Data.Type.Equality
-    ( testEquality )
 import Data.Vector.Shuffle
     ( shuffle )
 import Data.Word
@@ -380,8 +377,6 @@ import Safe
     ( lastMay )
 import Statistics.Quantile
     ( medianUnbiased, quantiles )
-import Type.Reflection
-    ( Typeable, typeRep )
 
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Random as Rnd
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Sequential as Seq
@@ -866,51 +861,85 @@ deleteWallet ctx wid = db & \DBLayer{..} -> do
   where
     db = ctx ^. dbLayer @s @k
 
--- | Fetch the reward balance of a given wallet.
+-- | Fetch the cached reward balance of a given wallet from the database.
+fetchRewardBalance
+    :: forall ctx s k.
+        ( HasDBLayer s k ctx
+        )
+    => ctx
+    -> WalletId
+    -> IO (Quantity "lovelace" Word64)
+fetchRewardBalance ctx wid = db & \DBLayer{..} ->
+    atomically $ readDelegationRewardBalance pk
+  where
+    pk = PrimaryKey wid
+    db = ctx ^. dbLayer @s @k
+
+-- | Query the node for the reward balance of a given wallet.
 --
 -- Rather than force all callers of 'readWallet' to wait for fetching the
 -- account balance (via the 'NetworkLayer'), we expose this function for it.
-fetchRewardBalance
+queryRewardBalance
     :: forall ctx s t k.
         ( HasDBLayer s k ctx
         , HasNetworkLayer t ctx
         , HasRewardAccount s k
-        , HasLogger WalletLog ctx
-        , Typeable k
         )
     => ctx
     -> WalletId
     -> ExceptT ErrFetchRewards IO (Quantity "lovelace" Word64)
-fetchRewardBalance ctx wid = db & \DBLayer{..} -> do
-    -- FIXME: issue #1750 re-enable querying reward balance when it's faster
-    if isShelleyKey then do
-        lift $ traceWith tr MsgTemporaryDisableFetchReward
-        pure $ Quantity 0
-    else do
-        let pk = PrimaryKey wid
-        cp <- withExceptT ErrFetchRewardsNoSuchWallet
-            . mapExceptT atomically
-            . withNoSuchWallet wid
-            $ readCheckpoint pk
-        mapExceptT (fmap handleErr)
-            . getAccountBalance nw
-            . toChimericAccount @s @k
-            . rewardAccountKey
-            $ getState cp
+queryRewardBalance ctx wid = db & \DBLayer{..} -> do
+    cp <- withExceptT ErrFetchRewardsNoSuchWallet
+        . mapExceptT atomically
+        . withNoSuchWallet wid
+        $ readCheckpoint pk
+    mapExceptT (fmap handleErr)
+        . getAccountBalance nw
+        . toChimericAccount @s @k
+        . rewardAccountKey
+        $ getState cp
   where
+    pk = PrimaryKey wid
     db = ctx ^. dbLayer @s @k
     nw = ctx ^. networkLayer @t
-    tr = ctx ^. logger
     handleErr = \case
         Right x -> Right x
         Left (ErrGetAccountBalanceAccountNotFound _) ->
             Right $ Quantity 0
         Left (ErrGetAccountBalanceNetworkUnreachable e) ->
             Left $ ErrFetchRewardsNetworkUnreachable e
 
-    isShelleyKey = isJust $ testEquality
-        (typeRep @(k 'RootK XPrv))
-        (typeRep @(ShelleyKey 'RootK XPrv))
+manageRewardBalance
+    :: forall ctx s t k.
+        ( HasLogger WalletLog ctx
+        , HasNetworkLayer t ctx
+        , HasDBLayer s k ctx
+        , HasRewardAccount s k
+        , ctx ~ WalletLayer s t k
+        )
+    => ctx
+    -> WalletId
+    -> IO ()
+manageRewardBalance ctx wid = db & \DBLayer{..} -> do
+    Left err <- runExceptT $ watchNodeTip $ \bh -> do
+         lift $ traceWith tr $ MsgRewardBalanceQuery bh
+         res <- lift $ runExceptT (queryRewardBalance @ctx @s @t @k ctx wid)
+         lift $ traceWith tr $ MsgRewardBalanceResult res
+         case res of
+            Right amt ->
+                withExceptT show $
+                mapExceptT atomically $ putDelegationRewardBalance pk amt
+            Left _err ->
+                -- Occasionaly failing to query is generally not fatal. It will
+                -- just update the balance next time the tip changes.
+                pure ()
+    traceWith tr $ MsgRewardBalanceFinish err
+
+  where
+    pk = PrimaryKey wid
+    db = ctx ^. dbLayer @s @k
+    NetworkLayer{watchNodeTip} = ctx ^. networkLayer @t
+    tr = ctx ^. logger @WalletLog
 
 {-------------------------------------------------------------------------------
                                     Address
@@ -1930,6 +1959,7 @@ data ErrQuitStakePool
 data ErrFetchRewards
     = ErrFetchRewardsNetworkUnreachable ErrNetworkUnavailable
     | ErrFetchRewardsNoSuchWallet ErrNoSuchWallet
+    deriving (Generic, Eq, Show)
 
 data ErrSelectForMigration
     = ErrSelectForMigrationNoSuchWallet ErrNoSuchWallet
@@ -2031,7 +2061,9 @@ data WalletLog
     | MsgDelegationCoinSelection CoinSelection
     | MsgPaymentCoinSelection CoinSelection
     | MsgPaymentCoinSelectionAdjusted CoinSelection
-    | MsgTemporaryDisableFetchReward
+    | MsgRewardBalanceQuery BlockHeader
+    | MsgRewardBalanceResult (Either ErrFetchRewards (Quantity "lovelace" Word64))
+    | MsgRewardBalanceFinish String
     deriving (Show, Eq)
 
 instance ToText WalletLog where
@@ -2073,8 +2105,15 @@ instance ToText WalletLog where
             "Coins selected for payment: \n" <> pretty sel
         MsgPaymentCoinSelectionAdjusted sel ->
             "Coins after fee adjustment: \n" <> pretty sel
-        MsgTemporaryDisableFetchReward ->
-            "FIXME: (issue #1750) fetching reward temporarily disabled."
+        MsgRewardBalanceQuery bh ->
+            "Updating the reward balance for block " <> pretty bh
+        MsgRewardBalanceResult (Right amt) ->
+            "The reward balance is " <> pretty amt
+        MsgRewardBalanceResult (Left err) ->
+            "Problem fetching reward balance. Will try again on next chain update. " <>
+            T.pack (show err)
+        MsgRewardBalanceFinish err ->
+            "Reward balance worker has finished due to: " <> T.pack err
 
 instance HasPrivacyAnnotation WalletLog
 instance HasSeverityAnnotation WalletLog where
@@ -2093,4 +2132,7 @@ instance HasSeverityAnnotation WalletLog where
         MsgDelegationCoinSelection _ -> Debug
         MsgPaymentCoinSelection _ -> Debug
         MsgPaymentCoinSelectionAdjusted _ -> Debug
-        MsgTemporaryDisableFetchReward -> Warning
+        MsgRewardBalanceQuery _ -> Debug
+        MsgRewardBalanceResult (Right _) -> Debug
+        MsgRewardBalanceResult (Left _) -> Notice
+        MsgRewardBalanceFinish _ -> Debug

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

@@ -83,6 +83,9 @@ module Cardano.Wallet.Api.Server
     , rndStateChange
     , assignMigrationAddresses
     , withWorkerCtx
+
+    -- * Workers
+    , manageRewardBalance
     ) where
 
 import Prelude
@@ -130,6 +133,7 @@ import Cardano.Wallet
     , HasLogger
     , WalletLog
     , genesisData
+    , manageRewardBalance
     )
 import Cardano.Wallet.Api
     ( ApiLayer (..)
@@ -275,6 +279,8 @@ import Cardano.Wallet.Unsafe
     ( unsafeRunExceptT )
 import Control.Arrow
     ( second )
+import Control.Concurrent.Async
+    ( concurrently_ )
 import Control.Exception
     ( IOException, bracket, throwIO, tryJust )
 import Control.Monad
@@ -368,8 +374,6 @@ import System.IO.Error
     )
 import System.Random
     ( getStdRandom, random )
-import Type.Reflection
-    ( Typeable )
 
 import qualified Cardano.Wallet as W
 import qualified Cardano.Wallet.Network as NW
@@ -512,7 +516,6 @@ postWallet
         , Bounded (Index (AddressIndexDerivationType k) 'AddressK)
         , HasDBFactory s k ctx
         , HasWorkerRegistry s k ctx
-        , Typeable k
         )
     => ctx
     -> ((SomeMnemonic, Maybe SomeMnemonic) -> Passphrase "encryption" -> k 'RootK XPrv)
@@ -535,7 +538,6 @@ postShelleyWallet
         , HasDBFactory s k ctx
         , HasWorkerRegistry s k ctx
         , HasRewardAccount s k
-        , Typeable k
         )
     => ctx
     -> ((SomeMnemonic, Maybe SomeMnemonic) -> Passphrase "encryption" -> k 'RootK XPrv)
@@ -546,6 +548,7 @@ postShelleyWallet ctx generateKey body = do
     void $ liftHandler $ initWorker @_ @s @k ctx wid
         (\wrk -> W.createWallet  @(WorkerCtx ctx) @s @k wrk wid wName state)
         (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @t @k wrk wid)
+        (\wrk -> W.manageRewardBalance @(WorkerCtx ctx) @s @t @k wrk wid)
     withWorkerCtx @_ @s @k ctx wid liftE liftE $ \wrk -> liftHandler $
         W.attachPrivateKeyFromPwd @_ @s @k wrk wid (rootXPrv, pwd)
     fst <$> getWallet ctx (mkShelleyWallet @_ @s @t @k) (ApiT wid)
@@ -579,6 +582,7 @@ postAccountWallet ctx mkWallet liftKey body = do
     void $ liftHandler $ initWorker @_ @s @k ctx wid
         (\wrk -> W.createWallet  @(WorkerCtx ctx) @s @k wrk wid wName state)
         (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @t @k wrk wid)
+        (\wrk -> W.manageRewardBalance @(WorkerCtx ctx) @s @t @k wrk wid)
     fst <$> getWallet ctx mkWallet (ApiT wid)
   where
     g = maybe defaultAddressPoolGap getApiT (body ^. #addressPoolGap)
@@ -592,14 +596,14 @@ mkShelleyWallet
         ( ctx ~ ApiLayer s t k
         , s ~ SeqState n k
         , IsOurs s Address
-        , HasRewardAccount s k
         , HasWorkerRegistry s k ctx
-        , Typeable k
         )
     => MkApiWallet ctx s ApiWallet
 mkShelleyWallet ctx wid cp meta pending progress = do
-    reward <- withWorkerCtx @_ @s @k ctx wid liftE liftE $ \wrk -> liftHandler $
-        W.fetchRewardBalance @_ @s @t @k wrk wid
+    reward <- withWorkerCtx @_ @s @k ctx wid liftE liftE $ \wrk ->
+        -- never fails - returns zero if balance not found
+        Handler $ ExceptT $ Right <$>
+        W.fetchRewardBalance @_ @s @k wrk wid
     pure ApiWallet
         { addressPoolGap = ApiT $ getState cp ^. #externalPool . #gap
         , balance = ApiT $ WalletBalance
@@ -665,6 +669,7 @@ postLegacyWallet
 postLegacyWallet ctx (rootXPrv, pwd) createWallet = do
     void $ liftHandler $ initWorker @_ @s @k ctx wid (`createWallet` wid)
         (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @t @k wrk wid)
+        (\_ -> pure ())
     withWorkerCtx ctx wid liftE liftE $ \wrk -> liftHandler $
         W.attachPrivateKeyFromPwd wrk wid (rootXPrv, pwd)
     fst <$> getWallet ctx mkLegacyWallet (ApiT wid)
@@ -757,6 +762,7 @@ postRandomWalletFromXPrv ctx body = do
     void $ liftHandler $ initWorker @_ @s @k ctx wid
         (\wrk -> W.createWallet @(WorkerCtx ctx) @s @k wrk wid wName s)
         (\wrk -> W.restoreWallet @(WorkerCtx ctx) @s @t @k wrk wid)
+        (\_ -> pure ())
     withWorkerCtx ctx wid liftE liftE $ \wrk -> liftHandler $
         W.attachPrivateKeyFromPwdHash wrk wid (byronKey, pwd)
     fst <$> getWallet ctx mkLegacyWallet (ApiT wid)
@@ -1454,8 +1460,10 @@ initWorker
         -- ^ Create action
     -> (WorkerCtx ctx -> ExceptT ErrNoSuchWallet IO ())
         -- ^ Restore action
+    -> (WorkerCtx ctx -> IO ())
+        -- ^ Action to run concurrently with restore
     -> ExceptT ErrCreateWallet IO WalletId
-initWorker ctx wid createWallet restoreWallet =
+initWorker ctx wid createWallet restoreWallet coworker =
     liftIO (Registry.lookup re wid) >>= \case
         Just _ ->
             throwE $ ErrCreateWalletAlreadyExists $ ErrWalletAlreadyExists wid
@@ -1475,7 +1483,9 @@ initWorker ctx wid createWallet restoreWallet =
         , workerMain = \ctx' _ -> do
             -- FIXME:
             -- Review error handling here
-            unsafeRunExceptT $ restoreWallet ctx'
+            concurrently_
+                (unsafeRunExceptT $ restoreWallet ctx')
+                (coworker ctx')
 
         , workerAfter =
             defaultWorkerAfter
@@ -1591,11 +1601,13 @@ newApiLayer
     -> NetworkLayer IO t Block
     -> TransactionLayer t k
     -> DBFactory IO s k
+    -> (WorkerCtx ctx -> WalletId -> IO ())
+        -- ^ Action to run concurrently with wallet restore
     -> IO ctx
-newApiLayer tr g0 nw tl df = do
+newApiLayer tr g0 nw tl df coworker = do
     re <- Registry.empty
     let ctx = ApiLayer tr g0 nw tl df re
-    listDatabases df >>= mapM_ (registerWorker ctx)
+    listDatabases df >>= mapM_ (registerWorker ctx coworker)
     return ctx
 
 -- | Register a restoration worker to the registry.
@@ -1606,9 +1618,10 @@ registerWorker
         , IsOurs s Address
         )
     => ApiLayer s t k
+    -> (WorkerCtx ctx -> WalletId -> IO ())
     -> WalletId
     -> IO ()
-registerWorker ctx wid =
+registerWorker ctx coworker wid =
     void $ Registry.register @_ @ctx re ctx wid config
   where
     (_, NetworkParameters gp _, _) = ctx ^. genesisData
@@ -1622,8 +1635,9 @@ registerWorker ctx wid =
         , workerMain = \ctx' _ -> do
             -- FIXME:
             -- Review error handling here
-            unsafeRunExceptT $
-                W.restoreWallet @(WorkerCtx ctx) @s @t ctx' wid
+            concurrently_
+                (unsafeRunExceptT $ W.restoreWallet @(WorkerCtx ctx) @s @t ctx' wid)
+                (coworker ctx' wid)
 
         , workerAfter =
             defaultWorkerAfter