State: Replay staking data during checkpoint copy

Author: endophysics

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

@@ -12,7 +12,7 @@
 //! each time the database format (column, serialization, etc) changes.
 
 use std::{
-    collections::{BTreeMap, HashMap},
+    collections::{BTreeMap, HashMap, HashSet},
     sync::Arc,
 };
 
@@ -268,25 +268,35 @@ impl DiskWriteBatch {
 
         // Apply bond rewards to the staking_bonded pool tally FIRST,
         // before unbonding moves value between pools (to match non-finalized state order)
-        let total_rewards: u64 = finalized.bond_rewards.iter().map(|(_, amount)| amount).sum();
+        let total_rewards: u64 = if finalized.bond_rewards.is_empty() {
+            checkpoint_replay_reward_total(db, finalized)
+        } else {
+            finalized
+                .bond_rewards
+                .iter()
+                .map(|(_, amount)| amount)
+                .sum()
+        };
         if total_rewards > 0 {
             let current_bonded = new_value_pool.staking_bonded_amount();
-            let new_bonded: Amount<NonNegative> = (current_bonded + Amount::try_from(total_rewards as i64)?)
-                .expect("staking_bonded pool should not overflow from rewards");
+            let new_bonded: Amount<NonNegative> = (current_bonded
+                + Amount::try_from(total_rewards as i64)?)
+            .expect("staking_bonded pool should not overflow from rewards");
             new_value_pool.set_staking_bonded_amount(new_bonded);
         }
 
         // Handle BeginDelegationUnbonding staking actions.
         // These move value from staking_bonded to staking_unbonded, but this transfer
         // is not captured by chain_value_pool_change (which only handles value entering/leaving pools).
-        // Use the pre-computed unbonding_amounts which include rewards from the non-finalized state.
-        for (_bond_key, bond_amount) in &finalized.unbonding_amounts {
-            let bond_amount: Amount<NonNegative> = Amount::try_from(*bond_amount as i64)?;
-
+        // Use the pre-computed unbonding_amounts from the non-finalized state when available.
+        // Checkpoint-style replay, used by copy-state and rollback-tip-height, bypasses the
+        // non-finalized state and therefore has to recover the amount from the finalized bond table.
+        let unbonding_amounts = finalized_unbonding_amounts(db, finalized)?;
+        for bond_amount in unbonding_amounts {
             // Move value from staking_bonded to staking_unbonded
             let current_bonded = new_value_pool.staking_bonded_amount();
-            let new_bonded: Amount<NonNegative> = (current_bonded - bond_amount)
-                .expect("staking_bonded pool should not underflow");
+            let new_bonded: Amount<NonNegative> =
+                (current_bonded - bond_amount).expect("staking_bonded pool should not underflow");
             new_value_pool.set_staking_bonded_amount(new_bonded);
 
             let current_unbonded = new_value_pool.staking_unbonded_amount();
@@ -314,3 +324,67 @@ impl DiskWriteBatch {
         Ok(())
     }
 }
+
+fn checkpoint_replay_reward_total(db: &ZebraDb, finalized: &FinalizedBlock) -> u64 {
+    let mut active_bonds: HashSet<_> = db
+        .all_bonds()
+        .filter_map(|(bond_key, _bond, status)| status.is_active().then_some(bond_key))
+        .collect();
+
+    for transaction in &finalized.block.transactions {
+        let Some(staking_action) = transaction.staking_action() else {
+            continue;
+        };
+
+        match staking_action.kind {
+            StakingActionKind::CreateNewDelegationBond => {
+                active_bonds.insert(staking_action.arg32_0);
+            }
+            StakingActionKind::BeginDelegationUnbonding
+            | StakingActionKind::WithdrawDelegationBond => {
+                active_bonds.remove(&staking_action.arg32_0);
+            }
+            _ => {}
+        }
+    }
+
+    if active_bonds.is_empty() {
+        0
+    } else {
+        500_000_000
+    }
+}
+
+fn finalized_unbonding_amounts(
+    db: &ZebraDb,
+    finalized: &FinalizedBlock,
+) -> Result<Vec<Amount<NonNegative>>, BoxError> {
+    if !finalized.unbonding_amounts.is_empty() {
+        return finalized
+            .unbonding_amounts
+            .iter()
+            .map(|(_bond_key, bond_amount)| {
+                let bond_amount = i64::try_from(*bond_amount)?;
+                Amount::try_from(bond_amount).map_err(Into::into)
+            })
+            .collect();
+    }
+
+    finalized
+        .block
+        .transactions
+        .iter()
+        .filter_map(|transaction| {
+            let staking_action = transaction.staking_action()?;
+            (staking_action.kind == StakingActionKind::BeginDelegationUnbonding)
+                .then_some(staking_action.arg32_0)
+        })
+        .map(|bond_key| {
+            db.delegation_bond(&bond_key)
+                .map(|bond| bond.amount)
+                .ok_or_else(|| {
+                    format!("bond {:?} not found while updating value pools", bond_key).into()
+                })
+        })
+        .collect()
+}

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

@@ -134,6 +134,74 @@ impl ZebraDb {
     }
 }
 
+fn checkpoint_replay_bond_rewards(
+    db: &ZebraDb,
+    bonds_modified_in_block: &HashMap<BondKey, DelegationBond>,
+    statuses_modified_in_block: &HashMap<BondKey, BondStatus>,
+) -> Vec<(BondKey, u64)> {
+    let mut active_bonds: Vec<(BondKey, DelegationBond)> = db
+        .all_bonds()
+        .filter_map(|(bond_key, bond, status)| {
+            let status = statuses_modified_in_block
+                .get(&bond_key)
+                .copied()
+                .unwrap_or(status);
+
+            status.is_active().then_some((
+                bond_key,
+                bonds_modified_in_block
+                    .get(&bond_key)
+                    .copied()
+                    .unwrap_or(bond),
+            ))
+        })
+        .collect();
+
+    for (bond_key, bond) in bonds_modified_in_block {
+        if db.delegation_bond(bond_key).is_none()
+            && statuses_modified_in_block
+                .get(bond_key)
+                .is_some_and(BondStatus::is_active)
+        {
+            active_bonds.push((*bond_key, *bond));
+        }
+    }
+
+    if active_bonds.is_empty() {
+        return Vec::new();
+    }
+
+    let total_staked_zats: u64 = active_bonds
+        .iter()
+        .map(|(_, bond)| u64::from(bond.amount))
+        .sum();
+    let Some(max_staker) = active_bonds
+        .iter()
+        .max_by_key(|(bond_key, bond)| (u64::from(bond.amount), std::cmp::Reverse(*bond_key)))
+        .map(|(bond_key, _)| *bond_key)
+    else {
+        return Vec::new();
+    };
+
+    let bond_reward_total = 500_000_000u64;
+    let mut paid_reward = 0u64;
+    let mut rewards = Vec::with_capacity(active_bonds.len());
+
+    for (bond_key, bond) in active_bonds {
+        if bond_key == max_staker {
+            continue;
+        }
+
+        let reward = ((u64::from(bond.amount) as u128) * (bond_reward_total as u128)
+            / (total_staked_zats as u128)) as u64;
+        paid_reward += reward;
+        rewards.push((bond_key, reward));
+    }
+
+    rewards.push((max_staker, bond_reward_total - paid_reward));
+    rewards
+}
+
 impl DiskWriteBatch {
     /// Prepare a database batch containing the delegation bonds in `finalized.block`,
     /// and return it (without actually writing anything).
@@ -150,9 +218,10 @@ impl DiskWriteBatch {
         // Process transactions to update bond state
         use zcash_primitives::transaction::StakingActionKind;
 
-        // Track bonds created in this block so we can apply rewards to them
-        // (they won't be in the DB yet when we apply rewards)
-        let mut bonds_created_in_block: HashMap<BondKey, DelegationBond> = HashMap::new();
+        // Track bonds modified in this block so we can apply rewards to the post-block bond state
+        // (they won't be in the DB yet when we apply rewards).
+        let mut bonds_modified_in_block: HashMap<BondKey, DelegationBond> = HashMap::new();
+        let mut statuses_modified_in_block: HashMap<BondKey, BondStatus> = HashMap::new();
 
         // Iterate through all transactions in the block
         for (transaction_index, transaction) in finalized.block.transactions.iter().enumerate() {
@@ -172,18 +241,41 @@ impl DiskWriteBatch {
                             DelegationBond::new(amount, target_finalizer, transaction_location);
 
                         // Track this bond for reward application
-                        bonds_created_in_block.insert(bond_key, bond.clone());
+                        bonds_modified_in_block.insert(bond_key, bond.clone());
+                        statuses_modified_in_block.insert(bond_key, BondStatus::Active);
 
                         // Insert new bond
                         self.prepare_new_delegation_bond(&db.db, bond_key, bond);
                     }
                     StakingActionKind::BeginDelegationUnbonding => {
                         // Mark bond as unbonding
-                        self.prepare_unbonding_delegation_bond(&db.db, db, bond_key, transaction_location)?;
+                        self.prepare_unbonding_delegation_bond(
+                            &db.db,
+                            db,
+                            bond_key,
+                            transaction_location,
+                        )?;
+                        statuses_modified_in_block.insert(
+                            bond_key,
+                            BondStatus::Unbonding {
+                                unbonded_at: transaction_location,
+                            },
+                        );
                     }
                     StakingActionKind::WithdrawDelegationBond => {
                         // Mark bond as withdrawn
-                        self.prepare_withdrawn_delegation_bond(&db.db, db, bond_key, transaction_location)?;
+                        self.prepare_withdrawn_delegation_bond(
+                            &db.db,
+                            db,
+                            bond_key,
+                            transaction_location,
+                        )?;
+                        statuses_modified_in_block.insert(
+                            bond_key,
+                            BondStatus::Withdrawn {
+                                withdrawn_at: transaction_location,
+                            },
+                        );
                     }
                     StakingActionKind::RetargetDelegationBond => {
                         // Update the bond's target_finalizer
@@ -193,7 +285,7 @@ impl DiskWriteBatch {
                             db,
                             bond_key,
                             new_target,
-                            &mut bonds_created_in_block,
+                            &mut bonds_modified_in_block,
                         )?;
                     }
                     // Other staking actions don't affect delegation bonds
@@ -202,12 +294,29 @@ impl DiskWriteBatch {
             }
         }
 
-        // Apply bond rewards accumulated in the non-finalized state
+        // Apply bond rewards accumulated in the non-finalized state.
+        // Checkpoint-style replay, used by copy-state and rollback-tip-height, bypasses the
+        // non-finalized state. In that case, derive the same deterministic rewards from the
+        // finalized bond table plus this block's pending bond/status changes.
+        let derived_bond_rewards;
+        let bond_rewards = if finalized.bond_rewards.is_empty() {
+            derived_bond_rewards = checkpoint_replay_bond_rewards(
+                db,
+                &bonds_modified_in_block,
+                &statuses_modified_in_block,
+            );
+            &derived_bond_rewards
+        } else {
+            &finalized.bond_rewards
+        };
+
         let delegation_bond_by_key_cf = db.db.cf_handle(DELEGATION_BOND_BY_KEY).unwrap();
-        for (bond_key, reward_amount) in &finalized.bond_rewards {
+        for (bond_key, reward_amount) in bond_rewards {
             // Get current bond from bonds modified in this block first, then fall back to DB.
             // This ensures we use the updated bond if it was retargeted/created in this block.
-            let bond_opt = bonds_created_in_block.get(bond_key).cloned()
+            let bond_opt = bonds_modified_in_block
+                .get(bond_key)
+                .cloned()
                 .or_else(|| db.delegation_bond(bond_key));
 
             if let Some(mut bond) = bond_opt {
@@ -226,7 +335,12 @@ impl DiskWriteBatch {
     /// Inserts into:
     /// - `delegation_bond_by_key`: stores the bond data
     /// - `bond_status_by_key`: stores Active status
-    fn prepare_new_delegation_bond(&mut self, db: &DiskDb, bond_key: BondKey, bond: DelegationBond) {
+    fn prepare_new_delegation_bond(
+        &mut self,
+        db: &DiskDb,
+        bond_key: BondKey,
+        bond: DelegationBond,
+    ) {
         let delegation_bond_by_key_cf = db.cf_handle(DELEGATION_BOND_BY_KEY).unwrap();
         let bond_status_by_key_cf = db.cf_handle(BOND_STATUS_BY_KEY).unwrap();
 
@@ -330,10 +444,13 @@ impl DiskWriteBatch {
 
         // Get current bond from bonds modified in this block first, then fall back to DB.
         // This ensures we use the updated bond if it was modified earlier in this block.
-        let bond_opt = bonds_created_in_block.get(&bond_key).cloned()
+        let bond_opt = bonds_created_in_block
+            .get(&bond_key)
+            .cloned()
             .or_else(|| zebra_db.delegation_bond(&bond_key));
 
-        let mut bond = bond_opt.ok_or_else(|| format!("bond {:?} not found for retarget", bond_key))?;
+        let mut bond =
+            bond_opt.ok_or_else(|| format!("bond {:?} not found for retarget", bond_key))?;
 
         // Update only target_finalizer (not created_at or amount)
         bond.target_finalizer = new_target;