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
This commit is contained in:
teor 2022-04-14 23:34:31 +10:00 committed by GitHub
parent 2041d69312
commit 8e29219565
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 270 additions and 111 deletions

View File

@ -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)]

View File

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

View File

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

View File

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

View File

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

View File

@ -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 {

View File

@ -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(
"unexpected missing transaction hash: transaction must already be indexed",
);
address_transfers.update_chain_tip_with(&(
&outpoint,
created_utxo,
transaction_location,
))?;
address_transfers.update_chain_tip_with(&(&outpoint, created_utxo))?;
}
}
@ -913,13 +968,7 @@ impl
.get_mut(&receiving_address)
.expect("block has previously been applied to the chain");
let transaction_location = self
.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);
address_transfers.revert_chain_with(&(&outpoint, created_utxo), position);
// Remove this transfer if it is now empty
if address_transfers.is_empty() {

View File

@ -6,13 +6,10 @@ use mset::MultiSet;
use zebra_chain::{
amount::{Amount, NegativeAllowed},
block::Height,
transaction, transparent,
};
use crate::{
request::ContextuallyValidBlock, OutputLocation, TransactionLocation, ValidateContextError,
};
use crate::{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): &(
&transparent::OutPoint,
&transparent::OrderedUtxo,
&TransactionLocation,
),
&(outpoint, created_utxo): &(&transparent::OutPoint, &transparent::OrderedUtxo),
) -> 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): &(
&transparent::OutPoint,
&transparent::OrderedUtxo,
&TransactionLocation,
),
&(outpoint, created_utxo): &(&transparent::OutPoint, &transparent::OrderedUtxo),
_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.
///
/// # Panics
///
/// If `chain_blocks` is missing some transaction hashes from this index.
/// Some of these outputs might already be spent.
/// [`TransparentTransfers::spent_utxos`] returns spent UTXOs.
#[allow(dead_code)]
pub fn created_utxos(
&self,
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()
pub fn created_utxos(&self) -> &BTreeMap<OutputLocation, transparent::Output> {
&self.created_utxos
}
/// Returns the [`OutputLocation`]s of the spent transparent outputs sent to this address,

View File

@ -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()
}