feat: store total consensus stake; garbage collect validator sets

apps/src/lib/node/ledger/shell/finalize_block.rs

@@ -101,8 +101,14 @@ where
                 &mut self.wl_storage,
                 current_epoch,
                 current_epoch + pos_params.pipeline_len,
-                &namada_proof_of_stake::consensus_validator_set_handle(),
-                &namada_proof_of_stake::below_capacity_validator_set_handle(),
+            )?;
+            namada_proof_of_stake::store_total_consensus_stake(
+                &mut self.wl_storage,
+                current_epoch,
+            )?;
+            namada_proof_of_stake::purge_validator_sets_for_old_epoch(
+                &mut self.wl_storage,
+                current_epoch,
             )?;
         }
 

proof_of_stake/src/epoched.rs

@@ -659,6 +659,29 @@ where
     }
 }
 
+/// Zero offset
+#[derive(
+    Debug,
+    Clone,
+    BorshDeserialize,
+    BorshSerialize,
+    BorshSchema,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+)]
+pub struct OffsetZero;
+impl EpochOffset for OffsetZero {
+    fn value(_paras: &PosParams) -> u64 {
+        0
+    }
+
+    fn dyn_offset() -> DynEpochOffset {
+        DynEpochOffset::Zero
+    }
+}
+
 /// Offset at pipeline length.
 #[derive(
     Debug,
@@ -731,6 +754,8 @@ impl EpochOffset for OffsetPipelinePlusUnbondingLen {
 /// Offset length dynamic choice.
 #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
 pub enum DynEpochOffset {
+    /// Zero offset
+    Zero,
     /// Offset at pipeline length - 1
     PipelineLenMinusOne,
     /// Offset at pipeline length.

proof_of_stake/src/lib.rs

@@ -55,7 +55,8 @@ use storage::{
     params_key, slashes_prefix, unbonds_for_source_prefix, unbonds_prefix,
     validator_address_raw_hash_key, validator_max_commission_rate_change_key,
     BondDetails, BondsAndUnbondsDetail, BondsAndUnbondsDetails,
-    ReverseOrdTokenAmount, RewardsAccumulator, UnbondDetails,
+    ReverseOrdTokenAmount, RewardsAccumulator, TotalConsensusStakes,
+    UnbondDetails,
 };
 use thiserror::Error;
 use types::{
@@ -87,6 +88,10 @@ pub fn staking_token_address(storage: &impl StorageRead) -> Address {
         .expect("Must be able to read native token address")
 }
 
+/// Number of epochs below the current epoch for which full validator sets are
+/// stored
+const STORE_VALIDATOR_SETS_LEN: u64 = 2;
+
 #[allow(missing_docs)]
 #[derive(Error, Debug)]
 pub enum GenesisError {
@@ -261,6 +266,12 @@ pub fn validator_consensus_key_handle(
     ValidatorConsensusKeys::open(key)
 }
 
+/// Get the storage handle to the total consensus validator stake
+pub fn total_consensus_stake_key_handle() -> TotalConsensusStakes {
+    let key = storage::total_consensus_stake_key();
+    TotalConsensusStakes::open(key)
+}
+
 /// Get the storage handle to a PoS validator's state
 pub fn validator_state_handle(validator: &Address) -> ValidatorStates {
     let key = storage::validator_state_key(validator);
@@ -451,6 +462,9 @@ where
         )?;
     }
 
+    // Store the total consensus validator stake to storage
+    store_total_consensus_stake(storage, current_epoch)?;
+
     // Write total deltas to storage
     total_deltas_handle().init_at_genesis(
         storage,
@@ -463,13 +477,7 @@ where
     credit_tokens(storage, &staking_token, &ADDRESS, total_bonded)?;
     // Copy the genesis validator set into the pipeline epoch as well
     for epoch in (current_epoch.next()).iter_range(params.pipeline_len) {
-        copy_validator_sets_and_positions(
-            storage,
-            current_epoch,
-            epoch,
-            &consensus_validator_set_handle(),
-            &below_capacity_validator_set_handle(),
-        )?;
+        copy_validator_sets_and_positions(storage, current_epoch, epoch)?;
     }
 
     tracing::debug!("Genesis initialized");
@@ -1503,13 +1511,18 @@ pub fn copy_validator_sets_and_positions<S>(
     storage: &mut S,
     current_epoch: Epoch,
     target_epoch: Epoch,
-    consensus_validator_set: &ConsensusValidatorSets,
-    below_capacity_validator_set: &BelowCapacityValidatorSets,
 ) -> storage_api::Result<()>
 where
     S: StorageRead + StorageWrite,
 {
     let prev_epoch = target_epoch.prev();
+    debug_assert!(
+        prev_epoch < target_epoch,
+        "Target epoch must be in the future."
+    );
+
+    let consensus_validator_set = consensus_validator_set_handle();
+    let below_capacity_validator_set = below_capacity_validator_set_handle();
 
     let (consensus, below_capacity) = (
         consensus_validator_set.at(&prev_epoch),
@@ -1572,28 +1585,31 @@ where
 
     // Copy validator positions
     let mut positions = HashMap::<Address, Position>::default();
-    let positions_handle = validator_set_positions_handle().at(&prev_epoch);
+    let validator_set_positions_handle = validator_set_positions_handle();
+    let positions_handle = validator_set_positions_handle.at(&prev_epoch);
+
     for result in positions_handle.iter(storage)? {
         let (validator, position) = result?;
         positions.insert(validator, position);
     }
-    let new_positions_handle =
-        validator_set_positions_handle().at(&target_epoch);
+
+    let new_positions_handle = validator_set_positions_handle.at(&target_epoch);
     for (validator, position) in positions {
         let prev = new_positions_handle.insert(storage, validator, position)?;
         debug_assert!(prev.is_none());
     }
-    validator_set_positions_handle().set_last_update(storage, current_epoch)?;
+    validator_set_positions_handle.set_last_update(storage, current_epoch)?;
 
     // Copy set of all validator addresses
     let mut all_validators = HashSet::<Address>::default();
-    let all_validators_handle = validator_addresses_handle().at(&prev_epoch);
+    let validator_addresses_handle = validator_addresses_handle();
+    let all_validators_handle = validator_addresses_handle.at(&prev_epoch);
     for result in all_validators_handle.iter(storage)? {
         let validator = result?;
         all_validators.insert(validator);
     }
     let new_all_validators_handle =
-        validator_addresses_handle().at(&target_epoch);
+        validator_addresses_handle.at(&target_epoch);
     for validator in all_validators {
         let was_in = new_all_validators_handle.insert(storage, validator)?;
         debug_assert!(!was_in);
@@ -1602,6 +1618,68 @@ where
     Ok(())
 }
 
+/// Compute total validator stake for the current epoch
+fn compute_total_consensus_stake<S>(
+    storage: &S,
+    epoch: Epoch,
+) -> storage_api::Result<token::Amount>
+where
+    S: StorageRead,
+{
+    consensus_validator_set_handle()
+        .at(&epoch)
+        .iter(storage)?
+        .fold(Ok(token::Amount::zero()), |acc, entry| {
+            let acc = acc?;
+            let (
+                NestedSubKey::Data {
+                    key: amount,
+                    nested_sub_key: _,
+                },
+                _validator,
+            ) = entry?;
+            Ok(acc + amount)
+        })
+}
+
+/// Store total consensus stake
+pub fn store_total_consensus_stake<S>(
+    storage: &mut S,
+    epoch: Epoch,
+) -> storage_api::Result<()>
+where
+    S: StorageRead + StorageWrite,
+{
+    let total = compute_total_consensus_stake(storage, epoch)?;
+    tracing::debug!(
+        "Computed total consensus stake for epoch {}: {}",
+        epoch,
+        total.to_string_native()
+    );
+    total_consensus_stake_key_handle().set(storage, total, epoch, 0)
+}
+
+/// Purge the validator sets from the epochs older than the current epoch minus
+/// `keep_validator_sets_len`
+pub fn purge_validator_sets_for_old_epoch<S>(
+    storage: &mut S,
+    epoch: Epoch,
+) -> storage_api::Result<()>
+where
+    S: StorageRead + StorageWrite,
+{
+    if Epoch(STORE_VALIDATOR_SETS_LEN) < epoch {
+        let old_epoch = epoch - STORE_VALIDATOR_SETS_LEN - 1;
+        consensus_validator_set_handle()
+            .get_data_handler()
+            .remove_all(storage, &old_epoch)?;
+        below_capacity_validator_set_handle()
+            .get_data_handler()
+            .remove_all(storage, &old_epoch)?;
+    }
+    Ok(())
+}
+
 /// Read the position of the validator in the subset of validators that have the
 /// same bonded stake. This information is held in its own epoched structure in
 /// addition to being inside the validator sets.
@@ -3187,9 +3265,9 @@ where
 
     for epoch in Epoch::iter_bounds_inclusive(start_epoch, end_epoch) {
         let consensus_stake =
-            Dec::from(get_total_consensus_stake(storage, epoch)?);
+            Dec::from(get_total_consensus_stake(storage, epoch, params)?);
         tracing::debug!(
-            "Consensus stake in epoch {}: {}",
+            "Total consensus stake in epoch {}: {}",
             epoch,
             consensus_stake
         );
@@ -3791,22 +3869,14 @@ where
 fn get_total_consensus_stake<S>(
     storage: &S,
     epoch: Epoch,
+    params: &PosParams,
 ) -> storage_api::Result<token::Amount>
 where
     S: StorageRead,
 {
-    let mut total = token::Amount::default();
-    for res in consensus_validator_set_handle().at(&epoch).iter(storage)? {
-        let (
-            NestedSubKey::Data {
-                key: bonded_stake,
-                nested_sub_key: _,
-            },
-            _validator,
-        ) = res?;
-        total += bonded_stake;
-    }
-    Ok(total)
+    total_consensus_stake_key_handle()
+        .get(storage, epoch, params)
+        .map(|o| o.expect("Total consensus stake could not be retrieved."))
 }
 
 /// Find slashes applicable to a validator with inclusive `start` and exclusive

proof_of_stake/src/storage.rs

@@ -34,6 +34,7 @@ const VALIDATOR_TOTAL_UNBONDED_STORAGE_KEY: &str = "total_unbonded";
 const VALIDATOR_SETS_STORAGE_PREFIX: &str = "validator_sets";
 const CONSENSUS_VALIDATOR_SET_STORAGE_KEY: &str = "consensus";
 const BELOW_CAPACITY_VALIDATOR_SET_STORAGE_KEY: &str = "below_capacity";
+const TOTAL_CONSENSUS_STAKE_STORAGE_KEY: &str = "total_consensus_stake";
 const TOTAL_DELTAS_STORAGE_KEY: &str = "total_deltas";
 const VALIDATOR_SET_POSITIONS_KEY: &str = "validator_set_positions";
 const CONSENSUS_KEYS: &str = "consensus_keys";
@@ -532,6 +533,21 @@ pub fn is_below_capacity_validator_set_key(key: &Key) -> bool {
     matches!(&key.segments[..], [DbKeySeg::AddressSeg(addr), DbKeySeg::StringSeg(key), DbKeySeg::StringSeg(set_type), DbKeySeg::StringSeg(lazy_map), DbKeySeg::StringSeg(data), DbKeySeg::StringSeg(_epoch), DbKeySeg::StringSeg(_), DbKeySeg::StringSeg(_amount), DbKeySeg::StringSeg(_), DbKeySeg::StringSeg(_position)] if addr == &ADDRESS && key == VALIDATOR_SETS_STORAGE_PREFIX && set_type == BELOW_CAPACITY_VALIDATOR_SET_STORAGE_KEY && lazy_map == LAZY_MAP_SUB_KEY && data == lazy_map::DATA_SUBKEY)
 }
 
+/// Storage key for total consensus stake
+pub fn total_consensus_stake_key() -> Key {
+    Key::from(ADDRESS.to_db_key())
+        .push(&TOTAL_CONSENSUS_STAKE_STORAGE_KEY.to_owned())
+        .expect("Cannot obtain a total consensus stake key")
+}
+
+/// Is storage key for the total consensus stake?
+pub fn is_total_consensus_stake_key(key: &Key) -> bool {
+    matches!(&key.segments[..], [
+                DbKeySeg::AddressSeg(addr),
+                DbKeySeg::StringSeg(key)
+            ] if addr == &ADDRESS && key == TOTAL_CONSENSUS_STAKE_STORAGE_KEY)
+}
+
 /// Storage key for total deltas of all validators.
 pub fn total_deltas_key() -> Key {
     Key::from(ADDRESS.to_db_key())

proof_of_stake/src/tests.rs

@@ -42,15 +42,16 @@ use crate::{
     copy_validator_sets_and_positions, find_validator_by_raw_hash,
     get_num_consensus_validators, init_genesis,
     insert_validator_into_validator_set, is_validator, process_slashes,
+    purge_validator_sets_for_old_epoch,
     read_below_capacity_validator_set_addresses_with_stake,
     read_below_threshold_validator_set_addresses,
     read_consensus_validator_set_addresses_with_stake, read_total_stake,
     read_validator_delta_value, read_validator_stake, slash,
-    staking_token_address, total_deltas_handle, unbond_handle, unbond_tokens,
-    unjail_validator, update_validator_deltas, update_validator_set,
-    validator_consensus_key_handle, validator_set_update_tendermint,
-    validator_slashes_handle, validator_state_handle, withdraw_tokens,
-    write_validator_address_raw_hash,
+    staking_token_address, store_total_consensus_stake, total_deltas_handle,
+    unbond_handle, unbond_tokens, unjail_validator, update_validator_deltas,
+    update_validator_set, validator_consensus_key_handle,
+    validator_set_update_tendermint, validator_slashes_handle,
+    validator_state_handle, withdraw_tokens, write_validator_address_raw_hash,
 };
 
 proptest! {
@@ -1961,14 +1962,14 @@ fn get_tendermint_set_updates(
 fn advance_epoch(s: &mut TestWlStorage, params: &PosParams) -> Epoch {
     s.storage.block.epoch = s.storage.block.epoch.next();
     let current_epoch = s.storage.block.epoch;
+    store_total_consensus_stake(s, current_epoch).unwrap();
     copy_validator_sets_and_positions(
         s,
         current_epoch,
         current_epoch + params.pipeline_len,
-        &consensus_validator_set_handle(),
-        &below_capacity_validator_set_handle(),
     )
     .unwrap();
+    purge_validator_sets_for_old_epoch(s, current_epoch).unwrap();
     // process_slashes(s, current_epoch).unwrap();
     // dbg!(current_epoch);
     current_epoch

proof_of_stake/src/types.rs

@@ -109,6 +109,10 @@ pub type BelowCapacityValidatorSets = crate::epoched::NestedEpoched<
     crate::epoched::OffsetPipelineLen,
 >;
 
+/// Epoched total consensus validator stake
+pub type TotalConsensusStakes =
+    crate::epoched::Epoched<Amount, crate::epoched::OffsetZero, U64_MAX>;
+
 /// Epoched validator's deltas.
 pub type ValidatorDeltas = crate::epoched::EpochedDelta<
     token::Change,