fix(state): Restore initial sync performance by avoiding RocksDB merge operations when the on-disk database format is up-to-date (#9973)

* Adds a `finished_format_upgrades` field to `DiskDb` and sets it to true if no db format upgrade is required or after required db format upgrades have been applied

* Updates `ZebraDb::write_block()` to insert new address balance locations instead of merging the changes once any db format upgrades have finished.

* Apply suggestions from code review

Co-authored-by: Marek <mail@marek.onl>

* Removes `#[inline]` attribute from private generic functions that should be inlined by the compiler anyway

Adds a comment explaining why we prefer insertions when updating address balance locations over merge operations when there's no risk that insertions could overwrite concurrently updated entries.

---------

Co-authored-by: Marek <mail@marek.onl>

Author: arya2

zebra-state/src/service/finalized_state/disk_db.rs

@@ -16,7 +16,10 @@ use std::{
     fs,
     ops::RangeBounds,
     path::Path,
-    sync::Arc,
+    sync::{
+        atomic::{self, AtomicBool},
+        Arc,
+    },
 };
 
 use itertools::Itertools;
@@ -90,6 +93,10 @@ pub struct DiskDb {
     /// If true, the database files are deleted on drop.
     ephemeral: bool,
 
+    /// A boolean flag indicating whether the db format change task has finished
+    /// applying any format changes that may have been required.
+    finished_format_upgrades: Arc<AtomicBool>,
+
     // Owned State
     //
     // Everything contained in this state must be shared by all clones, or read-only.
@@ -623,6 +630,19 @@ impl DiskDb {
         total_size_on_disk
     }
 
+    /// Sets `finished_format_upgrades` to true to indicate that Zebra has
+    /// finished applying any required db format upgrades.
+    pub fn mark_finished_format_upgrades(&self) {
+        self.finished_format_upgrades
+            .store(true, atomic::Ordering::SeqCst);
+    }
+
+    /// Returns true if the `finished_format_upgrades` flag has been set to true to
+    /// indicate that Zebra has finished applying any required db format upgrades.
+    pub fn finished_format_upgrades(&self) -> bool {
+        self.finished_format_upgrades.load(atomic::Ordering::SeqCst)
+    }
+
     /// When called with a secondary DB instance, tries to catch up with the primary DB instance
     pub fn try_catch_up_with_primary(&self) -> Result<(), rocksdb::Error> {
         self.db.try_catch_up_with_primary()
@@ -932,6 +952,7 @@ impl DiskDb {
                     network: network.clone(),
                     ephemeral: config.ephemeral,
                     db: Arc::new(db),
+                    finished_format_upgrades: Arc::new(AtomicBool::new(false)),
                 };
 
                 db.assert_default_cf_is_empty();

zebra-state/src/service/finalized_state/disk_format/transparent.rs

@@ -5,7 +5,7 @@
 //! [`crate::constants::state_database_format_version_in_code()`] must be incremented
 //! each time the database format (column, serialization, etc) changes.
 
-use std::{cmp::max, fmt::Debug};
+use std::{cmp::max, collections::HashMap, fmt::Debug};
 
 use zebra_chain::{
     amount::{self, Amount, Constraint, NegativeAllowed, NonNegative},
@@ -173,7 +173,7 @@ impl<C: Constraint + Copy + std::fmt::Debug> AddressBalanceLocationInner<C> {
     /// the first [`transparent::Output`] sent to an address.
     ///
     /// The returned value has a zero initial balance and received balance.
-    fn new(first_output: OutputLocation) -> Self {
+    pub(crate) fn new(first_output: OutputLocation) -> Self {
         Self {
             balance: Amount::zero(),
             received: 0,
@@ -212,6 +212,38 @@ impl<C: Constraint + Copy + std::fmt::Debug> AddressBalanceLocationInner<C> {
     pub fn height_mut(&mut self) -> &mut Height {
         &mut self.location.transaction_location.height
     }
+
+    /// Updates the current balance by adding the supplied output's value.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn receive_output(
+        &mut self,
+        unspent_output: &transparent::Output,
+    ) -> Result<(), amount::Error> {
+        self.balance = (self
+            .balance
+            .zatoshis()
+            .checked_add(unspent_output.value().zatoshis()))
+        .expect("ops handling taddr balances must not overflow")
+        .try_into()?;
+        self.received = self.received.saturating_add(unspent_output.value().into());
+        Ok(())
+    }
+
+    /// Updates the current balance by subtracting the supplied output's value.
+    #[allow(clippy::unwrap_in_result)]
+    pub fn spend_output(
+        &mut self,
+        spent_output: &transparent::Output,
+    ) -> Result<(), amount::Error> {
+        self.balance = (self
+            .balance
+            .zatoshis()
+            .checked_sub(spent_output.value().zatoshis()))
+        .expect("ops handling taddr balances must not underflow")
+        .try_into()?;
+
+        Ok(())
+    }
 }
 
 impl<C: Constraint + Copy + std::fmt::Debug> std::ops::Add for AddressBalanceLocationInner<C> {
@@ -238,49 +270,51 @@ impl<C: Constraint + Copy + std::fmt::Debug> std::ops::Add for AddressBalanceLoc
     }
 }
 
+impl From<AddressBalanceLocationInner<NonNegative>> for AddressBalanceLocation {
+    fn from(value: AddressBalanceLocationInner<NonNegative>) -> Self {
+        Self(value)
+    }
+}
+
+impl From<AddressBalanceLocationInner<NegativeAllowed>> for AddressBalanceLocationChange {
+    fn from(value: AddressBalanceLocationInner<NegativeAllowed>) -> Self {
+        Self(value)
+    }
+}
+
 /// Represents a change in the [`AddressBalanceLocation`] of a transparent address
 /// in the finalized state.
 pub struct AddressBalanceLocationChange(AddressBalanceLocationInner<NegativeAllowed>);
 
+/// Represents a set of updates to address balance locations in the database.
+pub enum AddressBalanceLocationUpdates {
+    /// A set of [`AddressBalanceLocationChange`]s that should be merged into the existing values in the database.
+    Merge(HashMap<transparent::Address, AddressBalanceLocationChange>),
+    /// A set of full [`AddressBalanceLocation`]s that should be inserted as the new values in the database.
+    Insert(HashMap<transparent::Address, AddressBalanceLocation>),
+}
+
+impl From<HashMap<transparent::Address, AddressBalanceLocation>> for AddressBalanceLocationUpdates {
+    fn from(value: HashMap<transparent::Address, AddressBalanceLocation>) -> Self {
+        Self::Insert(value)
+    }
+}
+
+impl From<HashMap<transparent::Address, AddressBalanceLocationChange>>
+    for AddressBalanceLocationUpdates
+{
+    fn from(value: HashMap<transparent::Address, AddressBalanceLocationChange>) -> Self {
+        Self::Merge(value)
+    }
+}
+
 impl AddressBalanceLocationChange {
     /// Creates a new [`AddressBalanceLocationChange`].
     ///
     /// See [`AddressBalanceLocationInner::new`] for more details.
     pub fn new(location: AddressLocation) -> Self {
         Self(AddressBalanceLocationInner::new(location))
     }
-
-    /// Updates the current balance by adding the supplied output's value.
-    #[allow(clippy::unwrap_in_result)]
-    pub fn receive_output(
-        &mut self,
-        unspent_output: &transparent::Output,
-    ) -> Result<(), amount::Error> {
-        self.balance = (self
-            .balance
-            .zatoshis()
-            .checked_add(unspent_output.value().zatoshis()))
-        .expect("adding two Amounts is always within an i64")
-        .try_into()?;
-        self.received = self.received.saturating_add(unspent_output.value().into());
-        Ok(())
-    }
-
-    /// Updates the current balance by subtracting the supplied output's value.
-    #[allow(clippy::unwrap_in_result)]
-    pub fn spend_output(
-        &mut self,
-        spent_output: &transparent::Output,
-    ) -> Result<(), amount::Error> {
-        self.balance = (self
-            .balance
-            .zatoshis()
-            .checked_sub(spent_output.value().zatoshis()))
-        .expect("subtracting two Amounts is always within an i64")
-        .try_into()?;
-
-        Ok(())
-    }
 }
 
 impl std::ops::Deref for AddressBalanceLocationChange {

zebra-state/src/service/finalized_state/disk_format/upgrade.rs

@@ -381,9 +381,18 @@ impl DbFormatChange {
         initial_tip_height: Option<Height>,
         cancel_receiver: &Receiver<CancelFormatChange>,
     ) -> Result<(), CancelFormatChange> {
+        // Mark the database as having finished applying any format upgrades if there are no
+        // format upgrades that need to be applied.
+        if !self.is_upgrade() {
+            db.mark_finished_format_upgrades();
+        }
+
         match self {
             // Perform any required upgrades, then mark the state as upgraded.
-            Upgrade { .. } => self.apply_format_upgrade(db, initial_tip_height, cancel_receiver)?,
+            Upgrade { .. } => {
+                self.apply_format_upgrade(db, initial_tip_height, cancel_receiver)?;
+                db.mark_finished_format_upgrades();
+            }
 
             NewlyCreated { .. } => {
                 Self::mark_as_newly_created(db);

zebra-state/src/service/finalized_state/zebra_db.rs

@@ -179,6 +179,18 @@ impl ZebraDb {
         self.format_change_handle = Some(format_change_handle);
     }
 
+    /// Sets `finished_format_upgrades` to true on the inner [`DiskDb`] to indicate that Zebra has
+    /// finished applying any required db format upgrades.
+    pub fn mark_finished_format_upgrades(&self) {
+        self.db.mark_finished_format_upgrades();
+    }
+
+    /// Returns true if the `finished_format_upgrades` flag has been set to true on the inner [`DiskDb`] to
+    /// indicate that Zebra has finished applying any required db format upgrades.
+    pub fn finished_format_upgrades(&self) -> bool {
+        self.db.finished_format_upgrades()
+    }
+
     /// Returns config for this database.
     pub fn config(&self) -> &Config {
         &self.config

zebra-state/src/service/finalized_state/zebra_db/block.rs

@@ -37,7 +37,7 @@ use crate::{
         disk_db::{DiskDb, DiskWriteBatch, ReadDisk, WriteDisk},
         disk_format::{
             block::TransactionLocation,
-            transparent::{AddressBalanceLocationChange, OutputLocation},
+            transparent::{AddressBalanceLocationUpdates, OutputLocation},
         },
         zebra_db::{metrics::block_precommit_metrics, ZebraDb},
         FromDisk, RawBytes,
@@ -517,18 +517,36 @@ impl ZebraDb {
             .collect();
 
         // Get the current address balances, before the transactions in this block
-        let address_balances: HashMap<transparent::Address, AddressBalanceLocationChange> =
+
+        fn read_addr_locs<T, F: Fn(&transparent::Address) -> Option<T>>(
+            changed_addresses: HashSet<transparent::Address>,
+            f: F,
+        ) -> HashMap<transparent::Address, T> {
             changed_addresses
                 .into_iter()
-                .filter_map(|address| {
-                    // # Correctness
-                    //
-                    // Address balances are updated with the `fetch_add_balance_and_received` merge operator, so
-                    // the values must represent the changes to the balance, not the final balance.
-                    let addr_loc = self.address_balance_location(&address)?.into_new_change();
-                    Some((address.clone(), addr_loc))
-                })
-                .collect();
+                .filter_map(|address| Some((address.clone(), f(&address)?)))
+                .collect()
+        }
+
+        // # Performance
+        //
+        // It's better to update entries in RocksDB with insertions over merge operations when there is no risk that
+        // insertions may overwrite values that are updated concurrently in database format upgrades as inserted values
+        // are quicker to read and require less background compaction.
+        //
+        // Reading entries that have been updated with merge ops often requires reading the latest fully-merged value,
+        // reading all of the pending merge operands (potentially hundreds), and applying pending merge operands to the
+        // fully-merged value such that it's much faster to read entries that have been updated with insertions than it
+        // is to read entries that have been updated with merge operations.
+        let address_balances: AddressBalanceLocationUpdates = if self.finished_format_upgrades() {
+            AddressBalanceLocationUpdates::Insert(read_addr_locs(changed_addresses, |addr| {
+                self.address_balance_location(addr)
+            }))
+        } else {
+            AddressBalanceLocationUpdates::Merge(read_addr_locs(changed_addresses, |addr| {
+                Some(self.address_balance_location(addr)?.into_new_change())
+            }))
+        };
 
         let mut batch = DiskWriteBatch::new();
 
@@ -602,7 +620,7 @@ impl DiskWriteBatch {
             transparent::OutPoint,
             OutputLocation,
         >,
-        address_balances: HashMap<transparent::Address, AddressBalanceLocationChange>,
+        address_balances: AddressBalanceLocationUpdates,
         value_pool: ValueBalance<NonNegative>,
         prev_note_commitment_trees: Option<NoteCommitmentTrees>,
     ) -> Result<(), BoxError> {