From 944f8beb3c130585564319649be16b846ed016b5 Mon Sep 17 00:00:00 2001 From: ar Date: Fri, 21 Feb 2025 22:02:41 -0500 Subject: [PATCH] Adds `received` field to `getaddressbalance` response --- zebra-rpc/src/methods.rs | 5 ++++- zebra-rpc/src/methods/tests/prop.rs | 4 ++-- zebra-state/src/response.rs | 12 ++++++++--- zebra-state/src/service.rs | 10 ++++----- .../upgrade/add_balance_received.rs | 3 +++ .../finalized_state/zebra_db/transparent.rs | 21 +++++++++++++------ .../src/service/read/address/balance.rs | 17 ++++++++------- 7 files changed, 47 insertions(+), 25 deletions(-) diff --git a/zebra-rpc/src/methods.rs b/zebra-rpc/src/methods.rs index af09203ff5d..70b68008835 100644 --- a/zebra-rpc/src/methods.rs +++ b/zebra-rpc/src/methods.rs @@ -752,8 +752,9 @@ where let response = state.oneshot(request).await.map_misc_error()?; match response { - zebra_state::ReadResponse::AddressBalance(balance) => Ok(AddressBalance { + zebra_state::ReadResponse::AddressBalance { balance, received } => Ok(AddressBalance { balance: u64::from(balance), + received: u64::from(received), }), _ => unreachable!("Unexpected response from state service: {response:?}"), } @@ -2003,6 +2004,8 @@ impl AddressStrings { pub struct AddressBalance { /// The total transparent balance. pub balance: u64, + /// The total received balance, including change. + pub received: u64, } /// A hex-encoded [`ConsensusBranchId`] string. diff --git a/zebra-rpc/src/methods/tests/prop.rs b/zebra-rpc/src/methods/tests/prop.rs index 0e8db0f36e6..5a7319264f3 100644 --- a/zebra-rpc/src/methods/tests/prop.rs +++ b/zebra-rpc/src/methods/tests/prop.rs @@ -665,7 +665,7 @@ proptest! { let state_query = state .expect_request(zebra_state::ReadRequest::AddressBalance(addresses)) .map_ok(|responder| { - responder.respond(zebra_state::ReadResponse::AddressBalance(balance)) + responder.respond(zebra_state::ReadResponse::AddressBalance { balance, received: Default::default() }) }); // Await the RPC call and the state query @@ -676,7 +676,7 @@ proptest! { // Check that response contains the expected balance let received_balance = response?; - prop_assert_eq!(received_balance, AddressBalance { balance: balance.into() }); + prop_assert_eq!(received_balance, AddressBalance { balance: balance.into(), received: Default::default() }); // Check no further requests were made during this test mempool.expect_no_requests().await?; diff --git a/zebra-state/src/response.rs b/zebra-state/src/response.rs index 0e847e9d0cf..acc32ca1035 100644 --- a/zebra-state/src/response.rs +++ b/zebra-state/src/response.rs @@ -221,8 +221,14 @@ pub enum ReadResponse { BTreeMap>, ), - /// Response to [`ReadRequest::AddressBalance`] with the total balance of the addresses. - AddressBalance(Amount), + /// Response to [`ReadRequest::AddressBalance`] with the total balance of the addresses, + /// and the total received funds, including change. + AddressBalance { + /// The total balance of the addresses. + balance: Amount, + /// The total received funds, including change. + received: Amount, + }, /// Response to [`ReadRequest::TransactionIdsByAddresses`] /// with the obtained transaction ids, in the order they appear in blocks. @@ -344,7 +350,7 @@ impl TryFrom for Response { | ReadResponse::OrchardTree(_) | ReadResponse::SaplingSubtrees(_) | ReadResponse::OrchardSubtrees(_) - | ReadResponse::AddressBalance(_) + | ReadResponse::AddressBalance { .. } | ReadResponse::AddressesTransactionIds(_) | ReadResponse::AddressUtxos(_) | ReadResponse::ChainInfo(_) => { diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs index 9e3fbfed2d4..8382d2dae2e 100644 --- a/zebra-state/src/service.rs +++ b/zebra-state/src/service.rs @@ -1691,20 +1691,20 @@ impl Service for ReadStateService { tokio::task::spawn_blocking(move || { span.in_scope(move || { - let balance = state.non_finalized_state_receiver.with_watch_data( - |non_finalized_state| { + let (balance, received) = state + .non_finalized_state_receiver + .with_watch_data(|non_finalized_state| { read::transparent_balance( non_finalized_state.best_chain().cloned(), &state.db, addresses, ) - }, - )?; + })?; // The work is done in the future. timer.finish(module_path!(), line!(), "ReadRequest::AddressBalance"); - Ok(ReadResponse::AddressBalance(balance)) + Ok(ReadResponse::AddressBalance { balance, received }) }) }) .wait_for_panics() diff --git a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_balance_received.rs b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_balance_received.rs index cd6279883f9..1d23a58a37f 100644 --- a/zebra-state/src/service/finalized_state/disk_format/upgrade/add_balance_received.rs +++ b/zebra-state/src/service/finalized_state/disk_format/upgrade/add_balance_received.rs @@ -202,6 +202,9 @@ impl DiskFormatUpgrade for AddAddressBalanceReceived { continue; } + // TODO: Retain items that don't have the right received balance instead? + // Note: This would require ignoring outputs sent to addresses that are not being tracked + // after the first disk write. let is_updated_on_disk = address_received_map.par_iter().all(|(address, &received)| { // short-circuit iteration and immediately return an error in the next iteration of the outer loop diff --git a/zebra-state/src/service/finalized_state/zebra_db/transparent.rs b/zebra-state/src/service/finalized_state/zebra_db/transparent.rs index 06a09d803e6..763e75cacb2 100644 --- a/zebra-state/src/service/finalized_state/zebra_db/transparent.rs +++ b/zebra-state/src/service/finalized_state/zebra_db/transparent.rs @@ -89,11 +89,14 @@ impl ZebraDb { self.db.zs_get(&balance_by_transparent_addr, address) } - /// Returns the balance for a [`transparent::Address`], + /// Returns the balance and received balance for a [`transparent::Address`], /// if it is in the finalized state. - pub fn address_balance(&self, address: &transparent::Address) -> Option> { + pub fn address_balance( + &self, + address: &transparent::Address, + ) -> Option<(Amount, Amount)> { self.address_balance_location(address) - .map(|abl| abl.balance()) + .map(|abl| (abl.balance(), abl.received())) } /// Returns the first output that sent funds to a [`transparent::Address`], @@ -304,11 +307,17 @@ impl ZebraDb { pub fn partial_finalized_transparent_balance( &self, addresses: &HashSet, - ) -> Amount { - let balance: amount::Result> = addresses + ) -> (Amount, Amount) { + let balance: amount::Result<(Amount, Amount)> = addresses .iter() .filter_map(|address| self.address_balance(address)) - .sum(); + .fold( + Ok((Amount::zero(), Amount::zero())), + |acc, (b_balance, b_received)| { + let (a_balance, a_received) = acc?; + Ok(((a_balance + b_balance)?, (a_received + b_received)?)) + }, + ); balance.expect( "unexpected amount overflow: value balances are valid, so partial sum should be valid", diff --git a/zebra-state/src/service/read/address/balance.rs b/zebra-state/src/service/read/address/balance.rs index 14ec49fcff1..554e9a715a0 100644 --- a/zebra-state/src/service/read/address/balance.rs +++ b/zebra-state/src/service/read/address/balance.rs @@ -26,14 +26,14 @@ use crate::{ BoxError, }; -/// Returns the total transparent balance for the supplied [`transparent::Address`]es. +/// Returns the total transparent balance and received 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>, db: &ZebraDb, addresses: HashSet, -) -> Result, BoxError> { +) -> Result<(Amount, Amount), BoxError> { let mut balance_result = finalized_transparent_balance(db, &addresses); // Retry the finalized balance query if it was interrupted by a finalizing block @@ -54,7 +54,7 @@ pub fn transparent_balance( let chain_balance_change = chain_transparent_balance_change(chain, &addresses, finalized_tip); - balance = apply_balance_change(balance, chain_balance_change).expect( + balance = apply_balance_change(balance, (chain_balance_change, Amount::zero())).expect( "unexpected amount overflow: value balances are valid, so partial sum should be valid", ); } @@ -71,7 +71,7 @@ pub fn transparent_balance( fn finalized_transparent_balance( db: &ZebraDb, addresses: &HashSet, -) -> Result<(Amount, Option), BoxError> { +) -> Result<((Amount, Amount), Option), BoxError> { // # Correctness // // The StateService can commit additional blocks while we are querying address balances. @@ -139,10 +139,11 @@ fn chain_transparent_balance_change( /// Add the supplied finalized and non-finalized balances together, /// and return the result. fn apply_balance_change( - finalized_balance: Amount, - chain_balance_change: Amount, -) -> amount::Result> { + (finalized_balance, finalized_received): (Amount, Amount), + (chain_balance_change, chain_received_change): (Amount, Amount), +) -> amount::Result<(Amount, Amount)> { let balance = finalized_balance.constrain()? + chain_balance_change; + let received = finalized_received.constrain()? + chain_received_change; - balance?.constrain() + Ok((balance?.constrain()?, received?.constrain()?)) }