Adds received field to getaddressbalance response

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.

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?;

zebra-state/src/response.rs

@@ -221,8 +221,14 @@ pub enum ReadResponse {
         BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>>,
     ),
 
-    /// Response to [`ReadRequest::AddressBalance`] with the total balance of the addresses.
-    AddressBalance(Amount<NonNegative>),
+    /// 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<NonNegative>,
+        /// The total received funds, including change.
+        received: Amount<NonNegative>,
+    },
 
     /// Response to [`ReadRequest::TransactionIdsByAddresses`]
     /// with the obtained transaction ids, in the order they appear in blocks.
@@ -344,7 +350,7 @@ impl TryFrom<ReadResponse> for Response {
             | ReadResponse::OrchardTree(_)
             | ReadResponse::SaplingSubtrees(_)
             | ReadResponse::OrchardSubtrees(_)
-            | ReadResponse::AddressBalance(_)
+            | ReadResponse::AddressBalance { .. }
             | ReadResponse::AddressesTransactionIds(_)
             | ReadResponse::AddressUtxos(_)
             | ReadResponse::ChainInfo(_) => {

zebra-state/src/service.rs

@@ -1691,20 +1691,20 @@ impl Service<ReadRequest> 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()

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

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<Amount<NonNegative>> {
+    pub fn address_balance(
+        &self,
+        address: &transparent::Address,
+    ) -> Option<(Amount<NonNegative>, Amount<NonNegative>)> {
         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<transparent::Address>,
-    ) -> Amount<NonNegative> {
-        let balance: amount::Result<Amount<NonNegative>> = addresses
+    ) -> (Amount<NonNegative>, Amount<NonNegative>) {
+        let balance: amount::Result<(Amount<NonNegative>, Amount<NonNegative>)> = 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",

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<Arc<Chain>>,
     db: &ZebraDb,
     addresses: HashSet<transparent::Address>,
-) -> Result<Amount<NonNegative>, BoxError> {
+) -> Result<(Amount<NonNegative>, 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
@@ -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<transparent::Address>,
-) -> Result<(Amount<NonNegative>, Option<Height>), BoxError> {
+) -> Result<((Amount<NonNegative>, Amount<NonNegative>), Option<Height>), 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<NonNegative>,
-    chain_balance_change: Amount<NegativeAllowed>,
-) -> amount::Result<Amount<NonNegative>> {
+    (finalized_balance, finalized_received): (Amount<NonNegative>, Amount<NonNegative>),
+    (chain_balance_change, chain_received_change): (Amount<NegativeAllowed>, Amount<NonNegative>),
+) -> amount::Result<(Amount<NonNegative>, Amount<NonNegative>)> {
     let balance = finalized_balance.constrain()? + chain_balance_change;
+    let received = finalized_received.constrain()? + chain_received_change;
 
-    balance?.constrain()
+    Ok((balance?.constrain()?, received?.constrain()?))
 }