Add way to fetch un-cached reward balances & use from listStakeKeys

1. Rename existing getAccountBalance to getCachedAccountBalance to make
its behaviour clearer.

2. Add new `fetchAccountBalances` to `NetworkLayer`.

3. Remove unused `ErrGetAccountBalance` error

The error wasn't used, so it was just a cause for confusion and
complexity.

4. Make listStakeKeys use use fetchAccountBalances instead of
`getCachedAccountBalance`.

getCachedAccountBalance will return 0 the first time it is called for a
new stake key. This might not be appropriate for listing stake keys.

Let's block and wait for new values from the node instead.

If we see any performance problems in the future we can re-consider / go
for a third option.

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 $ getCachedAccountBalance 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, fetchAccountBalances, 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
+                (fetchAccountBalances 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
+    , getCachedAccountBalance
         :: 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.
+
+    , fetchAccountBalances
+        :: 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"
+    , getCachedAccountBalance = error "getCachedAccountBalance: not implemented"
+    , fetchAccountBalances = error "fetchAccountBalances: 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 =
+        , getCachedAccountBalance =
             _getAccountBalance rewardsObserver
+        , fetchAccountBalances =
+            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
+    _getAccountBalance rewardsObserver k = do
         startObserving rewardsObserver k
         fromMaybe (W.Coin 0) <$> query rewardsObserver k
 
@@ -890,6 +892,7 @@ fetchRewardAccounts tr queryRewardQ accounts = do
             Map.map fromShelleyCoin rewardAccounts
         , [MsgAccountDelegationAndRewards deleg rewardAccounts]
         )
+
 data ObserverLog key value
     = MsgWillFetch (Set key)
     | MsgDidFetch (Map key value)