change(rpc): Update getaddressbalance RPC to return received field

zebra-consensus/src/transaction/tests.rs

@@ -712,7 +712,7 @@ async fn mempool_request_with_unmined_output_spends_is_accepted() {
     );
 }
 
-#[tokio::test]
+#[tokio::test(flavor = "multi_thread")]
 async fn skips_verification_of_block_transactions_in_mempool() {
     let mut state: MockService<_, _, _, _> = MockService::build().for_prop_tests();
     let mempool: MockService<_, _, _, _> = MockService::build().for_prop_tests();
@@ -797,6 +797,9 @@ async fn skips_verification_of_block_transactions_in_mempool() {
             .respond(mempool::Response::UnspentOutput(output));
     });
 
+    // Briefly yield and sleep so the spawned task can first expect an await output request.
+    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+
     let verifier_response = verifier
         .clone()
         .oneshot(Request::Mempool {
@@ -847,6 +850,9 @@ async fn skips_verification_of_block_transactions_in_mempool() {
         time: Utc::now(),
     };
 
+    // Briefly yield and sleep so the spawned task can first expect the requests.
+    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+
     let crate::transaction::Response::Block { .. } = verifier
         .clone()
         .oneshot(make_request.clone()(Arc::new([input_outpoint.hash].into())))

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,
             }),
             _ => unreachable!("Unexpected response from state service: {response:?}"),
         }
@@ -1929,11 +1930,33 @@ impl GetBlockChainInfo {
 /// This is used for the input parameter of [`RpcServer::get_address_balance`],
 /// [`RpcServer::get_address_tx_ids`] and [`RpcServer::get_address_utxos`].
 #[derive(Clone, Debug, Eq, PartialEq, Hash, serde::Deserialize)]
+#[serde(from = "DAddressStrings")]
 pub struct AddressStrings {
     /// A list of transparent address strings.
     addresses: Vec<String>,
 }
 
+impl From<DAddressStrings> for AddressStrings {
+    fn from(address_strings: DAddressStrings) -> Self {
+        match address_strings {
+            DAddressStrings::Addresses(addresses) => AddressStrings { addresses },
+            DAddressStrings::Address(address) => AddressStrings {
+                addresses: vec![address],
+            },
+        }
+    }
+}
+
+/// An intermediate type used to deserialize [`AddressStrings`].
+#[derive(Clone, Debug, Eq, PartialEq, Hash, serde::Deserialize)]
+#[serde(untagged)]
+enum DAddressStrings {
+    /// A list of address strings.
+    Addresses(Vec<String>),
+    /// A single address string.
+    Address(String),
+}
+
 impl AddressStrings {
     /// Creates a new `AddressStrings` given a vector.
     #[cfg(test)]
@@ -1981,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/constants.rs

@@ -55,7 +55,7 @@ const DATABASE_FORMAT_VERSION: u64 = 26;
 /// - adding new column families,
 /// - changing the format of a column family in a compatible way, or
 /// - breaking changes with compatibility code in all supported Zebra versions.
-const DATABASE_FORMAT_MINOR_VERSION: u64 = 0;
+const DATABASE_FORMAT_MINOR_VERSION: u64 = 1;
 
 /// The database format patch version, incremented each time the on-disk database format has a
 /// significant format compatibility fix.
@@ -78,13 +78,6 @@ pub fn state_database_format_version_in_code() -> Version {
     }
 }
 
-/// Returns the highest database version that modifies the subtree index format.
-///
-/// This version is used by tests to wait for the subtree upgrade to finish.
-pub fn latest_version_for_adding_subtrees() -> Version {
-    Version::parse("25.2.2").expect("Hardcoded version string should be valid.")
-}
-
 /// The name of the file containing the minor and patch database versions.
 ///
 /// Use [`Config::version_file_path()`] to get the path to this file.

zebra-state/src/lib.rs

@@ -81,9 +81,6 @@ pub use service::{
     init_test, init_test_services,
 };
 
-#[cfg(any(test, feature = "proptest-impl"))]
-pub use constants::latest_version_for_adding_subtrees;
-
 #[cfg(any(test, feature = "proptest-impl"))]
 pub use config::hidden::{
     write_database_format_version_to_disk, write_state_database_format_version_to_disk,

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 in zatoshis, including change.
+        received: u64,
+    },
 
     /// 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

@@ -1689,20 +1689,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/transparent.rs

@@ -225,10 +225,26 @@ pub struct AddressBalanceLocation {
     /// The total balance of all UTXOs sent to an address.
     balance: Amount<NonNegative>,
 
+    /// The total balance of all spent and unspent outputs sent to an address.
+    received: u64,
+
     /// The location of the first [`transparent::Output`] sent to an address.
     location: AddressLocation,
 }
 
+impl std::ops::Add for AddressBalanceLocation {
+    type Output = Result<Self, amount::Error>;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        Ok(AddressBalanceLocation {
+            balance: (self.balance + rhs.balance)?,
+            // Addresses could receive more than the max money supply by sending to themselves, use u64::MAX if the addition overflows.
+            received: self.received.checked_add(rhs.received).unwrap_or(u64::MAX),
+            location: self.location,
+        })
+    }
+}
+
 impl AddressBalanceLocation {
     /// Creates a new [`AddressBalanceLocation`] from the location of
     /// the first [`transparent::Output`] sent to an address.
@@ -237,6 +253,7 @@ impl AddressBalanceLocation {
     pub fn new(first_output: OutputLocation) -> AddressBalanceLocation {
         AddressBalanceLocation {
             balance: Amount::zero(),
+            received: 0,
             location: first_output,
         }
     }
@@ -246,17 +263,32 @@ impl AddressBalanceLocation {
         self.balance
     }
 
+    /// Returns the current received balance for the address.
+    pub fn received(&self) -> u64 {
+        self.received
+    }
+
     /// Returns a mutable reference to the current balance for the address.
     pub fn balance_mut(&mut self) -> &mut Amount<NonNegative> {
         &mut self.balance
     }
 
+    /// Returns a mutable reference to the current received balance for the address.
+    pub fn received_mut(&mut self) -> &mut u64 {
+        &mut self.received
+    }
+
     /// Updates the current balance by adding the supplied output's value.
     pub fn receive_output(
         &mut self,
         unspent_output: &transparent::Output,
     ) -> Result<(), amount::Error> {
         self.balance = (self.balance + unspent_output.value())?;
+        self.received = self
+            .received
+            .checked_add(unspent_output.value().into())
+            // Addresses could receive more than the max money supply by sending to themselves.
+            .unwrap_or(u64::MAX);
 
         Ok(())
     }
@@ -645,13 +677,14 @@ impl FromDisk for OutputLocation {
 }
 
 impl IntoDisk for AddressBalanceLocation {
-    type Bytes = [u8; BALANCE_DISK_BYTES + OUTPUT_LOCATION_DISK_BYTES];
+    type Bytes = [u8; BALANCE_DISK_BYTES + OUTPUT_LOCATION_DISK_BYTES + size_of::<u64>()];
 
     fn as_bytes(&self) -> Self::Bytes {
         let balance_bytes = self.balance().as_bytes().to_vec();
         let address_location_bytes = self.address_location().as_bytes().to_vec();
+        let received_bytes = self.received().to_be_bytes().to_vec();
 
-        [balance_bytes, address_location_bytes]
+        [balance_bytes, address_location_bytes, received_bytes]
             .concat()
             .try_into()
             .unwrap()
@@ -660,14 +693,19 @@ impl IntoDisk for AddressBalanceLocation {
 
 impl FromDisk for AddressBalanceLocation {
     fn from_bytes(disk_bytes: impl AsRef<[u8]>) -> Self {
-        let (balance_bytes, address_location_bytes) =
-            disk_bytes.as_ref().split_at(BALANCE_DISK_BYTES);
+        let (balance_bytes, rest) = disk_bytes.as_ref().split_at(BALANCE_DISK_BYTES);
+        let (received_bytes, address_location_bytes) = rest.split_at(BALANCE_DISK_BYTES);
 
         let balance = Amount::from_bytes(balance_bytes.try_into().unwrap()).unwrap();
         let address_location = AddressLocation::from_bytes(address_location_bytes);
+        // # Backwards Compatibility
+        //
+        // If the value is missing a `received` field, default to 0.
+        let received = u64::from_be_bytes(received_bytes.try_into().unwrap_or_default());
 
         let mut address_balance_location = AddressBalanceLocation::new(address_location);
         *address_balance_location.balance_mut() = balance;
+        *address_balance_location.received_mut() = received;
 
         address_balance_location
     }