zebra_state/service/read/address/
balance.rs

Help
1//! Reading address balances.
2//!
3//! In the functions in this module:
4//!
5//! The block write task commits blocks to the finalized state before updating
6//! `chain` with a cached copy of the best non-finalized chain from
7//! `NonFinalizedState.chain_set`. Then the block commit task can commit additional blocks to
8//! the finalized state after we've cloned the `chain`.
9//!
10//! This means that some blocks can be in both:
11//! - the cached [`Chain`], and
12//! - the shared finalized [`ZebraDb`] reference.
13
14use std::{collections::HashSet, sync::Arc};
15
16use zebra_chain::{
17    amount::{self, Amount, NegativeAllowed, NonNegative},
18    block::Height,
19    transparent,
20};
21
22use crate::{
23    service::{
24        finalized_state::ZebraDb, non_finalized_state::Chain, read::FINALIZED_STATE_QUERY_RETRIES,
25    },
26    BoxError,
27};
28
29/// Returns the total transparent balance and received balance for the supplied [`transparent::Address`]es.
30///
31/// If the addresses do not exist in the non-finalized `chain` or finalized `db`, returns zero.
32pub fn transparent_balance(
33    chain: Option<Arc<Chain>>,
34    db: &ZebraDb,
35    addresses: HashSet<transparent::Address>,
36) -> Result<(Amount<NonNegative>, u64), BoxError> {
37    let mut balance_result = finalized_transparent_balance(db, &addresses);
38
39    // Retry the finalized balance query if it was interrupted by a finalizing block
40    //
41    // TODO: refactor this into a generic retry(finalized_closure, process_and_check_closure) fn
42    for _ in 0..FINALIZED_STATE_QUERY_RETRIES {
43        if balance_result.is_ok() {
44            break;
45        }
46
47        balance_result = finalized_transparent_balance(db, &addresses);
48    }
49
50    let (mut balance, finalized_tip) = balance_result?;
51
52    // Apply the non-finalized balance changes
53    if let Some(chain) = chain {
54        let chain_balance_change =
55            chain_transparent_balance_change(chain, &addresses, finalized_tip);
56
57        balance = apply_balance_change(balance, chain_balance_change).expect(
58            "unexpected amount overflow: value balances are valid, so partial sum should be valid",
59        );
60    }
61
62    Ok(balance)
63}
64
65/// Returns the total transparent balance for `addresses` in the finalized chain,
66/// and the finalized tip height the balances were queried at.
67///
68/// If the addresses do not exist in the finalized `db`, returns zero.
69//
70// TODO: turn the return type into a struct?
71fn finalized_transparent_balance(
72    db: &ZebraDb,
73    addresses: &HashSet<transparent::Address>,
74) -> Result<((Amount<NonNegative>, u64), Option<Height>), BoxError> {
75    // # Correctness
76    //
77    // The StateService can commit additional blocks while we are querying address balances.
78
79    // Check if the finalized state changed while we were querying it
80    let original_finalized_tip = db.tip();
81
82    let finalized_balance = db.partial_finalized_transparent_balance(addresses);
83
84    let finalized_tip = db.tip();
85
86    if original_finalized_tip != finalized_tip {
87        // Correctness: Some balances might be from before the block, and some after
88        return Err("unable to get balance: state was committing a block".into());
89    }
90
91    let finalized_tip = finalized_tip.map(|(height, _hash)| height);
92
93    Ok((finalized_balance, finalized_tip))
94}
95
96/// Returns the total transparent balance change for `addresses` in the non-finalized chain,
97/// matching the balance for the `finalized_tip`.
98///
99/// If the addresses do not exist in the non-finalized `chain`, returns zero.
100fn chain_transparent_balance_change(
101    mut chain: Arc<Chain>,
102    addresses: &HashSet<transparent::Address>,
103    finalized_tip: Option<Height>,
104) -> (Amount<NegativeAllowed>, u64) {
105    // # Correctness
106    //
107    // Find the balance adjustment that corrects for overlapping finalized and non-finalized blocks.
108
109    // Check if the finalized and non-finalized states match
110    let required_chain_root = finalized_tip
111        .map(|tip| (tip + 1).unwrap())
112        .unwrap_or(Height(0));
113
114    let chain = Arc::make_mut(&mut chain);
115
116    assert!(
117        chain.non_finalized_root_height() <= required_chain_root,
118        "unexpected chain gap: the best chain is updated after its previous root is finalized"
119    );
120
121    let chain_tip = chain.non_finalized_tip_height();
122
123    // If we've already committed this entire chain, ignore its balance changes.
124    // This is more likely if the non-finalized state is just getting started.
125    if chain_tip < required_chain_root {
126        return (Amount::zero(), 0);
127    }
128
129    // Correctness: some balances might have duplicate creates or spends,
130    // so we pop root blocks from `chain` until the chain root is a child of the finalized tip.
131    while chain.non_finalized_root_height() < required_chain_root {
132        // TODO: just revert the transparent balances, to improve performance
133        chain.pop_root();
134    }
135
136    chain.partial_transparent_balance_change(addresses)
137}
138
139/// Add the supplied finalized and non-finalized balances together,
140/// and return the result.
141fn apply_balance_change(
142    (finalized_balance, finalized_received): (Amount<NonNegative>, u64),
143    (chain_balance_change, chain_received_change): (Amount<NegativeAllowed>, u64),
144) -> amount::Result<(Amount<NonNegative>, u64)> {
145    let balance = finalized_balance.constrain()? + chain_balance_change;
146    // Addresses could receive more than the max money supply by sending to themselves,
147    // use u64::MAX if the addition overflows.
148    let received = finalized_received.saturating_add(chain_received_change);
149    Ok((balance?.constrain()?, received))
150}
zebra_state/service/read/address/balance.rs

zebra_state/service/read/address/
balance.rs
