Merge #2694

2694: Factor `getAccountBalance` into `getCachedAccountBalance` and `fetchAccountBalances` r=Anviking a=Anviking

# Issue Number

<!-- Put here a reference to the issue that this PR relates to and which requirements it tackles. Jira issues of the form ADP- will be auto-linked. -->

Split off from #2684 


# Overview

<!-- Detail in a few bullet points the work accomplished in this PR -->

- [x] Rename `getAccountBalance` to `getCachedAccountBalance` for clarity
- [x] Add `fetchAccountBalances` function for un-cached behaviour
- [x] `fetchAccountBalances` from `listStakeKeys`  


# Comments

<!-- Additional comments or screenshots to attach if any -->

<!--
Don't forget to:

 ✓ Self-review your changes to make sure nothing unexpected slipped through
 ✓ Assign yourself to the PR
 ✓ Assign one or several reviewer(s)
 ✓ Jira will detect and link to this PR once created, but you can also link this PR in the description of the corresponding ticket
 ✓ Acknowledge any changes required to the Wiki
 ✓ Finally, in the PR description delete any empty sections and all text commented in <!--, so that this text does not appear in merge commit messages.
-->


Co-authored-by: Johannes Lund <johannes.lund@iohk.io>

lib/core/src/Cardano/Wallet.hs

@@ -216,8 +216,7 @@ import Cardano.Wallet.DB
 import Cardano.Wallet.Logging
     ( BracketLog, bracketTracer, traceWithExceptT, unliftIOTracer )
 import Cardano.Wallet.Network
-    ( ErrGetAccountBalance (..)
-    , ErrPostTx (..)
+    ( ErrPostTx (..)
     , FollowAction (..)
     , FollowExceptionRecovery (..)
     , FollowLog (..)
@@ -1057,13 +1056,9 @@ queryRewardBalance
     -> RewardAccount
     -> ExceptT ErrFetchRewards IO Coin
 queryRewardBalance ctx acct = do
-    mapExceptT (fmap handleErr) $ getAccountBalance nw acct
+    liftIO $ getCachedRewardAccountBalance nw acct
   where
     nw = ctx ^. networkLayer
-    handleErr = \case
-        Right x -> Right x
-        Left (ErrGetAccountBalanceAccountNotFound _) ->
-            Right $ Coin 0
 
 manageRewardBalance
     :: forall ctx s k (n :: NetworkDiscriminant).

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

@@ -273,7 +273,7 @@ import Cardano.Wallet.Compat
 import Cardano.Wallet.DB
     ( DBFactory (..) )
 import Cardano.Wallet.Network
-    ( NetworkLayer, getAccountBalance, timeInterpreter )
+    ( NetworkLayer, fetchRewardAccountBalances, timeInterpreter )
 import Cardano.Wallet.Primitive.AddressDerivation
     ( DelegationAddress (..)
     , Depth (..)
@@ -2086,7 +2086,7 @@ listStakeKeys'
         -- ^ The wallet's UTxO
     -> (Address -> Maybe RewardAccount)
         -- ^ Lookup reward account of addr
-    -> (Set RewardAccount -> m (Map RewardAccount (Maybe Coin)))
+    -> (Set RewardAccount -> m (Map RewardAccount Coin))
         -- ^ Batch fetch of rewards
     -> [(RewardAccount, Natural, ApiWalletDelegation)]
         -- ^ The wallet's known stake keys, along with derivation index, and
@@ -2102,13 +2102,12 @@ listStakeKeys' utxo lookupStakeRef fetchRewards ourKeysWithInfo = do
         let allKeys = ourKeys <> stakeKeysInUTxO
 
         -- If we wanted to know whether a stake key is registered or not, we
-        -- could look at the difference between @Nothing@ and
-        -- @Just (Coin 0)@ from the response here, instead of hiding the
-        -- difference.
+        -- could expose the difference between `Nothing` and `Just 0` in the
+        -- `NetworkLayer` interface.
         rewardsMap <- fetchRewards $ Set.fromList allKeys
 
         let rewards acc = fromMaybe (Coin 0) $
-                join $ Map.lookup acc rewardsMap
+                Map.lookup acc rewardsMap
 
         let mkOurs (acc, ix, deleg) = ApiOurStakeKey
                 { _index = ix
@@ -2165,18 +2164,13 @@ listStakeKeys lookupStakeRef ctx (ApiT wid) = do
                     Just acc -> [(acc, 0, ourApiDelegation)]
                     Nothing -> []
 
-            let fetchRewards = flip lookupUsing rewardsOfAccount . Set.toList
-            liftIO $ listStakeKeys' @n utxo lookupStakeRef fetchRewards ourKeys
-
+            liftIO $ listStakeKeys' @n
+                utxo
+                lookupStakeRef
+                (fetchRewardAccountBalances nl)
+                ourKeys
   where
-    lookupUsing
-        :: (Traversable t, Monad m, Ord a) => t a -> (a -> m b) -> m (Map a b)
-    lookupUsing xs f =
-        Map.fromList . F.toList <$> forM xs (\x -> f x >>= \x' -> pure (x,x') )
-
-    rewardsOfAccount :: forall m. MonadIO m => RewardAccount -> m (Maybe Coin)
-    rewardsOfAccount acc = fmap eitherToMaybe <$> liftIO . runExceptT $
-        getAccountBalance (ctx ^. networkLayer) acc
+    nl = ctx ^. networkLayer
 
 {-------------------------------------------------------------------------------
                                 Migrations

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

@@ -28,7 +28,6 @@ module Cardano.Wallet.Network
 
     -- * Errors
     , ErrPostTx (..)
-    , ErrGetAccountBalance (..)
 
     -- * Logging
     , FollowLog (..)
@@ -84,8 +83,12 @@ import Data.Functor
     ( ($>) )
 import Data.List.NonEmpty
     ( NonEmpty (..) )
+import Data.Map
+    ( Map )
 import Data.Quantity
     ( Quantity (..) )
+import Data.Set
+    ( Set )
 import Data.Text
     ( Text )
 import Data.Text.Class
@@ -172,9 +175,20 @@ data NetworkLayer m block = NetworkLayer
         :: Coin -- Stake to consider for rewards
         -> m StakePoolsSummary
 
-    , getAccountBalance
+    , getCachedRewardAccountBalance
         :: RewardAccount
-        -> ExceptT ErrGetAccountBalance m Coin
+        -> m Coin
+        -- ^ Return the cached reward balance of an account.
+        --
+        -- If there is no cached value, it will return `Coin 0`, and add the
+        -- account to the internal set of observed account, such that it will be
+        -- fetched later.
+
+    , fetchRewardAccountBalances
+        :: Set RewardAccount
+        -> m (Map RewardAccount Coin)
+        -- ^ Fetch the reward account balance of a set of accounts without
+        -- any caching.
 
     , timeInterpreter
         :: TimeInterpreter (ExceptT PastHorizonException m)
@@ -202,10 +216,6 @@ instance ToText ErrPostTx where
         ErrPostTxBadRequest msg -> msg
         ErrPostTxProtocolFailure msg -> msg
 
-newtype ErrGetAccountBalance
-    = ErrGetAccountBalanceAccountNotFound RewardAccount
-    deriving (Generic, Eq, Show)
-
 {-------------------------------------------------------------------------------
                               Initialization
 -------------------------------------------------------------------------------}

lib/core/test/shared/Cardano/Wallet/DummyTarget/Primitive/Types.hs

@@ -165,7 +165,8 @@ dummyNetworkLayer = NetworkLayer
     , currentSlottingParameters = error "currentSlottingParameters: not implemented"
     , postTx = error "postTx: not implemented"
     , stakeDistribution = error "stakeDistribution: not implemented"
-    , getAccountBalance = error "getAccountBalance: not implemented"
+    , getCachedRewardAccountBalance = error "getRewardCachedAccountBalance: not implemented"
+    , fetchRewardAccountBalances = error "fetchRewardAccountBalances: not implemented"
     , timeInterpreter = error "timeInterpreter: not implemented"
     , syncProgress = error "syncProgress: not implemented"
     }

lib/shelley/src/Cardano/Wallet/Shelley/Network.hs

@@ -355,8 +355,10 @@ withNetworkLayerBase tr np conn (versionData, _) tol action = do
             _postTx localTxSubmissionQ era sealed
         , stakeDistribution =
             _stakeDistribution queryRewardQ
-        , getAccountBalance =
-            _getAccountBalance rewardsObserver
+        , getCachedRewardAccountBalance =
+            _getCachedRewardAccountBalance rewardsObserver
+        , fetchRewardAccountBalances =
+            fetchRewardAccounts tr queryRewardQ
         , timeInterpreter =
             _timeInterpreter (contramap MsgInterpreterLog tr) interpreterVar
         , syncProgress = _syncProgress interpreterVar
@@ -545,7 +547,7 @@ withNetworkLayerBase tr np conn (versionData, _) tol action = do
 
     -- TODO(#2042): Make wallets call manually, with matching
     -- stopObserving.
-    _getAccountBalance rewardsObserver k = liftIO $ do
+    _getCachedRewardAccountBalance rewardsObserver k = do
         startObserving rewardsObserver k
         fromMaybe (W.Coin 0) <$> query rewardsObserver k
 
@@ -855,9 +857,16 @@ newRewardBalanceFetcher tr readNodeTip queryRewardQ = do
     fetch _tip accounts = do
         -- NOTE: We no longer need the tip to run LSQ queries. The local state
         -- query client will automatically acquire the latest tip.
+        Just <$> fetchRewardAccounts tr queryRewardQ accounts
 
+fetchRewardAccounts
+    :: Tracer IO NetworkLayerLog
+    -> TQueue IO (LocalStateQueryCmd (CardanoBlock StandardCrypto) IO)
+    -> Set W.RewardAccount
+    -> IO (Map W.RewardAccount W.Coin)
+fetchRewardAccounts tr queryRewardQ accounts = do
         liftIO $ traceWith tr $
-            MsgGetRewardAccountBalance accounts
+            MsgFetchRewardAccountBalance accounts
 
         let qry = byronOrShelleyBased (pure (byronValue, [])) $
                    fmap fromBalanceResult
@@ -867,11 +876,10 @@ newRewardBalanceFetcher tr readNodeTip queryRewardQ = do
 
         (res,logs) <- bracketQuery "queryRewards" tr (send queryRewardQ (SomeLSQ qry))
         liftIO $ mapM_ (traceWith tr) logs
-        return $ Just res
-
-      where
-        byronValue :: Map W.RewardAccount W.Coin
-        byronValue = Map.fromList . map (, minBound) $ Set.toList accounts
+        return res
+  where
+    byronValue :: Map W.RewardAccount W.Coin
+    byronValue = Map.fromList . map (, minBound) $ Set.toList accounts
 
     fromBalanceResult
         :: ( Map (SL.Credential 'SL.Staking crypto)
@@ -1146,7 +1154,7 @@ data NetworkLayerLog where
     MsgProtocolParameters :: W.ProtocolParameters -> W.SlottingParameters -> NetworkLayerLog
     MsgLocalStateQueryError :: QueryClientName -> String -> NetworkLayerLog
     MsgLocalStateQueryEraMismatch :: MismatchEraInfo (CardanoEras StandardCrypto) -> NetworkLayerLog
-    MsgGetRewardAccountBalance
+    MsgFetchRewardAccountBalance
         :: Set W.RewardAccount
         -> NetworkLayerLog
     MsgAccountDelegationAndRewards
@@ -1226,7 +1234,7 @@ instance ToText NetworkLayerLog where
         MsgLocalStateQueryEraMismatch mismatch ->
             "Local state query for the wrong era - this is fine. " <>
             T.pack (show mismatch)
-        MsgGetRewardAccountBalance accts -> T.unwords
+        MsgFetchRewardAccountBalance accts -> T.unwords
             [ "Querying the reward account balance for"
             , fmt $ listF accts
             ]
@@ -1295,7 +1303,7 @@ instance HasSeverityAnnotation NetworkLayerLog where
             | isSlowQuery qry dt           -> Notice
             | otherwise                    -> Debug
         MsgInterpreterLog msg              -> getSeverityAnnotation msg
-        MsgGetRewardAccountBalance{}       -> Debug
+        MsgFetchRewardAccountBalance{}       -> Debug
         MsgObserverLog (MsgDidChange _)    -> Notice
         MsgObserverLog{}                   -> Debug