zebra_state/service/read/address/
balance.rs

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)?;
58    }
59
60    Ok(balance)
61}
62
63/// Returns the total transparent balance for `addresses` in the finalized chain,
64/// and the finalized tip height the balances were queried at.
65///
66/// If the addresses do not exist in the finalized `db`, returns zero.
67//
68// TODO: turn the return type into a struct?
69fn finalized_transparent_balance(
70    db: &ZebraDb,
71    addresses: &HashSet<transparent::Address>,
72) -> Result<((Amount<NonNegative>, u64), Option<Height>), BoxError> {
73    // # Correctness
74    //
75    // The StateService can commit additional blocks while we are querying address balances.
76
77    // Check if the finalized state changed while we were querying it
78    let original_finalized_tip = db.tip();
79
80    let finalized_balance = db.partial_finalized_transparent_balance(addresses);
81
82    let finalized_tip = db.tip();
83
84    if original_finalized_tip != finalized_tip {
85        // Correctness: Some balances might be from before the block, and some after
86        return Err("unable to get balance: state was committing a block".into());
87    }
88
89    let finalized_tip = finalized_tip.map(|(height, _hash)| height);
90
91    Ok((finalized_balance, finalized_tip))
92}
93
94/// Returns the total transparent balance change for `addresses` in the non-finalized chain,
95/// matching the balance for the `finalized_tip`.
96///
97/// If the addresses do not exist in the non-finalized `chain`, returns zero.
98fn chain_transparent_balance_change(
99    mut chain: Arc<Chain>,
100    addresses: &HashSet<transparent::Address>,
101    finalized_tip: Option<Height>,
102) -> (Amount<NegativeAllowed>, u64) {
103    // # Correctness
104    //
105    // Find the balance adjustment that corrects for overlapping finalized and non-finalized blocks.
106
107    // Check if the finalized and non-finalized states match
108    let required_chain_root = finalized_tip
109        .map(|tip| (tip + 1).unwrap())
110        .unwrap_or(Height(0));
111
112    let chain = Arc::make_mut(&mut chain);
113
114    assert!(
115        chain.non_finalized_root_height() <= required_chain_root,
116        "unexpected chain gap: the best chain is updated after its previous root is finalized"
117    );
118
119    let chain_tip = chain.non_finalized_tip_height();
120
121    // If we've already committed this entire chain, ignore its balance changes.
122    // This is more likely if the non-finalized state is just getting started.
123    if chain_tip < required_chain_root {
124        return (Amount::zero(), 0);
125    }
126
127    // Correctness: some balances might have duplicate creates or spends,
128    // so we pop root blocks from `chain` until the chain root is a child of the finalized tip.
129    while chain.non_finalized_root_height() < required_chain_root {
130        // TODO: just revert the transparent balances, to improve performance
131        chain.pop_root();
132    }
133
134    chain.partial_transparent_balance_change(addresses)
135}
136
137/// Add the supplied finalized and non-finalized balances together,
138/// and return the result.
139fn apply_balance_change(
140    (finalized_balance, finalized_received): (Amount<NonNegative>, u64),
141    (chain_balance_change, chain_received_change): (Amount<NegativeAllowed>, u64),
142) -> amount::Result<(Amount<NonNegative>, u64)> {
143    let balance = finalized_balance.constrain()? + chain_balance_change;
144    // Addresses could receive more than the max money supply by sending to themselves,
145    // use u64::MAX if the addition overflows.
146    let received = finalized_received.saturating_add(chain_received_change);
147    Ok((balance?.constrain()?, received))
148}
zebra_state/service/read/address/balance.rs

zebra_state/service/read/address/
balance.rs
