Return leaf when fetching account state so we can always add it to th…

…e HotShot state map

sequencer/src/api.rs

@@ -4,13 +4,13 @@ use anyhow::{bail, Context};
 use async_once_cell::Lazy;
 use async_std::sync::{Arc, RwLock};
 use async_trait::async_trait;
-use committable::Commitment;
+use committable::{Commitment, Committable};
 use data_source::{CatchupDataSource, SubmitDataSource};
 use derivative::Derivative;
 use espresso_types::{
     retain_accounts, v0::traits::SequencerPersistence, v0_3::ChainConfig, AccountQueryData,
     BlockMerkleTree, FeeAccount, FeeAccountProof, FeeMerkleTree, MockSequencerVersions, NodeState,
-    PubKey, Transaction,
+    PubKey, Transaction, ValidatedState,
 };
 use futures::{
     future::{BoxFuture, Future, FutureExt},
@@ -26,17 +26,17 @@ use hotshot_types::{
     data::ViewNumber,
     event::Event,
     light_client::StateSignatureRequestBody,
-    traits::{network::ConnectedNetwork, node_implementation::Versions},
-    utils::ViewInner,
+    traits::{network::ConnectedNetwork, node_implementation::Versions, ValidatedState as _},
+    utils::{View, ViewInner},
 };
 use jf_merkle_tree::MerkleTreeScheme;
 
 use self::data_source::{
     HotShotConfigDataSource, NodeStateDataSource, PublicNetworkConfig, StateSignatureDataSource,
 };
 use crate::{
-    context::Consensus, network, state_signature::StateSigner, SeqTypes, SequencerApiVersion,
-    SequencerContext,
+    catchup::CatchupStorage, context::Consensus, network, state_signature::StateSigner, SeqTypes,
+    SequencerApiVersion, SequencerContext,
 };
 
 pub mod data_source;
@@ -204,7 +204,7 @@ impl<
         N: ConnectedNetwork<PubKey>,
         V: Versions,
         P: SequencerPersistence,
-        D: CatchupDataSource + Send + Sync,
+        D: CatchupStorage + Send + Sync,
     > CatchupDataSource for StorageState<N, P, D, V>
 {
     #[tracing::instrument(skip(self, instance))]
@@ -228,25 +228,27 @@ impl<
         }
 
         // Try storage.
-        let tree = self
+        let (tree, leaf) = self
             .inner()
             .get_accounts(instance, height, view, accounts)
             .await
             .context("accounts not in memory, and could not fetch from storage")?;
         // If we successfully fetched accounts from storage, try to add them back into the in-memory
         // state.
         let handle = self.as_ref().consensus().await;
-        let consensus = handle.read().await.consensus();
+        let handle = handle.read().await;
+        let consensus = handle.consensus();
         let mut consensus = consensus.write().await;
-        if let Some(v) = consensus.validated_state_map().get(&view) {
-            // Clone the view info so we can update it.
-            let mut v = v.clone();
-            if let ViewInner::Leaf { state, .. } = &mut v.view_inner {
-                // Clone the state so we can update it.
-                let mut updated_state = (**state).clone();
+        let (state, delta, leaf_commit) = match consensus.validated_state_map().get(&view) {
+            Some(View {
+                view_inner: ViewInner::Leaf { state, delta, leaf },
+            }) => {
+                let mut state = (**state).clone();
+
+                // Add the fetched accounts to the state.
                 for account in accounts {
                     if let Some((proof, _)) = FeeAccountProof::prove(&tree, (*account).into()) {
-                        if let Err(err) = proof.remember(&mut updated_state.fee_merkle_tree) {
+                        if let Err(err) = proof.remember(&mut state.fee_merkle_tree) {
                             tracing::warn!(
                                 ?view,
                                 %account,
@@ -257,25 +259,36 @@ impl<
                         tracing::warn!(?view, %account, "cannot update fetched account state because account is not in the merkle tree");
                     };
                 }
-                // Update the state in the view.
-                *state = Arc::new(updated_state);
-                // Put the updated view back into the state map.
-                if let Err(err) = consensus.update_validated_state_map(view, v) {
-                    tracing::warn!(?view, "cannot update fetched account state: {err:#}");
-                }
-                tracing::info!(?view, "updated with fetched account state");
-            } else {
-                tracing::warn!(
-                    ?view,
-                    "cannot cache fetched account state because view has no state",
-                );
+
+                (Arc::new(state), delta.clone(), *leaf)
             }
-        } else {
-            tracing::warn!(
-                ?view,
-                "cannot cache fetched account state because view is not available"
-            );
+            _ => {
+                // If we don't already have a leaf for this view, or if we don't have the view
+                // at all, we can create a new view based on the recovered leaf and add it to
+                // our state map. In this case, we must also add the leaf to the saved leaves
+                // map to ensure consistency.
+                let mut state = ValidatedState::from_header(leaf.block_header());
+                state.fee_merkle_tree = tree.clone();
+                let res = (Arc::new(state), None, Committable::commit(&leaf));
+                consensus
+                    .update_saved_leaves(leaf, &handle.hotshot.upgrade_lock)
+                    .await;
+                res
+            }
+        };
+        if let Err(err) = consensus.update_validated_state_map(
+            view,
+            View {
+                view_inner: ViewInner::Leaf {
+                    state,
+                    delta,
+                    leaf: leaf_commit,
+                },
+            },
+        ) {
+            tracing::warn!(?view, "cannot update fetched account state: {err:#}");
         }
+        tracing::info!(?view, "updated with fetched account state");
 
         Ok(tree)
     }

sequencer/src/api/data_source.rs

@@ -1,6 +1,6 @@
 use std::{num::NonZeroUsize, time::Duration};
 
-use anyhow::{bail, Context};
+use anyhow::Context;
 use async_trait::async_trait;
 use committable::Commitment;
 use espresso_types::{
@@ -14,7 +14,7 @@ use hotshot_orchestrator::config::{
 };
 use hotshot_query_service::{
     availability::AvailabilityDataSource,
-    data_source::{MetricsDataSource, VersionedDataSource},
+    data_source::VersionedDataSource,
     fetching::provider::{AnyProvider, QueryServiceProvider},
     node::NodeDataSource,
     status::StatusDataSource,
@@ -145,19 +145,11 @@ pub(crate) trait CatchupDataSource: Sync {
     /// decided view.
     fn get_accounts(
         &self,
-        _instance: &NodeState,
-        _height: u64,
-        _view: ViewNumber,
-        _accounts: &[FeeAccount],
-    ) -> impl Send + Future<Output = anyhow::Result<FeeMerkleTree>> {
-        // Merklized state catchup is only supported by persistence backends that provide merklized
-        // state storage. This default implementation is overridden for those that do. Otherwise,
-        // catchup can still be provided by fetching undecided merklized state from consensus
-        // memory.
-        async {
-            bail!("merklized state catchup is not supported for this data source");
-        }
-    }
+        instance: &NodeState,
+        height: u64,
+        view: ViewNumber,
+        accounts: &[FeeAccount],
+    ) -> impl Send + Future<Output = anyhow::Result<FeeMerkleTree>>;
 
     /// Get the blocks Merkle tree frontier.
     ///
@@ -167,30 +159,16 @@ pub(crate) trait CatchupDataSource: Sync {
     /// decided view.
     fn get_frontier(
         &self,
-        _height: u64,
-        _view: ViewNumber,
-    ) -> impl Send + Future<Output = anyhow::Result<BlocksFrontier>> {
-        // Merklized state catchup is only supported by persistence backends that provide merklized
-        // state storage. This default implementation is overridden for those that do. Otherwise,
-        // catchup can still be provided by fetching undecided merklized state from consensus
-        // memory.
-        async {
-            bail!("merklized state catchup is not supported for this data source");
-        }
-    }
+        height: u64,
+        view: ViewNumber,
+    ) -> impl Send + Future<Output = anyhow::Result<BlocksFrontier>>;
 
     fn get_chain_config(
         &self,
-        _commitment: Commitment<ChainConfig>,
-    ) -> impl Send + Future<Output = anyhow::Result<ChainConfig>> {
-        async {
-            bail!("chain config catchup is not supported for this data source");
-        }
-    }
+        commitment: Commitment<ChainConfig>,
+    ) -> impl Send + Future<Output = anyhow::Result<ChainConfig>>;
 }
 
-impl CatchupDataSource for MetricsDataSource {}
-
 /// This struct defines the public Hotshot validator configuration.
 /// Private key and state key pairs are excluded for security reasons.
 

sequencer/src/api/fs.rs

@@ -3,8 +3,8 @@ use std::path::Path;
 use async_trait::async_trait;
 use hotshot_query_service::data_source::FileSystemDataSource;
 
-use super::data_source::{CatchupDataSource, Provider, SequencerDataSource};
-use crate::{persistence::fs::Options, SeqTypes};
+use super::data_source::{Provider, SequencerDataSource};
+use crate::{catchup::CatchupStorage, persistence::fs::Options, SeqTypes};
 
 pub type DataSource = FileSystemDataSource<SeqTypes, Provider>;
 
@@ -26,7 +26,7 @@ impl SequencerDataSource for DataSource {
     }
 }
 
-impl CatchupDataSource for DataSource {}
+impl CatchupStorage for DataSource {}
 
 #[cfg(test)]
 mod impl_testable_data_source {

sequencer/src/api/options.rs

@@ -35,6 +35,7 @@ use super::{
     ApiState, StorageState,
 };
 use crate::{
+    catchup::CatchupStorage,
     context::{SequencerContext, TaskList},
     persistence,
     state::update_state_storage_loop,
@@ -265,7 +266,7 @@ impl Options {
     where
         N: ConnectedNetwork<PubKey>,
         P: SequencerPersistence,
-        D: SequencerDataSource + CatchupDataSource + Send + Sync + 'static,
+        D: SequencerDataSource + CatchupStorage + Send + Sync + 'static,
         for<'a> D::Transaction<'a>: UpdateDataSource<SeqTypes>,
     {
         let metrics = ds.populate_metrics();