-
Notifications
You must be signed in to change notification settings - Fork 94
/
balance.rs
146 lines (118 loc) · 4.99 KB
/
balance.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
//! Reading address balances.
//!
//! In the functions in this module:
//!
//! The block write task commits blocks to the finalized state before updating
//! `chain` with a cached copy of the best non-finalized chain from
//! `NonFinalizedState.chain_set`. Then the block commit task can commit additional blocks to
//! the finalized state after we've cloned the `chain`.
//!
//! This means that some blocks can be in both:
//! - the cached [`Chain`], and
//! - the shared finalized [`ZebraDb`] reference.
use std::{collections::HashSet, sync::Arc};
use zebra_chain::{
amount::{self, Amount, NegativeAllowed, NonNegative},
block::Height,
transparent,
};
use crate::{
service::{finalized_state::ZebraDb, non_finalized_state::Chain},
BoxError,
};
use super::FINALIZED_ADDRESS_INDEX_RETRIES;
/// 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.
pub 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 interrupted 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
//
// Find the balance adjustment that corrects for overlapping finalized and non-finalized blocks.
// 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()
}