staking: propagate api changes throughout

crates/core/component/stake/src/component/stake.rs

@@ -2,8 +2,8 @@ use crate::params::StakeParameters;
 use crate::rate::BaseRateData;
 use crate::validator::{self, Validator};
 use crate::{
-    state_key, CurrentConsensusKeys, Delegate, DelegationChanges, DelegationToken, FundingStreams,
-    IdentityKey, Penalty, Undelegate,
+    state_key, CurrentConsensusKeys, Delegate, DelegationChanges, FundingStreams, IdentityKey,
+    Penalty, Undelegate,
 };
 use anyhow::Context;
 use anyhow::{anyhow, Result};
@@ -14,7 +14,6 @@ use futures::{StreamExt, TryStreamExt};
 use penumbra_num::Amount;
 use penumbra_proto::{StateReadProto, StateWriteProto};
 use penumbra_sct::component::clock::EpochRead;
-use penumbra_shielded_pool::component::SupplyRead;
 use sha2::{Digest, Sha256};
 use std::pin::Pin;
 use std::str::FromStr;
@@ -418,9 +417,14 @@ pub(crate) trait InternalStakingData: StateRead {
             }
 
             let delegation_token_supply = self
-                .token_supply(&DelegationToken::from(validator_identity).id())
-                .await?
-                .expect("delegation token should be known");
+                .get_validator_pool_size(&validator_identity)
+                .await
+                .ok_or_else(|| {
+                anyhow::anyhow!(
+                    "validator delegation pool not found for {}",
+                    validator_identity
+                )
+            })?;
 
             let validator_rate = self
                 .get_validator_rate(&validator_identity)

crates/core/component/stake/src/component/validator_handler/validator_manager.rs

@@ -14,7 +14,7 @@ use async_trait::async_trait;
 use futures::StreamExt as _;
 use penumbra_num::Amount;
 use penumbra_sct::component::clock::{EpochManager, EpochRead};
-use penumbra_shielded_pool::component::{SupplyRead as _, SupplyWrite};
+use penumbra_shielded_pool::component::SupplyWrite;
 use sha2::{Digest as _, Sha256};
 use tendermint::abci::types::{CommitInfo, Misbehavior};
 use tokio::task::JoinSet;
@@ -25,6 +25,7 @@ use cnidarium::StateWrite;
 use penumbra_proto::StateWriteProto;
 use tracing::{instrument, Instrument};
 
+use crate::component::validator_handler::validator_store::ValidatorPoolTracker;
 use crate::{
     component::validator_handler::ValidatorDataRead,
     component::StateReadExt as _,
@@ -393,6 +394,7 @@ pub trait ValidatorManager: StateWrite {
             // All genesis validators start in the "Bonded" bonding state:
             validator::BondingState::Bonded,
             power,
+            total_delegation_tokens,
         )
         .await?;
 
@@ -421,7 +423,8 @@ pub trait ValidatorManager: StateWrite {
             rate_data,
             validator::State::Defined,
             validator::BondingState::Unbonded,
-            0u128.into(),
+            Amount::zero(),
+            Amount::zero(),
         )
         .await
     }
@@ -442,6 +445,7 @@ pub trait ValidatorManager: StateWrite {
         initial_state: validator::State,
         initial_bonding_state: validator::BondingState,
         initial_voting_power: Amount,
+        initial_delegation_pool_size: Amount,
     ) -> Result<()> {
         tracing::debug!("adding validator");
         if !matches!(initial_state, State::Defined | State::Active) {
@@ -473,6 +477,7 @@ pub trait ValidatorManager: StateWrite {
         self.set_initial_validator_state(&validator_identity, initial_state)?;
         self.set_validator_power(&validator_identity, initial_voting_power)?;
         self.set_validator_bonding_state(&validator_identity, initial_bonding_state);
+        self.set_validator_pool_size(&validator_identity, initial_delegation_pool_size);
 
         // Finally, update metrics for the new validator.
         match initial_state {
@@ -517,9 +522,9 @@ pub trait ValidatorManager: StateWrite {
                     .await?
                     .ok_or_else(|| anyhow::anyhow!("updated validator not found in JMT"))?;
                 let delegation_token_supply = self
-                    .token_supply(&DelegationToken::from(id).id())
-                    .await?
-                    .ok_or_else(|| anyhow::anyhow!("updated validator not found in JMT"))?;
+                    .get_validator_pool_size(id)
+                    .await
+                    .unwrap_or_else(Amount::zero);
                 let unbonded_amount =
                     current_validator_rate.unbonded_amount(delegation_token_supply);
 
@@ -538,9 +543,9 @@ pub trait ValidatorManager: StateWrite {
                     .await?
                     .ok_or_else(|| anyhow::anyhow!("updated validator not found in JMT"))?;
                 let delegation_pool_size = self
-                    .token_supply(&DelegationToken::from(id).id())
-                    .await?
-                    .ok_or_else(|| anyhow::anyhow!("updated validator not found in JMT"))?;
+                    .get_validator_pool_size(id)
+                    .await
+                    .unwrap_or_else(Amount::zero);
 
                 let unbonded_pool_size = validator_rate_data.unbonded_amount(delegation_pool_size);