8. feat(state): add a query function for transparent address balances (#4097)

* Make address index types consistent

* Simplify non-finalized address index updates

* Update snapshots for address index queries

* Simplify non-finalized UTXO query

* Add a query method for non-finalized address balance changes

* Add a query method for finalized state address balances

* Add a query function for address balances

* Refactor balance queries to make them repeatable

* Retry interrupted finalized balance queries

* Pop chain root blocks until it matches the finalized tip

* Avoid cloning the chain

It has already been cloned by the watch receiver

* Refactor and fix documentation of the balance query code

zebra-chain/src/amount.rs

@@ -22,7 +22,8 @@ pub mod arbitrary;
 #[cfg(test)]
 mod tests;
 
-type Result<T, E = Error> = std::result::Result<T, E>;
+/// The result of an amount operation.
+pub type Result<T, E = Error> = std::result::Result<T, E>;
 
 /// A runtime validated type for representing amounts of zatoshis
 #[derive(Clone, Copy, Serialize, Deserialize)]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/address_utxo_data@mainnet_1.snap

@@ -4,13 +4,9 @@ expression: stored_address_utxos
 ---
 [
   ("t3Vz22vK5z2LcKEdg16Yv4FFneEL1zg9ojd", [
-    Utxo(
+    Output(
-      output: Output(
+      value: Amount(12500),
-        value: Amount(12500),
+      lock_script: Script("a9147d46a730d31f97b1930d3368a967c309bd4d136a87"),
-        lock_script: Script("a9147d46a730d31f97b1930d3368a967c309bd4d136a87"),
-      ),
-      height: Height(1),
-      from_coinbase: true,
     ),
   ]),
 ]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/address_utxo_data@mainnet_2.snap

@@ -4,21 +4,13 @@ expression: stored_address_utxos
 ---
 [
   ("t3Vz22vK5z2LcKEdg16Yv4FFneEL1zg9ojd", [
-    Utxo(
+    Output(
-      output: Output(
+      value: Amount(12500),
-        value: Amount(12500),
+      lock_script: Script("a9147d46a730d31f97b1930d3368a967c309bd4d136a87"),
-        lock_script: Script("a9147d46a730d31f97b1930d3368a967c309bd4d136a87"),
-      ),
-      height: Height(1),
-      from_coinbase: true,
     ),
-    Utxo(
+    Output(
-      output: Output(
+      value: Amount(25000),
-        value: Amount(25000),
+      lock_script: Script("a9147d46a730d31f97b1930d3368a967c309bd4d136a87"),
-        lock_script: Script("a9147d46a730d31f97b1930d3368a967c309bd4d136a87"),
-      ),
-      height: Height(2),
-      from_coinbase: true,
     ),
   ]),
 ]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/address_utxo_data@testnet_1.snap

@@ -4,13 +4,9 @@ expression: stored_address_utxos
 ---
 [
   ("t2UNzUUx8mWBCRYPRezvA363EYXyEpHokyi", [
-    Utxo(
+    Output(
-      output: Output(
+      value: Amount(12500),
-        value: Amount(12500),
+      lock_script: Script("a914ef775f1f997f122a062fff1a2d7443abd1f9c64287"),
-        lock_script: Script("a914ef775f1f997f122a062fff1a2d7443abd1f9c64287"),
-      ),
-      height: Height(1),
-      from_coinbase: true,
     ),
   ]),
 ]

zebra-state/src/service/finalized_state/zebra_db/block/tests/snapshots/address_utxo_data@testnet_2.snap

@@ -4,21 +4,13 @@ expression: stored_address_utxos
 ---
 [
   ("t2UNzUUx8mWBCRYPRezvA363EYXyEpHokyi", [
-    Utxo(
+    Output(
-      output: Output(
+      value: Amount(12500),
-        value: Amount(12500),
+      lock_script: Script("a914ef775f1f997f122a062fff1a2d7443abd1f9c64287"),
-        lock_script: Script("a914ef775f1f997f122a062fff1a2d7443abd1f9c64287"),
-      ),
-      height: Height(1),
-      from_coinbase: true,
     ),
-    Utxo(
+    Output(
-      output: Output(
+      value: Amount(25000),
-        value: Amount(25000),
+      lock_script: Script("a914ef775f1f997f122a062fff1a2d7443abd1f9c64287"),
-        lock_script: Script("a914ef775f1f997f122a062fff1a2d7443abd1f9c64287"),
-      ),
-      height: Height(2),
-      from_coinbase: true,
     ),
   ]),
 ]

zebra-state/src/service/finalized_state/zebra_db/transparent.rs

@@ -11,10 +11,10 @@
 //! The [`crate::constants::DATABASE_FORMAT_VERSION`] constant must
 //! be incremented each time the database format (column, serialization, etc) changes.
 
-use std::collections::{BTreeMap, BTreeSet, HashMap};
+use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet};
 
 use zebra_chain::{
-    amount::{Amount, NonNegative},
+    amount::{self, Amount, NonNegative},
     transaction, transparent,
 };
 
@@ -49,7 +49,6 @@ impl ZebraDb {
 
     /// Returns the balance for a [`transparent::Address`],
     /// if it is in the finalized state.
-    #[allow(dead_code)]
     pub fn address_balance(&self, address: &transparent::Address) -> Option<Amount<NonNegative>> {
         self.address_balance_location(address)
             .map(|abl| abl.balance())
@@ -108,7 +107,7 @@ impl ZebraDb {
     pub fn address_utxos(
         &self,
         address: &transparent::Address,
-    ) -> BTreeMap<OutputLocation, transparent::Utxo> {
+    ) -> BTreeMap<OutputLocation, transparent::Output> {
         let address_location = match self.address_location(address) {
             Some(address_location) => address_location,
             None => return BTreeMap::new(),
@@ -123,7 +122,8 @@ impl ZebraDb {
                 Some((
                     addr_out_loc.unspent_output_location(),
                     self.utxo_by_location(addr_out_loc.unspent_output_location())?
-                        .utxo,
+                        .utxo
+                        .output,
                 ))
             })
             .collect()
@@ -244,6 +244,32 @@ impl ZebraDb {
 
         addr_transactions
     }
+
+    // Address index queries
+
+    /// Returns the total transparent balance for `addresses` in the finalized chain.
+    ///
+    /// If none of the addresses has a balance, returns zero.
+    ///
+    /// # Correctness
+    ///
+    /// Callers should apply the non-finalized balance change for `addresses` to the returned balance.
+    ///
+    /// The total balance will only be correct if the non-finalized chain matches the finalized state.
+    /// Specifically, the root of the partial non-finalized chain must be a child block of the finalized tip.
+    pub fn partial_finalized_transparent_balance(
+        &self,
+        addresses: &HashSet<transparent::Address>,
+    ) -> Amount<NonNegative> {
+        let balance: amount::Result<Amount<NonNegative>> = addresses
+            .iter()
+            .filter_map(|address| self.address_balance(address))
+            .sum();
+
+        balance.expect(
+            "unexpected amount overflow: value balances are valid, so partial sum should be valid",
+        )
+    }
 }
 
 impl DiskWriteBatch {

zebra-state/src/service/non_finalized_state/chain.rs

@@ -12,7 +12,7 @@ use mset::MultiSet;
 use tracing::instrument;
 
 use zebra_chain::{
-    amount::{NegativeAllowed, NonNegative},
+    amount::{Amount, NegativeAllowed, NonNegative},
     block,
     history_tree::HistoryTree,
     orchard,
@@ -236,7 +236,7 @@ impl Chain {
     /// Remove the lowest height block of the non-finalized portion of a chain.
     #[instrument(level = "debug", skip(self))]
     pub(crate) fn pop_root(&mut self) -> ContextuallyValidBlock {
-        let block_height = self.lowest_height();
+        let block_height = self.non_finalized_root_height();
 
         // remove the lowest height block from self.blocks
         let block = self
@@ -251,7 +251,8 @@ impl Chain {
         block
     }
 
-    fn lowest_height(&self) -> block::Height {
+    /// Returns the height of the chain root.
+    pub fn non_finalized_root_height(&self) -> block::Height {
         self.blocks
             .keys()
             .next()
@@ -383,6 +384,15 @@ impl Chain {
             .get(tx_loc.index.as_usize())
     }
 
+    /// Returns the non-finalized tip block hash and height.
+    #[allow(dead_code)]
+    pub fn non_finalized_tip(&self) -> (block::Hash, block::Height) {
+        (
+            self.non_finalized_tip_hash(),
+            self.non_finalized_tip_height(),
+        )
+    }
+
     /// Returns the block hash of the tip block.
     pub fn non_finalized_tip_hash(&self) -> block::Hash {
         self.blocks
@@ -392,6 +402,15 @@ impl Chain {
             .hash
     }
 
+    /// Returns the non-finalized root block hash and height.
+    #[allow(dead_code)]
+    pub fn non_finalized_root(&self) -> (block::Hash, block::Height) {
+        (
+            self.non_finalized_root_hash(),
+            self.non_finalized_root_height(),
+        )
+    }
+
     /// Returns the block hash of the non-finalized root block.
     pub fn non_finalized_root_hash(&self) -> block::Hash {
         self.blocks
@@ -403,7 +422,7 @@ impl Chain {
 
     /// Returns the block hash of the `n`th block from the non-finalized root.
     ///
-    /// This is the block at `lowest_height() + n`.
+    /// This is the block at `non_finalized_root_height() + n`.
     #[allow(dead_code)]
     pub fn non_finalized_nth_hash(&self, n: usize) -> Option<block::Hash> {
         self.blocks.values().nth(n).map(|block| block.hash)
@@ -472,6 +491,52 @@ impl Chain {
         unspent_utxos
     }
 
+    // Address index queries
+
+    /// Returns the transparent transfers for `addresses` in this non-finalized chain.
+    ///
+    /// If none of the addresses have an address index, returns an empty iterator.
+    ///
+    /// # Correctness
+    ///
+    /// Callers should apply the returned indexes to the corresponding finalized state indexes.
+    ///
+    /// The combined result will only be correct if the chains match.
+    /// The exact type of match varies by query.
+    pub fn partial_transparent_indexes<'a>(
+        &'a self,
+        addresses: &'a HashSet<transparent::Address>,
+    ) -> impl Iterator<Item = &TransparentTransfers> {
+        addresses
+            .iter()
+            .copied()
+            .flat_map(|address| self.partial_transparent_transfers.get(&address))
+    }
+
+    /// Returns the transparent balance change for `addresses` in this non-finalized chain.
+    ///
+    /// If the balance doesn't change for any of the addresses, returns zero.
+    ///
+    /// # Correctness
+    ///
+    /// Callers should apply this balance change to the finalized state balance for `addresses`.
+    ///
+    /// The total balance will only be correct if this partial chain matches the finalized state.
+    /// Specifically, the root of this partial chain must be a child block of the finalized tip.
+    pub fn partial_transparent_balance_change(
+        &self,
+        addresses: &HashSet<transparent::Address>,
+    ) -> Amount<NegativeAllowed> {
+        let balance_change: Result<Amount<NegativeAllowed>, _> = self
+            .partial_transparent_indexes(addresses)
+            .map(|transfers| transfers.balance())
+            .sum();
+
+        balance_change.expect(
+            "unexpected amount overflow: value balances are valid, so partial sum should be valid",
+        )
+    }
+
     /// Clone the Chain but not the history and note commitment trees, using
     /// the specified trees instead.
     ///
@@ -762,10 +827,8 @@ impl UpdateWith<ContextuallyValidBlock> for Chain {
             };
 
             // remove the utxos this produced
-            // uses `tx_by_hash`
             self.revert_chain_with(&(outputs, transaction_hash, new_outputs), position);
             // remove the utxos this consumed
-            // uses `tx_by_hash`
             self.revert_chain_with(&(inputs, transaction_hash, spent_outputs), position);
 
             // remove `transaction.hash` from `tx_by_hash`
@@ -866,15 +929,7 @@ impl
                     .entry(receiving_address)
                     .or_default();
 
-                let transaction_location = self.tx_by_hash.get(&outpoint.hash).expect(
+                address_transfers.update_chain_tip_with(&(&outpoint, created_utxo))?;
-                    "unexpected missing transaction hash: transaction must already be indexed",
-                );
-
-                address_transfers.update_chain_tip_with(&(
-                    &outpoint,
-                    created_utxo,
-                    transaction_location,
-                ))?;
             }
         }
 
@@ -913,13 +968,7 @@ impl
                     .get_mut(&receiving_address)
                     .expect("block has previously been applied to the chain");
 
-                let transaction_location = self
+                address_transfers.revert_chain_with(&(&outpoint, created_utxo), position);
-                    .tx_by_hash
-                    .get(&outpoint.hash)
-                    .expect("transaction is reverted after its UTXOs are reverted");
-
-                address_transfers
-                    .revert_chain_with(&(&outpoint, created_utxo, transaction_location), position);
 
                 // Remove this transfer if it is now empty
                 if address_transfers.is_empty() {

zebra-state/src/service/non_finalized_state/chain/index.rs

@@ -6,13 +6,10 @@ use mset::MultiSet;
 
 use zebra_chain::{
     amount::{Amount, NegativeAllowed},
-    block::Height,
     transaction, transparent,
 };
 
-use crate::{
+use crate::{OutputLocation, TransactionLocation, ValidateContextError};
-    request::ContextuallyValidBlock, OutputLocation, TransactionLocation, ValidateContextError,
-};
 
 use super::{RevertPosition, UpdateWith};
 
@@ -67,23 +64,20 @@ impl
         // The location of the UTXO
         &transparent::OutPoint,
         // The UTXO data
+        // Includes the location of the transaction that created the output
         &transparent::OrderedUtxo,
-        // The location of the transaction that creates the UTXO
-        &TransactionLocation,
     )> for TransparentTransfers
 {
     fn update_chain_tip_with(
         &mut self,
-        &(outpoint, created_utxo, transaction_location): &(
+        &(outpoint, created_utxo): &(&transparent::OutPoint, &transparent::OrderedUtxo),
-            &transparent::OutPoint,
-            &transparent::OrderedUtxo,
-            &TransactionLocation,
-        ),
     ) -> Result<(), ValidateContextError> {
         self.balance =
             (self.balance + created_utxo.utxo.output.value().constrain().unwrap()).unwrap();
 
-        let output_location = OutputLocation::from_outpoint(*transaction_location, outpoint);
+        let transaction_location = transaction_location(created_utxo);
+        let output_location = OutputLocation::from_outpoint(transaction_location, outpoint);
+
         let previous_entry = self
             .created_utxos
             .insert(output_location, created_utxo.utxo.output.clone());
@@ -99,17 +93,15 @@ impl
 
     fn revert_chain_with(
         &mut self,
-        &(outpoint, created_utxo, transaction_location): &(
+        &(outpoint, created_utxo): &(&transparent::OutPoint, &transparent::OrderedUtxo),
-            &transparent::OutPoint,
-            &transparent::OrderedUtxo,
-            &TransactionLocation,
-        ),
         _position: RevertPosition,
     ) {
         self.balance =
             (self.balance - created_utxo.utxo.output.value().constrain().unwrap()).unwrap();
 
-        let output_location = OutputLocation::from_outpoint(*transaction_location, outpoint);
+        let transaction_location = transaction_location(created_utxo);
+        let output_location = OutputLocation::from_outpoint(transaction_location, outpoint);
+
         let removed_entry = self.created_utxos.remove(&output_location);
         assert!(
             removed_entry.is_some(),
@@ -240,29 +232,14 @@ impl TransparentTransfers {
             .collect()
     }
 
-    /// Returns the unspent transparent outputs sent to this address,
+    /// Returns the new transparent outputs sent to this address,
     /// in this partial chain, in chain order.
     ///
-    /// `chain_blocks` should be the `blocks` field from the [`Chain`] containing this index.
+    /// Some of these outputs might already be spent.
-    ///
+    /// [`TransparentTransfers::spent_utxos`] returns spent UTXOs.
-    /// # Panics
-    ///
-    /// If `chain_blocks` is missing some transaction hashes from this index.
     #[allow(dead_code)]
-    pub fn created_utxos(
+    pub fn created_utxos(&self) -> &BTreeMap<OutputLocation, transparent::Output> {
-        &self,
+        &self.created_utxos
-        chain_blocks: &BTreeMap<Height, ContextuallyValidBlock>,
-    ) -> BTreeMap<OutputLocation, (transparent::Output, transaction::Hash)> {
-        self.created_utxos
-            .iter()
-            .map(|(output_location, output)| {
-                let tx_loc = output_location.transaction_location();
-                let transaction_hash =
-                    chain_blocks[&tx_loc.height].transaction_hashes[tx_loc.index.as_usize()];
-
-                (*output_location, (output.clone(), transaction_hash))
-            })
-            .collect()
     }
 
     /// Returns the [`OutputLocation`]s of the spent transparent outputs sent to this address,

zebra-state/src/service/read.rs

@@ -4,21 +4,30 @@
 //! to read from the best [`Chain`] in the [`NonFinalizedState`],
 //! and the database in the [`FinalizedState`].
 
-use std::sync::Arc;
+use std::{collections::HashSet, sync::Arc};
 
 use zebra_chain::{
-    block::{self, Block},
+    amount::{self, Amount, NegativeAllowed, NonNegative},
+    block::{self, Block, Height},
     transaction::{self, Transaction},
+    transparent,
 };
 
 use crate::{
     service::{finalized_state::ZebraDb, non_finalized_state::Chain},
-    HashOrHeight,
+    BoxError, HashOrHeight,
 };
 
 #[cfg(test)]
 mod tests;
 
+/// If the transparent address index queries are interrupted by a new finalized block,
+/// retry this many times.
+///
+/// Once we're at the tip, we expect up to 2 blocks to arrive at the same time.
+/// If any more arrive, the client should wait until we're synchronised with our peers.
+const FINALIZED_ADDRESS_INDEX_RETRIES: usize = 3;
+
 /// Returns the [`Block`] with [`block::Hash`](zebra_chain::block::Hash) or
 /// [`Height`](zebra_chain::block::Height),
 /// if it exists in the non-finalized `chain` or finalized `db`.
@@ -73,3 +82,124 @@ where
         })
         .or_else(|| db.transaction(hash))
 }
+
+/// Returns the total transparent balance for the supplied [`transparent::Address`]es.
+///
+/// If the addresses do not exist in the non-finalized `chain` or finalized `db`, returns zero.
+#[allow(dead_code)]
+pub(crate) fn transparent_balance(
+    chain: Option<Arc<Chain>>,
+    db: &ZebraDb,
+    addresses: HashSet<transparent::Address>,
+) -> Result<Amount<NonNegative>, BoxError> {
+    let mut balance_result = finalized_transparent_balance(db, &addresses);
+
+    // Retry the finalized balance query if it was interruped by a finalizing block
+    for _ in 0..FINALIZED_ADDRESS_INDEX_RETRIES {
+        if balance_result.is_ok() {
+            break;
+        }
+
+        balance_result = finalized_transparent_balance(db, &addresses);
+    }
+
+    let (mut balance, finalized_tip) = balance_result?;
+
+    // Apply the non-finalized balance changes
+    if let Some(chain) = chain {
+        let chain_balance_change =
+            chain_transparent_balance_change(chain, &addresses, finalized_tip);
+
+        balance = apply_balance_change(balance, chain_balance_change).expect(
+            "unexpected amount overflow: value balances are valid, so partial sum should be valid",
+        );
+    }
+
+    Ok(balance)
+}
+
+/// Returns the total transparent balance for `addresses` in the finalized chain,
+/// and the finalized tip height the balances were queried at.
+///
+/// If the addresses do not exist in the finalized `db`, returns zero.
+//
+// TODO: turn the return type into a struct?
+fn finalized_transparent_balance(
+    db: &ZebraDb,
+    addresses: &HashSet<transparent::Address>,
+) -> Result<(Amount<NonNegative>, Option<Height>), BoxError> {
+    // # Correctness
+    //
+    // The StateService can commit additional blocks while we are querying address balances.
+
+    // Check if the finalized state changed while we were querying it
+    let original_finalized_tip = db.tip();
+
+    let finalized_balance = db.partial_finalized_transparent_balance(addresses);
+
+    let finalized_tip = db.tip();
+
+    if original_finalized_tip != finalized_tip {
+        // Correctness: Some balances might be from before the block, and some after
+        return Err("unable to get balance: state was committing a block".into());
+    }
+
+    let finalized_tip = finalized_tip.map(|(height, _hash)| height);
+
+    Ok((finalized_balance, finalized_tip))
+}
+
+/// Returns the total transparent balance change for `addresses` in the non-finalized chain,
+/// matching the balance for the `finalized_tip`.
+///
+/// If the addresses do not exist in the non-finalized `chain`, returns zero.
+fn chain_transparent_balance_change(
+    mut chain: Arc<Chain>,
+    addresses: &HashSet<transparent::Address>,
+    finalized_tip: Option<Height>,
+) -> Amount<NegativeAllowed> {
+    // # Correctness
+    //
+    // The StateService commits blocks to the finalized state before updating the latest chain,
+    // and it can commit additional blocks after we've cloned this `chain` variable.
+
+    // Check if the finalized and non-finalized states match
+    let required_chain_root = finalized_tip
+        .map(|tip| (tip + 1).unwrap())
+        .unwrap_or(Height(0));
+
+    let chain = Arc::make_mut(&mut chain);
+
+    assert!(
+        chain.non_finalized_root_height() <= required_chain_root,
+        "unexpected chain gap: the best chain is updated after its previous root is finalized"
+    );
+
+    let chain_tip = chain.non_finalized_tip_height();
+
+    // If we've already committed this entire chain, ignore its balance changes.
+    // This is more likely if the non-finalized state is just getting started.
+    if chain_tip < required_chain_root {
+        return Amount::zero();
+    }
+
+    // Correctness: some balances might have duplicate creates or spends,
+    // so we pop root blocks from `chain` until the chain root is a child of the finalized tip.
+    while chain.non_finalized_root_height() < required_chain_root {
+        // TODO: just revert the transparent balances, to improve performance
+        chain.pop_root();
+    }
+
+    chain.partial_transparent_balance_change(addresses)
+}
+
+/// Add the supplied finalized and non-finalized balances together,
+/// and return the result.
+fn apply_balance_change(
+    finalized_balance: Amount<NonNegative>,
+    chain_balance_change: Amount<NegativeAllowed>,
+) -> amount::Result<Amount<NonNegative>> {
+    let balance = finalized_balance.constrain()? + chain_balance_change;
+
+    balance?.constrain()
+}