staking: 🎁 add additional tracing telemetry (#4017)

this is a group of improvements to the telemetry in the staking
component, cherry-picked out of #4001. this improves the logs seen in
tests, and additionally takes the time to make some other small
gardening tweaks: consolidating imports, adding documentation comments,
etc. this is another branch in the spirit of #4009.

the changes here are itemized into distinct, small code transformations
to facilitate review, and demonstrate that these are noop changes. you
can start
[here](f7f78b9).

* f331964 - feat(staking): ❗ traces for invalid or missing parameters 
* d33598d - docs(staking): 🤝 add a doc comment 
* f6c60a8 - staking: 🙂 remove needless `Ok(_)` 
* fd4b8da - feat(staking): 💬 log when delegation changes are not found
* fe597ff - feat(staking): 🔍 trace-level spans for `StateReadExt`
methods
* 0447600 - refactor(staking): 🚡 use `self::` prefix in reexports 
* 67893c9 - refactor(staking): 🙂 remove import from reexports 
* 043222e - refactor(staking): 🎁 consolidate public submodules 
* 6376796 - refactor(staking): 🎁 consolidate submodules 
* b5bd240 - refactor(staking): 👯 component flag reexports
StateWriteExt
* d57beaa - chore(staking): 🥡 consolidate component reexports 
* aa52f6d - chore(staking): 🧹 tidy epoch_handler imports 
* 765ed75 - refactor(compact-block): 🧹 tidy imports 
* 7cb93f7 - feat(staking): 🌐 log why epoch ended 
* 081d041 - refactor(staking): add context to `state.end_epoch` call 
* 5701295 - feat(app): 🦠 tracing telemetry for `App::end_block` 
* f7f78b9 - docs(staking): 📕 add doc comment for `end_epoch`

crates/core/app/src/app/mod.rs

@@ -372,9 +372,11 @@ impl App {
         Ok(state_tx.apply().1)
     }
 
+    #[tracing::instrument(skip_all, fields(height = %end_block.height))]
     pub async fn end_block(&mut self, end_block: &request::EndBlock) -> Vec<abci::Event> {
         let state_tx = StateDelta::new(self.state.clone());
 
+        tracing::debug!("running app components' `end_block` hooks");
         let mut arc_state_tx = Arc::new(state_tx);
         ShieldedPool::end_block(&mut arc_state_tx, end_block).await;
         Distributions::end_block(&mut arc_state_tx, end_block).await;
@@ -387,6 +389,7 @@ impl App {
         Funding::end_block(&mut arc_state_tx, end_block).await;
         let mut state_tx = Arc::try_unwrap(arc_state_tx)
             .expect("components did not retain copies of shared state");
+        tracing::debug!("finished app components' `end_block` hooks");
 
         // Validate governance proposals here. We must do this here because proposals can affect
         // the entirety of application state, and the governance component does not have access to
@@ -474,7 +477,7 @@ impl App {
             .expect("able to detect upgrade heights");
 
         if is_end_epoch || is_chain_upgrade {
-            tracing::info!(?current_height, "ending epoch");
+            tracing::info!(%is_end_epoch, %is_chain_upgrade, ?current_height, "ending epoch");
 
             let mut arc_state_tx = Arc::new(state_tx);
 

crates/core/component/compact-block/src/component/manager.rs

@@ -1,15 +1,12 @@
-use anyhow::Context;
-use anyhow::Result;
+use anyhow::{Context, Result};
 use async_trait::async_trait;
 use cnidarium::StateWrite;
-use penumbra_dex::component::StateReadExt;
-use penumbra_dex::component::SwapManager as _;
+use penumbra_dex::component::{StateReadExt, SwapManager as _};
 use penumbra_fee::component::StateReadExt as _;
 use penumbra_governance::StateReadExt as _;
 use penumbra_proto::DomainType;
 use penumbra_sct::component::clock::EpochRead;
-use penumbra_sct::component::tree::SctManager as _;
-use penumbra_sct::component::tree::SctRead;
+use penumbra_sct::component::tree::{SctManager as _, SctRead};
 use penumbra_shielded_pool::component::NoteManager as _;
 use tracing::instrument;
 

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

@@ -6,10 +6,8 @@ use penumbra_num::Amount;
 use penumbra_sct::component::clock::EpochRead;
 
 use crate::{
-    component::{validator_handler::ValidatorDataRead, StateWriteExt as _},
-    event,
-    validator::State::*,
-    Delegate, StateReadExt as _,
+    component::validator_handler::ValidatorDataRead, event, validator::State::*, Delegate,
+    StateReadExt as _, StateWriteExt as _,
 };
 
 #[async_trait]

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

@@ -1,38 +1,31 @@
-use penumbra_distributions::component::StateReadExt as _;
-use penumbra_sct::{component::clock::EpochRead, epoch::Epoch};
-use std::collections::{BTreeMap, BTreeSet};
-
+use crate::{
+    component::{
+        stake::{
+            ConsensusIndexRead, ConsensusIndexWrite, ConsensusUpdateWrite, InternalStakingData,
+            RateDataWrite,
+        },
+        validator_handler::{ValidatorDataRead, ValidatorDataWrite, ValidatorManager},
+        SlashingData,
+    },
+    rate::BaseRateData,
+    state_key, validator, CurrentConsensusKeys, DelegationToken, FundingStreams, IdentityKey,
+    Penalty, StateReadExt, StateWriteExt, BPS_SQUARED_SCALING_FACTOR,
+};
 use anyhow::{Context, Result};
 use async_trait::async_trait;
-use futures::StreamExt;
-use penumbra_asset::STAKING_TOKEN_ASSET_ID;
-
 use cnidarium::StateWrite;
-use futures::TryStreamExt;
+use futures::{StreamExt, TryStreamExt};
+use penumbra_asset::STAKING_TOKEN_ASSET_ID;
+use penumbra_distributions::component::StateReadExt as _;
 use penumbra_num::{fixpoint::U128x128, Amount};
 use penumbra_proto::{StateReadProto, StateWriteProto};
+use penumbra_sct::{component::clock::EpochRead, epoch::Epoch};
 use penumbra_shielded_pool::component::{SupplyRead, SupplyWrite};
-use tendermint::validator::Update;
-use tendermint::PublicKey;
+use std::collections::{BTreeMap, BTreeSet};
+use tendermint::{validator::Update, PublicKey};
 use tokio::task::JoinSet;
 use tracing::instrument;
 
-use crate::state_key;
-use crate::BPS_SQUARED_SCALING_FACTOR;
-use crate::{
-    component::{
-        stake::{ConsensusUpdateWrite, InternalStakingData, RateDataWrite},
-        validator_handler::{ValidatorDataRead, ValidatorDataWrite, ValidatorManager},
-        SlashingData,
-    },
-    rate::BaseRateData,
-    validator, CurrentConsensusKeys, DelegationToken, FundingStreams, IdentityKey, Penalty,
-    StateReadExt,
-};
-
-use super::StateWriteExt;
-use crate::component::stake::{ConsensusIndexRead, ConsensusIndexWrite};
-
 #[async_trait]
 pub trait EpochHandler: StateWriteExt + ConsensusIndexRead {
     #[instrument(skip(self, epoch_to_end), fields(index = epoch_to_end.index))]

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

@@ -5,11 +5,11 @@ pub mod rpc;
 pub mod stake;
 pub mod validator_handler;
 
-pub use self::metrics::register_metrics;
-pub use stake::Staking;
+pub use self::{
+    metrics::register_metrics,
+    stake::{ConsensusIndexRead, SlashingData, Staking, StateReadExt, StateWriteExt},
+};
 
 // Max validator power is 1152921504606846975 (i64::MAX / 8)
 // https://github.com/tendermint/tendermint/blob/master/types/validator_set.go#L25
 pub const MAX_VOTING_POWER: u128 = 1152921504606846975;
-pub use stake::{ConsensusIndexRead, SlashingData};
-pub use stake::{StateReadExt, StateWriteExt};

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

@@ -19,11 +19,11 @@ use sha2::{Digest, Sha256};
 use std::pin::Pin;
 use std::str::FromStr;
 use std::{collections::BTreeMap, sync::Arc};
-use tap::Tap;
+use tap::{Tap, TapFallible, TapOptional};
 use tendermint::v0_37::abci;
 use tendermint::validator::Update;
 use tendermint::{block, PublicKey};
-use tracing::{instrument, trace};
+use tracing::{error, instrument, trace};
 
 use crate::component::epoch_handler::EpochHandler;
 use crate::component::validator_handler::{ValidatorDataRead, ValidatorManager};
@@ -149,14 +149,18 @@ impl Component for Staking {
         state.set_delegation_changes(height, changes).await;
     }
 
+    /// Writes validator updates for this block.
     #[instrument(name = "staking", skip(state))]
     async fn end_epoch<S: StateWrite + 'static>(state: &mut Arc<S>) -> anyhow::Result<()> {
         let state = Arc::get_mut(state).context("state should be unique")?;
         let epoch_ending = state
             .get_current_epoch()
             .await
             .context("should be able to get current epoch during end_epoch")?;
-        state.end_epoch(epoch_ending).await?;
+        state
+            .end_epoch(epoch_ending)
+            .await
+            .context("should be able to write end_epoch")?;
         // Since we only update the validator set at epoch boundaries,
         // we only need to build the validator set updates here in end_epoch.
         state
@@ -195,57 +199,70 @@ impl<T: StateWrite + ?Sized> ConsensusUpdateWrite for T {}
 #[async_trait]
 pub trait StateReadExt: StateRead {
     /// Gets the stake parameters from the JMT.
+    #[instrument(skip(self), level = "trace")]
     async fn get_stake_params(&self) -> Result<StakeParameters> {
         self.get(state_key::parameters::key())
             .await
+            .tap_err(|err| error!(?err, "could not deserialize stake parameters"))
             .expect("no deserialization error should happen")
+            .tap_none(|| error!("could not find stake parameters"))
             .ok_or_else(|| anyhow!("Missing StakeParameters"))
     }
 
     /// Indicates if the stake parameters have been updated in this block.
+    #[instrument(skip(self), level = "trace")]
     fn stake_params_updated(&self) -> bool {
         self.object_get::<()>(state_key::parameters::updated_flag())
             .is_some()
     }
 
+    #[instrument(skip(self), level = "trace")]
     async fn signed_blocks_window_len(&self) -> Result<u64> {
         Ok(self.get_stake_params().await?.signed_blocks_window_len)
     }
 
+    #[instrument(skip(self), level = "trace")]
     async fn missed_blocks_maximum(&self) -> Result<u64> {
         Ok(self.get_stake_params().await?.missed_blocks_maximum)
     }
 
     /// Delegation changes accumulated over the course of this block, to be
     /// persisted at the end of the block for processing at the end of the next
     /// epoch.
+    #[instrument(skip(self), level = "trace")]
     fn get_delegation_changes_tally(&self) -> DelegationChanges {
         self.object_get(state_key::chain::delegation_changes::key())
             .unwrap_or_default()
     }
 
+    #[instrument(skip(self), level = "trace")]
     async fn get_current_base_rate(&self) -> Result<BaseRateData> {
         self.get(state_key::chain::base_rate::current())
             .await
             .map(|rate_data| rate_data.expect("rate data must be set after init_chain"))
     }
 
+    #[instrument(skip(self), level = "trace")]
     fn get_previous_base_rate(&self) -> Option<BaseRateData> {
         self.object_get(state_key::chain::base_rate::previous())
     }
 
     /// Returns the funding queue from object storage (end-epoch).
+    #[instrument(skip(self), level = "trace")]
     fn get_funding_queue(&self) -> Option<Vec<(IdentityKey, FundingStreams, Amount)>> {
         self.object_get(state_key::validators::rewards::staking())
     }
 
+    /// Returns the [`DelegationChanges`] at the given [`Height`][block::Height].
+    #[instrument(skip(self), level = "trace")]
     async fn get_delegation_changes(&self, height: block::Height) -> Result<DelegationChanges> {
-        Ok(self
-            .get(&state_key::chain::delegation_changes::by_height(
-                height.value(),
-            ))
-            .await?
-            .ok_or_else(|| anyhow!("missing delegation changes for block {}", height))?)
+        self.get(&state_key::chain::delegation_changes::by_height(
+            height.value(),
+        ))
+        .await
+        .tap_err(|err| error!(?err, "delegation changes for block exist but are invalid"))?
+        .tap_none(|| error!("could not find delegation changes for block"))
+        .ok_or_else(|| anyhow!("missing delegation changes for block {}", height))
     }
 }
 

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

@@ -5,49 +5,46 @@
 
 mod changes;
 mod current_consensus_keys;
+mod delegation_token;
 mod event;
+mod governance_key;
+mod identity_key;
+mod penalty;
+mod unbonding_token;
 mod uptime;
 
-#[cfg(feature = "component")]
-pub mod component;
-
-#[cfg(feature = "component")]
-pub use component::StateReadExt;
-
-pub static BPS_SQUARED_SCALING_FACTOR: Lazy<penumbra_num::fixpoint::U128x128> =
-    Lazy::new(|| 1_0000_0000u128.into());
-
 pub mod delegate;
 pub mod funding_stream;
+pub mod genesis;
+pub mod params;
 pub mod rate;
 pub mod state_key;
 pub mod undelegate;
 pub mod undelegate_claim;
 pub mod validator;
 
-pub use delegate::Delegate;
-use once_cell::sync::Lazy;
-pub use undelegate::Undelegate;
-pub use undelegate_claim::{
-    UndelegateClaim, UndelegateClaimBody, UndelegateClaimPlan, UndelegateClaimProof,
-};
+#[cfg(feature = "component")]
+pub mod component;
 
-mod delegation_token;
-mod governance_key;
-mod identity_key;
-mod penalty;
-mod unbonding_token;
+#[cfg(feature = "component")]
+pub use component::{StateReadExt, StateWriteExt};
 
-pub use delegation_token::DelegationToken;
-pub use governance_key::GovernanceKey;
-pub use identity_key::IdentityKey;
-pub use penalty::Penalty;
-pub use unbonding_token::UnbondingToken;
+pub static BPS_SQUARED_SCALING_FACTOR: once_cell::sync::Lazy<penumbra_num::fixpoint::U128x128> =
+    once_cell::sync::Lazy::new(|| 1_0000_0000u128.into());
 
-pub use changes::DelegationChanges;
-pub use current_consensus_keys::CurrentConsensusKeys;
-pub use funding_stream::{FundingStream, FundingStreams};
-pub use uptime::Uptime;
+pub use self::delegate::Delegate;
+pub use self::undelegate::Undelegate;
+pub use self::undelegate_claim::{
+    UndelegateClaim, UndelegateClaimBody, UndelegateClaimPlan, UndelegateClaimProof,
+};
 
-pub mod genesis;
-pub mod params;
+pub use self::delegation_token::DelegationToken;
+pub use self::governance_key::GovernanceKey;
+pub use self::identity_key::IdentityKey;
+pub use self::penalty::Penalty;
+pub use self::unbonding_token::UnbondingToken;
+
+pub use self::changes::DelegationChanges;
+pub use self::current_consensus_keys::CurrentConsensusKeys;
+pub use self::funding_stream::{FundingStream, FundingStreams};
+pub use self::uptime::Uptime;