Refactor internal wallet and keep the tip around

input-output-hk · Nov 30, 2022 · 4f29047 · 4f29047
1 parent 914fd52
commit 4f29047
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 28 deletions.
diff --git a/hydra-node/src/Hydra/Chain/Direct.hs b/hydra-node/src/Hydra/Chain/Direct.hs
@@ -25,6 +25,7 @@ import Cardano.Ledger.Shelley.Rules.Ledger (LedgerPredicateFailure (UtxowFailure
 import Cardano.Ledger.Shelley.Rules.Utxow (UtxowPredicateFailure (UtxoFailure))
 import Cardano.Ledger.Slot (EpochInfo)
 import Cardano.Slotting.EpochInfo (hoistEpochInfo)
+import CardanoClient (QueryPoint (..))
 import Control.Exception (IOException)
 import Control.Monad.Class.MonadSTM (
   newEmptyTMVar,
@@ -88,6 +89,7 @@ import Hydra.Chain.Direct.Util (
  )
 import Hydra.Chain.Direct.Wallet (
   TinyWallet (..),
+  WalletInfoOnChain (..),
   getTxId,
   newTinyWallet,
  )
@@ -176,6 +178,7 @@ withDirectChain ::
   Tracer IO DirectChainLog ->
   ChainConfig ->
   ChainContext ->
+  -- | Last known point on chain as loaded from persistence.
   Maybe ChainPoint ->
   ChainComponent Tx IO a
 withDirectChain tracer config ctx persistedPoint callback action = do
@@ -186,7 +189,7 @@ withDirectChain tracer config ctx persistedPoint callback action = do
     (min <$> startChainFrom <*> persistedPoint)
       <|> persistedPoint
       <|> startChainFrom
-  wallet <- newTinyWallet (contramap Wallet tracer) networkId keyPair chainPoint queryUTxOEtc
+  wallet <- newTinyWallet (contramap Wallet tracer) networkId keyPair chainPoint queryWalletInfo
   let chainHandle =
         mkChain
           tracer
@@ -215,12 +218,15 @@ withDirectChain tracer config ctx persistedPoint callback action = do
  where
   DirectChainConfig{networkId, nodeSocket, cardanoSigningKey, startChainFrom} = config
 
-  queryUTxOEtc queryPoint address = do
-    utxo <- Ledger.unUTxO . toLedgerUTxO <$> queryUTxO networkId nodeSocket queryPoint [address]
-    pparams <- toLedgerPParams (shelleyBasedEra @Api.Era) <$> queryProtocolParameters networkId nodeSocket queryPoint
-    systemStart <- querySystemStart networkId nodeSocket queryPoint
-    epochInfo <- toEpochInfo <$> queryEraHistory networkId nodeSocket queryPoint
-    pure (utxo, pparams, systemStart, epochInfo)
+  queryWalletInfo queryPoint address = do
+    point <- case queryPoint of
+      QueryAt point -> pure point
+      QueryTip -> queryTip networkId nodeSocket
+    walletUTxO <- Ledger.unUTxO . toLedgerUTxO <$> queryUTxO networkId nodeSocket (QueryAt point) [address]
+    pparams <- toLedgerPParams (shelleyBasedEra @Api.Era) <$> queryProtocolParameters networkId nodeSocket (QueryAt point)
+    systemStart <- querySystemStart networkId nodeSocket (QueryAt point)
+    epochInfo <- toEpochInfo <$> queryEraHistory networkId nodeSocket (QueryAt point)
+    pure $ WalletInfoOnChain{walletUTxO, pparams, systemStart, epochInfo, tip = point}
 
   toEpochInfo :: EraHistory CardanoMode -> EpochInfo (Either Text)
   toEpochInfo (EraHistory _ interpreter) =

diff --git a/hydra-node/src/Hydra/Chain/Direct/Wallet.hs b/hydra-node/src/Hydra/Chain/Direct/Wallet.hs
@@ -94,16 +94,16 @@ data TinyWallet m = TinyWallet
     update :: Block -> m ()
   }
 
-type ChainQuery m =
-  ( QueryPoint ->
-    Api.Address ShelleyAddr ->
-    m
-      ( Map TxIn TxOut
-      , PParams LedgerEra
-      , SystemStart
-      , EpochInfo (Either Text)
-      )
-  )
+data WalletInfoOnChain = WalletInfoOnChain
+  { walletUTxO :: Map TxIn TxOut
+  , pparams :: PParams LedgerEra
+  , systemStart :: SystemStart
+  , epochInfo :: EpochInfo (Either Text)
+  , -- | Latest point on chain the wallet knows of.
+    tip :: ChainPoint
+  }
+
+type ChainQuery m = QueryPoint -> Api.Address ShelleyAddr -> m WalletInfoOnChain
 
 -- | Get a single, marked as "fuel" UTxO.
 getFuelUTxO :: MonadSTM m => TinyWallet m -> STM m (Maybe (TxIn, TxOut))
@@ -129,28 +129,27 @@ newTinyWallet ::
   -- node. Initially and on demand later.
   ChainQuery IO ->
   IO (TinyWallet IO)
-newTinyWallet tracer networkId (vk, sk) chainPoint queryUTxOEtc = do
-  utxoVar <- newTVarIO =<< queryUTxOEtc (QueryAt chainPoint) address
+newTinyWallet tracer networkId (vk, sk) chainPoint queryWalletInfo = do
+  walletInfoVar <- newTVarIO =<< queryWalletInfo (QueryAt chainPoint) address
   pure
     TinyWallet
-      { getUTxO =
-          (\(u, _, _, _) -> u) <$> readTVar utxoVar
+      { getUTxO = readTVar walletInfoVar <&> walletUTxO
       , sign = Util.signWith sk
       , coverFee = \lookupUTxO partialTx -> do
-          (walletUTxO, pparams, systemStart, epochInfo) <- readTVar utxoVar
+          WalletInfoOnChain{walletUTxO, pparams, systemStart, epochInfo} <- readTVar walletInfoVar
           pure $ coverFee_ pparams systemStart epochInfo lookupUTxO walletUTxO partialTx
       , reset = \point -> do
           traceWith tracer $ BeginInitialize{point}
-          res@(initialUTxO, _, _, _) <- queryUTxOEtc (QueryAt point) address
-          atomically $ writeTVar utxoVar res
-          traceWith tracer $ EndInitialize{initialUTxO}
+          walletInfo <- queryWalletInfo (QueryAt point) address
+          atomically $ writeTVar walletInfoVar walletInfo
+          traceWith tracer $ EndInitialize{initialUTxO = walletUTxO walletInfo}
       , update = \block -> do
           let point = fromConsensusPointHF $ blockPoint block
           traceWith tracer $ BeginUpdate{point}
           utxo' <- atomically $ do
-            (utxo, pparams, systemStart, epochInfo) <- readTVar utxoVar
-            let utxo' = applyBlock block (== ledgerAddress) utxo
-            writeTVar utxoVar (utxo', pparams, systemStart, epochInfo)
+            walletInfo@WalletInfoOnChain{walletUTxO} <- readTVar walletInfoVar
+            let utxo' = applyBlock block (== ledgerAddress) walletUTxO
+            writeTVar walletInfoVar walletInfo
             pure utxo'
           traceWith tracer $ EndUpdate utxo'
       }