Merge branch 'aleks/clear-out-validator-sets-for-old-epochs' (#1665)

* origin/aleks/clear-out-validator-sets-for-old-epochs:
  refactor: fix formatting
  add changelog
  refactor: simplify code
  refactor: introduce name for magic constant
  add test case
  feat: store total consensus stake; garbage collect validator sets

.changelog/unreleased/improvements/1129-clear-out-validator-sets-for-old-epochs.md

@@ -0,0 +1,2 @@
+- PoS: purge validator sets for old epochs from the storage; store total
+  validator stake ([\#1129](https://github.com/anoma/namada/issues/1129))

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

@@ -105,8 +105,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

@@ -56,8 +56,9 @@ use storage::{
     validator_address_raw_hash_key, validator_last_slash_key,
     validator_max_commission_rate_change_key, BondDetails,
     BondsAndUnbondsDetail, BondsAndUnbondsDetails, EpochedSlashes,
-    ReverseOrdTokenAmount, RewardsAccumulator, SlashedAmount, UnbondDetails,
-    ValidatorAddresses, ValidatorUnbondRecords,
+    ReverseOrdTokenAmount, RewardsAccumulator, SlashedAmount,
+    TotalConsensusStakes, UnbondDetails, ValidatorAddresses,
+    ValidatorUnbondRecords,
 };
 use thiserror::Error;
 use types::{
@@ -84,6 +85,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 {
@@ -274,6 +279,12 @@ pub fn validator_eth_cold_key_handle(
     ValidatorEthColdKeys::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);
@@ -476,6 +487,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,
@@ -488,13 +502,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");
@@ -1528,14 +1536,15 @@ 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();
 
+    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),
         below_capacity_validator_set.at(&prev_epoch),
@@ -1597,28 +1606,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);
@@ -1627,6 +1639,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
+/// `STORE_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.
@@ -3288,9 +3362,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
         );
@@ -3299,6 +3373,7 @@ where
         let infracting_stake = slashes.iter(storage)?.fold(
             Ok(Dec::zero()),
             |acc: storage_api::Result<Dec>, res| {
+                let acc = acc?;
                 let (
                     NestedSubKey::Data {
                         key: validator,
@@ -3312,11 +3387,7 @@ where
                         .unwrap_or_default();
                 // println!("Val {} stake: {}", &validator, validator_stake);
 
-                if let Ok(inner) = acc {
-                    Ok(inner + Dec::from(validator_stake))
-                } else {
-                    acc
-                }
+                Ok(acc + Dec::from(validator_stake))
                 // TODO: does something more complex need to be done
                 // here in the event some of these slashes correspond to
                 // the same validator?
@@ -3892,22 +3963,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

@@ -36,6 +36,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";
@@ -584,6 +585,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())