espresso_node/

catchup.rs

use std::{
    cmp::Ordering,
    collections::HashMap,
    fmt::{Debug, Display},
    sync::Arc,
    time::Duration,
};

use anyhow::{Context, anyhow, bail, ensure};
use async_lock::RwLock;
use async_trait::async_trait;
use committable::{Commitment, Committable};
use espresso_types::{
    BackoffParams, BlockMerkleTree, Certificate2, FeeAccount, FeeAccountProof, FeeMerkleCommitment,
    FeeMerkleTree, Leaf2, NodeState, SeqTypes, ValidatedState,
    config::PublicNetworkConfig,
    v0::traits::StateCatchup,
    v0_3::{
        ChainConfig, RewardAccountProofV1, RewardAccountV1, RewardMerkleCommitmentV1,
        RewardMerkleTreeV1,
    },
    v0_4::{
        PermittedRewardMerkleTreeV2, RewardAccountProofV2, RewardAccountV2,
        RewardMerkleCommitmentV2, RewardMerkleTreeV2, forgotten_accounts_include,
    },
};
use futures::{
    StreamExt,
    future::{Future, FutureExt, TryFuture, TryFutureExt},
    stream::FuturesUnordered,
};
use hotshot_new_protocol::{storage::NewProtocolStorage, utils::verify_new_protocol_leaf_chain};
use hotshot_types::{
    ValidatorConfig,
    data::{EpochNumber, ViewNumber},
    epoch_membership::EpochMembershipCoordinator,
    message::UpgradeLock,
    network::NetworkConfig,
    simple_certificate::LightClientStateUpdateCertificateV2,
    traits::{
        ValidatedState as ValidatedStateTrait,
        metrics::{Counter, CounterFamily, Metrics},
    },
    utils::{epoch_from_block_number, verify_leaf_chain},
};
use itertools::Itertools;
use jf_merkle_tree_compat::{ForgetableMerkleTreeScheme, MerkleTreeScheme, prelude::MerkleNode};
use parking_lot::Mutex;
use priority_queue::PriorityQueue;
use serde::de::DeserializeOwned;
use surf_disco::Request;
use tide_disco::error::ServerError;
use tokio::time::timeout;
use tokio_util::task::AbortOnDropHandle;
use url::Url;
use vbs::version::StaticVersionType;
use versions::{EPOCH_VERSION, NEW_PROTOCOL_VERSION};

use crate::{
    api::{BlocksFrontier, RewardMerkleTreeDataSource, RewardMerkleTreeV2Data},
    consensus_handle::ConsensusHandle,
};

// This newtype is probably not worth having. It's only used to be able to log
// URLs before doing requests.
#[derive(Debug, Clone)]
struct Client<ServerError, ApiVer: StaticVersionType> {
    inner: surf_disco::Client<ServerError, ApiVer>,
    url: Url,
    requests: Arc<Box<dyn Counter>>,
    failures: Arc<Box<dyn Counter>>,
}

impl<ApiVer: StaticVersionType> Client<ServerError, ApiVer> {
    pub fn new(
        url: Url,
        requests: &(impl CounterFamily + ?Sized),
        failures: &(impl CounterFamily + ?Sized),
    ) -> Self {
        Self {
            inner: surf_disco::Client::new(url.clone()),
            requests: Arc::new(requests.create(vec![url.to_string()])),
            failures: Arc::new(failures.create(vec![url.to_string()])),
            url,
        }
    }

    pub fn get<T: DeserializeOwned>(&self, route: &str) -> Request<T, ServerError, ApiVer> {
        self.inner.get(route)
    }
}

/// A score of a catchup peer, based on our interactions with that peer.
///
/// The score accounts for malicious peers -- i.e. peers that gave us an invalid response to a
/// verifiable request -- and faulty/unreliable peers -- those that fail to respond to requests at
/// all. The score has a comparison function where higher is better, or in other words `p1 > p2`
/// means we believe we are more likely to successfully catch up using `p1` than `p2`. This makes it
/// convenient and efficient to collect peers in a priority queue which we can easily convert to a
/// list sorted by reliability.
#[derive(Clone, Copy, Debug, Default)]
struct PeerScore {
    requests: usize,
    failures: usize,
}

impl Ord for PeerScore {
    fn cmp(&self, other: &Self) -> Ordering {
        // Compare failure rates: `self` is better than `other` if
        //      self.failures / self.requests < other.failures / other.requests
        // or equivalently
        //      other.failures * self.requests > self.failures * other.requests
        (other.failures * self.requests).cmp(&(self.failures * other.requests))
    }
}

impl PartialOrd for PeerScore {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for PeerScore {
    fn eq(&self, other: &Self) -> bool {
        self.cmp(other).is_eq()
    }
}

impl Eq for PeerScore {}

#[derive(Debug, Clone, Default)]
pub struct StatePeers<ApiVer: StaticVersionType> {
    // Peer IDs, ordered by reliability score. Each ID is an index into `clients`.
    scores: Arc<RwLock<PriorityQueue<usize, PeerScore>>>,
    clients: Vec<Client<ServerError, ApiVer>>,
    backoff: BackoffParams,
    /// Base timeout for per peer catchup request
    base_timeout: Duration,
}

impl<ApiVer: StaticVersionType> StatePeers<ApiVer> {
    async fn fetch<Fut>(
        &self,
        retry: usize,
        f: impl Fn(Client<ServerError, ApiVer>) -> Fut,
    ) -> anyhow::Result<Fut::Ok>
    where
        Fut: TryFuture<Error: Display>,
    {
        // Since we have generally have multiple peers we can catch up from, we want a fairly
        // aggressive timeout for requests: if a peer is not responding quickly, we're better off
        // just trying the next one rather than waiting, and this prevents a malicious peer from
        // delaying catchup for a long time.
        //
        // However, if we set the timeout _too_ aggressively, we might fail to catch up even from an
        // honest peer, and thus never make progress. Thus, we start with a base timeout (default
        // 2s), which is reasonable for an HTTP request. If that fails with all of our peers, we
        // increase the timeout by the base amount for each successive retry, until we eventually
        // succeed. The base timeout is configurable via ESPRESSO_NODE_CATCHUP_BASE_TIMEOUT.
        let timeout_dur = self.base_timeout * (retry as u32 + 1);

        // Keep track of which peers we make requests to and which succeed (`true`) or fail (`false`),
        // so we can update reliability scores at the end.
        let mut requests = HashMap::new();
        let mut res = Err(anyhow!("failed fetching from every peer"));

        // Try each peer in order of reliability score, until we succeed. We clone out of
        // `self.scores` because it is small (contains only numeric IDs and scores), so this clone
        // is a lot cheaper than holding the read lock the entire time we are making requests (which
        // could be a while).
        let mut scores = { (*self.scores.read().await).clone() };
        let mut logs = vec![format!("Fetching failed.\n")];
        while let Some((id, score)) = scores.pop() {
            let client = &self.clients[id];
            tracing::info!("fetching from {}", client.url);
            match timeout(timeout_dur, TryFutureExt::into_future(f(client.clone()))).await {
                Ok(Ok(t)) => {
                    requests.insert(id, true);
                    res = Ok(t);
                    logs = Vec::new();
                    break;
                },
                Ok(Err(err)) => {
                    tracing::debug!(id, ?score, peer = %client.url, "error from peer: {err:#}");
                    logs.push(format!(
                        "Error from peer {} with id {id} and score {score:?}: {err:#}",
                        client.url
                    ));
                    requests.insert(id, false);
                },
                Err(_) => {
                    tracing::debug!(id, ?score, peer = %client.url, ?timeout_dur, "request timed out");
                    logs.push(format!(
                        "Error from peer {} with id {id} and score {score:?}: request timed out",
                        client.url
                    ));
                    requests.insert(id, false);
                },
            }
        }

        if !logs.is_empty() {
            tracing::warn!("{}", logs.join("\n"));
        }

        // Update client scores.
        let mut scores = self.scores.write().await;
        for (id, success) in requests {
            scores.change_priority_by(&id, |score| {
                score.requests += 1;
                self.clients[id].requests.add(1);
                if !success {
                    score.failures += 1;
                    self.clients[id].failures.add(1);
                }
            });
        }

        res
    }

    pub fn from_urls(
        urls: Vec<Url>,
        backoff: BackoffParams,
        base_timeout: Duration,
        metrics: &(impl Metrics + ?Sized),
    ) -> Self {
        if urls.is_empty() {
            panic!("Cannot create StatePeers with no peers");
        }

        let metrics = metrics.subgroup("catchup".into());
        let requests = metrics.counter_family("requests".into(), vec!["peer".into()]);
        let failures = metrics.counter_family("request_failures".into(), vec!["peer".into()]);

        let scores = urls
            .iter()
            .enumerate()
            .map(|(i, _)| (i, PeerScore::default()))
            .collect();
        let clients = urls
            .into_iter()
            .map(|url| Client::new(url, &*requests, &*failures))
            .collect();

        Self {
            clients,
            scores: Arc::new(RwLock::new(scores)),
            backoff,
            base_timeout,
        }
    }

    #[tracing::instrument(skip(self, my_own_validator_config))]
    pub async fn fetch_config(
        &self,
        my_own_validator_config: ValidatorConfig<SeqTypes>,
    ) -> anyhow::Result<NetworkConfig<SeqTypes>> {
        self.backoff()
            .retry(self, move |provider, retry| {
                let my_own_validator_config = my_own_validator_config.clone();
                async move {
                    let cfg: PublicNetworkConfig = provider
                        .fetch(retry, |client| {
                            let url = client.url.join("config/hotshot").unwrap();

                            reqwest::get(url.clone())
                        })
                        .await?
                        .json()
                        .await?;
                    cfg.into_network_config(my_own_validator_config)
                        .context("fetched config, but failed to convert to private config")
                }
                .boxed()
            })
            .await
    }
}

/// Verify a legacy (pre-V6) leaf chain
pub(crate) async fn verify_legacy_leaf_chain(
    leaf_chain: Vec<Leaf2>,
    coordinator: &EpochMembershipCoordinator<SeqTypes>,
    height: u64,
) -> anyhow::Result<Leaf2> {
    let upgrade_lock = UpgradeLock::<SeqTypes>::new(versions::Upgrade::trivial(EPOCH_VERSION));
    let epoch = EpochNumber::new(epoch_from_block_number(height, *coordinator.epoch_height()));
    let membership = coordinator
        .stake_table_for_epoch(Some(epoch))
        .map_err(|err| anyhow!("no stake table available for epoch {epoch}: {err:?}"))?;
    let stake_table: Vec<_> = membership.stake_table().cloned().collect();
    verify_leaf_chain(
        leaf_chain,
        &stake_table,
        membership.success_threshold(),
        height,
        &upgrade_lock,
    )
    .await
    .with_context(|| format!("failed to verify leaf chain at height {height}"))
}

#[async_trait]
impl<ApiVer: StaticVersionType> StateCatchup for StatePeers<ApiVer> {
    #[tracing::instrument(skip(self, _instance))]
    async fn try_fetch_accounts(
        &self,
        retry: usize,
        _instance: &NodeState,
        height: u64,
        view: ViewNumber,
        fee_merkle_tree_root: FeeMerkleCommitment,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        self.fetch(retry, |client| async move {
            let tree = client
                .inner
                .post::<FeeMerkleTree>(&format!("catchup/{height}/{}/accounts", view.u64()))
                .body_binary(&accounts.to_vec())?
                .send()
                .await?;

            // Verify proofs.
            let mut proofs = Vec::new();
            for account in accounts {
                let (proof, _) = FeeAccountProof::prove(&tree, (*account).into())
                    .context(format!("response missing fee account {account}"))?;
                proof.verify(&fee_merkle_tree_root).context(format!(
                    "invalid proof for fee account {account}, root: {fee_merkle_tree_root}"
                ))?;
                proofs.push(proof);
            }

            anyhow::Ok(proofs)
        })
        .await
    }

    #[tracing::instrument(skip(self, _instance, mt))]
    async fn try_remember_blocks_merkle_tree(
        &self,
        retry: usize,
        _instance: &NodeState,
        height: u64,
        view: ViewNumber,
        mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        *mt = self
            .fetch(retry, |client| {
                let mut mt = mt.clone();
                async move {
                    let frontier = client
                        .get::<BlocksFrontier>(&format!("catchup/{height}/{}/blocks", view.u64()))
                        .send()
                        .await?;
                    let elem = frontier
                        .elem()
                        .context("provided frontier is missing leaf element")?;
                    mt.remember(mt.num_leaves() - 1, *elem, &frontier)
                        .context("verifying block proof")?;
                    anyhow::Ok(mt)
                }
            })
            .await?;
        Ok(())
    }

    async fn try_fetch_chain_config(
        &self,
        retry: usize,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        self.fetch(retry, |client| async move {
            let cf = client
                .get::<ChainConfig>(&format!("catchup/chain-config/{commitment}"))
                .send()
                .await?;
            ensure!(
                cf.commit() == commitment,
                "received chain config with mismatched commitment: expected {commitment}, got {}",
                cf.commit()
            );
            Ok(cf)
        })
        .await
    }

    async fn try_fetch_leaf(
        &self,
        retry: usize,
        coordinator: EpochMembershipCoordinator<SeqTypes>,
        height: u64,
    ) -> anyhow::Result<Leaf2> {
        // Fetch the leaf chain. For new protocol heights this is a leaf range
        // `[height..=cert2_height]`
        // for legacy-protocol heights it's a 3-chain.
        let leaf_chain = self
            .fetch(retry, |client| async move {
                let chain = client
                    .get::<Vec<Leaf2>>(&format!("catchup/{height}/leafchain"))
                    .send()
                    .await?;
                anyhow::Ok(chain)
            })
            .await
            .with_context(|| format!("failed to fetch leaf chain at height {height}"))?;

        let first = leaf_chain
            .first()
            .ok_or_else(|| anyhow!("empty leaf chain returned for height {height}"))?;

        if first.block_header().version() >= NEW_PROTOCOL_VERSION {
            let upgrade_lock =
                UpgradeLock::<SeqTypes>::new(versions::Upgrade::trivial(NEW_PROTOCOL_VERSION));

            // The chain terminates at the leaf the cert2 finalizes, so fetch that cert2 by its
            // exact height
            let cert2_height = leaf_chain
                .last()
                .ok_or_else(|| anyhow!("empty leaf chain returned for height {height}"))?
                .height();
            let cert2 = self
                .fetch(retry, |client| async move {
                    let cert2 = client
                        .get::<Certificate2<SeqTypes>>(&format!("catchup/{cert2_height}/cert2"))
                        .send()
                        .await?;
                    anyhow::Ok(cert2)
                })
                .await
                .with_context(|| format!("failed to fetch cert2 for height {cert2_height}"))?;

            verify_new_protocol_leaf_chain(leaf_chain, &coordinator, height, &upgrade_lock, cert2)
                .await
                .with_context(|| {
                    format!("failed to verify leaf chain with cert2 at height {height}")
                })
        } else {
            verify_legacy_leaf_chain(leaf_chain, &coordinator, height).await
        }
    }

    async fn try_fetch_reward_merkle_tree_v2(
        &self,
        retry: usize,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV2,
        accounts: Arc<Vec<RewardAccountV2>>,
    ) -> anyhow::Result<PermittedRewardMerkleTreeV2> {
        let result = self
            .fetch(retry, |client| async move {
                // Try the catchup endpoint first which returns tree from consensuss decided state
                // if not present, then fall back to
                // the reward-state-v2 endpoint which returns from storage decided state
                let tree_bytes = match client
                    .inner
                    .get::<Vec<u8>>(&format!("catchup/reward-merkle-tree-v2/{height}/{}", *view))
                    .send()
                    .await
                {
                    Ok(bytes) => bytes,
                    Err(err) => {
                        tracing::info!(
                            "catchup endpoint failed, falling back to reward-state-v2: {err:#}"
                        );
                        client
                            .inner
                            .get::<Vec<u8>>(&format!(
                                "reward-state-v2/reward-merkle-tree-v2/{height}"
                            ))
                            .send()
                            .await?
                    },
                };

                Ok::<Vec<u8>, anyhow::Error>(tree_bytes)
            })
            .await
            .context("Fetching from peer failed")?;

        let tree_data = bincode::deserialize::<RewardMerkleTreeV2Data>(&result)
            .context("Failed to deserialize merkle tree from catchup")?;

        let tree: PermittedRewardMerkleTreeV2 =
            PermittedRewardMerkleTreeV2::try_from_kv_set(tree_data.balances).await?;

        ensure!(
            tree.tree.commitment() == reward_merkle_tree_root,
            "RewardMerkleTreeV2 from peer failed commitment check."
        );
        ensure!(!forgotten_accounts_include(&tree, &accounts));

        Ok(tree)
    }

    #[tracing::instrument(skip(self, _instance))]
    async fn try_fetch_reward_accounts_v1(
        &self,
        retry: usize,
        _instance: &NodeState,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV1,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        self.fetch(retry, |client| async move {
            let tree = client
                .inner
                .post::<RewardMerkleTreeV1>(&format!(
                    "catchup/{height}/{}/reward-accounts",
                    view.u64()
                ))
                .body_binary(&accounts.to_vec())?
                .send()
                .await?;

            // Verify proofs.
            let mut proofs = Vec::new();
            for account in accounts {
                let (proof, _) = RewardAccountProofV1::prove(&tree, (*account).into())
                    .context(format!("response missing reward account {account}"))?;
                proof.verify(&reward_merkle_tree_root).context(format!(
                    "invalid proof for v1 reward account {account}, root: \
                     {reward_merkle_tree_root} height {height} view {view}"
                ))?;
                proofs.push(proof);
            }

            anyhow::Ok(proofs)
        })
        .await
    }

    async fn try_fetch_state_cert(
        &self,
        retry: usize,
        epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        self.fetch(retry, |client| async move {
            client
                .get::<LightClientStateUpdateCertificateV2<SeqTypes>>(&format!(
                    "catchup/{epoch}/state-cert"
                ))
                .send()
                .await
        })
        .await
    }

    fn backoff(&self) -> &BackoffParams {
        &self.backoff
    }

    fn name(&self) -> String {
        format!(
            "StatePeers({})",
            self.clients
                .iter()
                .map(|client| client.url.to_string())
                .join(",")
        )
    }

    fn is_local(&self) -> bool {
        false
    }
}

pub(crate) trait CatchupStorage: Sync {
    /// Get the state of the requested `accounts`.
    ///
    /// The state is fetched from a snapshot at the given height and view, which _must_ correspond!
    /// `height` is provided to simplify lookups for backends where data is not indexed by view.
    /// This function is intended to be used for catchup, so `view` should be no older than the last
    /// decided view.
    ///
    /// If successful, this function also returns the leaf from `view`, if it is available. This can
    /// be used to add the recovered state to HotShot's state map, so that future requests can get
    /// the state from memory rather than storage.
    fn get_accounts(
        &self,
        _instance: &NodeState,
        _height: u64,
        _view: ViewNumber,
        _accounts: &[FeeAccount],
    ) -> impl Send + Future<Output = anyhow::Result<(FeeMerkleTree, Leaf2)>> {
        // 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");
        }
    }

    fn get_reward_accounts_v1(
        &self,
        _instance: &NodeState,
        _height: u64,
        _view: ViewNumber,
        _accounts: &[RewardAccountV1],
    ) -> impl Send + Future<Output = anyhow::Result<(RewardMerkleTreeV1, Leaf2)>> {
        async {
            bail!("merklized state catchup is not supported for this data source");
        }
    }

    fn get_reward_accounts_v2(
        &self,
        _instance: &NodeState,
        _height: u64,
        _view: ViewNumber,
        _accounts: &[RewardAccountV2],
    ) -> impl Send + Future<Output = anyhow::Result<(RewardMerkleTreeV2, Leaf2)>> {
        async {
            bail!("merklized state catchup is not supported for this data source");
        }
    }

    /// Get the blocks Merkle tree frontier.
    ///
    /// The state is fetched from a snapshot at the given height and view, which _must_ correspond!
    /// `height` is provided to simplify lookups for backends where data is not indexed by view.
    /// This function is intended to be used for catchup, so `view` should be no older than the last
    /// decided view.
    fn get_frontier(
        &self,
        _instance: &NodeState,
        _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");
        }
    }

    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");
        }
    }

    fn get_leaf_chain(
        &self,
        _height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Vec<Leaf2>>> {
        async {
            bail!("leaf chain catchup is not supported for this data source");
        }
    }

    /// Load the cert2 stored at exactly `height`, if one exists.
    fn load_cert2(
        &self,
        _height: u64,
    ) -> impl Send + Future<Output = anyhow::Result<Option<Certificate2<SeqTypes>>>> {
        async { Ok(None) }
    }

    /// Load a decided leaf at the given height.
    fn get_leaf(&self, _height: u64) -> impl Send + Future<Output = anyhow::Result<Leaf2>> {
        async {
            bail!("leaf fetch is not supported for this data source");
        }
    }
}

impl CatchupStorage for hotshot_query_service::data_source::MetricsDataSource {}

impl<T, S> CatchupStorage for hotshot_query_service::data_source::ExtensibleDataSource<T, S>
where
    T: CatchupStorage,
    S: Sync,
{
    async fn get_accounts(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<(FeeMerkleTree, Leaf2)> {
        self.inner()
            .get_accounts(instance, height, view, accounts)
            .await
    }

    async fn get_reward_accounts_v2(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV2],
    ) -> anyhow::Result<(RewardMerkleTreeV2, Leaf2)> {
        self.inner()
            .get_reward_accounts_v2(instance, height, view, accounts)
            .await
    }

    async fn get_reward_accounts_v1(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<(RewardMerkleTreeV1, Leaf2)> {
        self.inner()
            .get_reward_accounts_v1(instance, height, view, accounts)
            .await
    }

    async fn get_frontier(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
    ) -> anyhow::Result<BlocksFrontier> {
        self.inner().get_frontier(instance, height, view).await
    }

    async fn get_chain_config(
        &self,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        self.inner().get_chain_config(commitment).await
    }
    async fn get_leaf_chain(&self, height: u64) -> anyhow::Result<Vec<Leaf2>> {
        self.inner().get_leaf_chain(height).await
    }

    async fn load_cert2(&self, height: u64) -> anyhow::Result<Option<Certificate2<SeqTypes>>> {
        self.inner().load_cert2(height).await
    }

    async fn get_leaf(&self, height: u64) -> anyhow::Result<Leaf2> {
        self.inner().get_leaf(height).await
    }
}

#[derive(Debug)]
pub(crate) struct SqlStateCatchup<T> {
    db: Arc<T>,
    backoff: BackoffParams,
}

impl<T> SqlStateCatchup<T> {
    pub(crate) fn new(db: Arc<T>, backoff: BackoffParams) -> Self {
        Self { db, backoff }
    }
}

#[async_trait]
impl<T> StateCatchup for SqlStateCatchup<T>
where
    T: CatchupStorage + RewardMerkleTreeDataSource + Send + Sync,
{
    async fn try_fetch_leaf(
        &self,
        _retry: usize,
        _coordinator: EpochMembershipCoordinator<SeqTypes>,
        height: u64,
    ) -> anyhow::Result<Leaf2> {
        // Leaves in our local DB were verified before they were stored, so we can return the leaf
        // at `height` directly without re-verifying.
        self.db
            .get_leaf(height)
            .await
            .with_context(|| format!("failed to load leaf at height {height} from DB"))
    }

    // TODO: add a test for the account proof validation
    // issue # 2102 (https://github.com/EspressoSystems/espresso-network/issues/2102)
    #[tracing::instrument(skip(self, _retry, instance))]
    async fn try_fetch_accounts(
        &self,
        _retry: usize,
        instance: &NodeState,
        block_height: u64,
        view: ViewNumber,
        fee_merkle_tree_root: FeeMerkleCommitment,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        // Get the accounts
        let (fee_merkle_tree_from_db, _) = self
            .db
            .get_accounts(instance, block_height, view, accounts)
            .await
            .with_context(|| "failed to get fee accounts from DB")?;

        // Verify the accounts
        let mut proofs = Vec::new();
        for account in accounts {
            let (proof, _) = FeeAccountProof::prove(&fee_merkle_tree_from_db, (*account).into())
                .context(format!("response missing account {account}"))?;
            proof.verify(&fee_merkle_tree_root).context(format!(
                "invalid proof for fee account {account}, root: {fee_merkle_tree_root}"
            ))?;
            proofs.push(proof);
        }

        Ok(proofs)
    }

    #[tracing::instrument(skip(self, _retry, instance, mt))]
    async fn try_remember_blocks_merkle_tree(
        &self,
        _retry: usize,
        instance: &NodeState,
        bh: u64,
        view: ViewNumber,
        mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        if bh == 0 {
            return Ok(());
        }

        let proof = self.db.get_frontier(instance, bh, view).await?;
        match proof
            .proof
            .first()
            .context(format!("empty proof for frontier at height {bh}"))?
        {
            MerkleNode::Leaf { pos, elem, .. } => mt
                .remember(pos, elem, proof.clone())
                .context("failed to remember proof"),
            _ => bail!("invalid proof"),
        }
    }

    async fn try_fetch_chain_config(
        &self,
        _retry: usize,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        let cf = self.db.get_chain_config(commitment).await?;

        if cf.commit() != commitment {
            panic!(
                "Critical error: Mismatched chain config detected. Expected chain config: {:?}, \
                 but got: {:?}.
                This may indicate a compromised database",
                commitment,
                cf.commit()
            )
        }

        Ok(cf)
    }

    async fn try_fetch_reward_merkle_tree_v2(
        &self,
        _retry: usize,
        height: u64,
        _view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV2,
        accounts: Arc<Vec<RewardAccountV2>>,
    ) -> anyhow::Result<PermittedRewardMerkleTreeV2> {
        let tree: PermittedRewardMerkleTreeV2 = self.db.load_reward_merkle_tree_v2(height).await?;

        ensure!(tree.tree.commitment() == reward_merkle_tree_root);
        ensure!(!forgotten_accounts_include(&tree, &accounts));

        Ok(tree)
    }

    #[tracing::instrument(skip(self, _retry, instance))]
    async fn try_fetch_reward_accounts_v1(
        &self,
        _retry: usize,
        instance: &NodeState,
        block_height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV1,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        // Get the accounts
        let (reward_merkle_tree_from_db, _) = self
            .db
            .get_reward_accounts_v1(instance, block_height, view, accounts)
            .await
            .with_context(|| "failed to get reward accounts from DB")?;
        // Verify the accounts
        let mut proofs = Vec::new();
        for account in accounts {
            let (proof, _) =
                RewardAccountProofV1::prove(&reward_merkle_tree_from_db, (*account).into())
                    .context(format!("response missing account {account}"))?;
            proof.verify(&reward_merkle_tree_root).context(format!(
                "invalid proof for v1 reward account {account}, root: {reward_merkle_tree_root}"
            ))?;
            proofs.push(proof);
        }

        Ok(proofs)
    }

    async fn try_fetch_state_cert(
        &self,
        _retry: usize,
        _epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        bail!("state cert catchup not supported for SqlStateCatchup");
    }

    fn backoff(&self) -> &BackoffParams {
        &self.backoff
    }

    fn name(&self) -> String {
        "SqlStateCatchup".into()
    }

    fn is_local(&self) -> bool {
        true
    }
}

/// Disable catchup entirely.
#[derive(Clone, Debug)]
pub struct NullStateCatchup {
    backoff: BackoffParams,
    chain_configs: HashMap<Commitment<ChainConfig>, ChainConfig>,
}

impl Default for NullStateCatchup {
    fn default() -> Self {
        Self {
            backoff: BackoffParams::disabled(),
            chain_configs: Default::default(),
        }
    }
}

impl NullStateCatchup {
    /// Add a chain config preimage which can be fetched by hash during STF evaluation.
    ///
    /// [`NullStateCatchup`] is used to disable catchup entirely when evaluating the STF, which
    /// requires the [`ValidatedState`](espresso_types::ValidatedState) to be pre-seeded with all
    /// the dependencies of STF evaluation. However, the STF also depends on having the preimage of
    /// various [`ChainConfig`] commitments, which are not stored in the
    /// [`ValidatedState`](espresso_types::ValidatedState), but which instead must be supplied by a
    /// separate preimage oracle. Thus, [`NullStateCatchup`] may be populated with a set of
    /// [`ChainConfig`]s, which it can feed to the STF during evaluation.
    pub fn add_chain_config(&mut self, cf: ChainConfig) {
        self.chain_configs.insert(cf.commit(), cf);
    }
}

#[async_trait]
impl StateCatchup for NullStateCatchup {
    async fn try_fetch_leaf(
        &self,
        _retry: usize,
        _coordinator: EpochMembershipCoordinator<SeqTypes>,
        _height: u64,
    ) -> anyhow::Result<Leaf2> {
        bail!("state catchup is disabled")
    }

    async fn try_fetch_accounts(
        &self,
        _retry: usize,
        _instance: &NodeState,
        _height: u64,
        _view: ViewNumber,
        _fee_merkle_tree_root: FeeMerkleCommitment,
        _account: &[FeeAccount],
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        bail!("state catchup is disabled");
    }

    async fn try_remember_blocks_merkle_tree(
        &self,
        _retry: usize,
        _instance: &NodeState,
        _height: u64,
        _view: ViewNumber,
        _mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        bail!("state catchup is disabled");
    }

    async fn try_fetch_chain_config(
        &self,
        _retry: usize,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        self.chain_configs
            .get(&commitment)
            .copied()
            .context(format!("chain config {commitment} not available"))
    }

    async fn try_fetch_reward_merkle_tree_v2(
        &self,
        _retry: usize,
        _height: u64,
        _view: ViewNumber,
        _reward_merkle_tree_root: RewardMerkleCommitmentV2,
        _accounts: Arc<Vec<RewardAccountV2>>,
    ) -> anyhow::Result<PermittedRewardMerkleTreeV2> {
        bail!("state catchup is disabled");
    }

    async fn try_fetch_reward_accounts_v1(
        &self,
        _retry: usize,
        _instance: &NodeState,
        _height: u64,
        _view: ViewNumber,
        _fee_merkle_tree_root: RewardMerkleCommitmentV1,
        _account: &[RewardAccountV1],
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        bail!("state catchup is disabled");
    }

    async fn try_fetch_state_cert(
        &self,
        _retry: usize,
        _epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        bail!("state catchup is disabled");
    }

    fn backoff(&self) -> &BackoffParams {
        &self.backoff
    }

    fn name(&self) -> String {
        "NullStateCatchup".into()
    }

    fn is_local(&self) -> bool {
        true
    }
}

/// A catchup implementation that parallelizes requests to many providers.
/// It returns the result of the first non-erroring provider to complete.
#[derive(Clone)]
pub struct ParallelStateCatchup {
    providers: Arc<Mutex<Vec<Arc<dyn StateCatchup>>>>,
    backoff: BackoffParams,
    /// Timeout for local provider requests
    local_timeout: Duration,
}

impl ParallelStateCatchup {
    /// Create a new [`ParallelStateCatchup`] with the given providers and local timeout.
    pub fn new(providers: &[Arc<dyn StateCatchup>], local_timeout: Duration) -> Self {
        Self {
            providers: Arc::new(Mutex::new(providers.to_vec())),
            backoff: BackoffParams::disabled(),
            local_timeout,
        }
    }

    /// Add a provider to the list of providers
    pub fn add_provider(&self, provider: Arc<dyn StateCatchup>) {
        self.providers.lock().push(provider);
    }

    /// Perform an async operation on all local providers, returning the first result to succeed.
    ///
    /// A timeout is applied so that a slow local lookup does not prevent the node from
    /// falling back to remote providers in time to vote within the current view.
    pub async fn on_local_providers<C, F, RT>(&self, closure: C) -> anyhow::Result<RT>
    where
        C: Fn(Arc<dyn StateCatchup>) -> F + Clone + Send + Sync + 'static,
        F: Future<Output = anyhow::Result<RT>> + Send + 'static,
        RT: Send + Sync + 'static,
    {
        let local_timeout = self.local_timeout;
        match timeout(
            local_timeout,
            self.on_providers(|provider| provider.is_local(), closure),
        )
        .await
        {
            Ok(result) => result,
            Err(_) => {
                let err = format!("local provider timed out after {local_timeout:?}");
                tracing::warn!("{err}");
                Err(anyhow::anyhow!(err))
            },
        }
    }

    /// Perform an async operation on all remote providers, returning the first result to succeed
    pub async fn on_remote_providers<C, F, RT>(&self, closure: C) -> anyhow::Result<RT>
    where
        C: Fn(Arc<dyn StateCatchup>) -> F + Clone + Send + Sync + 'static,
        F: Future<Output = anyhow::Result<RT>> + Send + 'static,
        RT: Send + Sync + 'static,
    {
        self.on_providers(|provider| !provider.is_local(), closure)
            .await
    }

    /// Perform an async operation on all providers matching the given predicate, returning the first result to succeed
    pub async fn on_providers<P, C, F, RT>(&self, predicate: P, closure: C) -> anyhow::Result<RT>
    where
        P: Fn(&Arc<dyn StateCatchup>) -> bool + Clone + Send + Sync + 'static,
        C: Fn(Arc<dyn StateCatchup>) -> F + Clone + Send + Sync + 'static,
        F: Future<Output = anyhow::Result<RT>> + Send + 'static,
        RT: Send + Sync + 'static,
    {
        // Make sure we have at least one provider
        let providers = self.providers.lock().clone();
        if providers.is_empty() {
            return Err(anyhow::anyhow!("no providers were initialized"));
        }

        // Filter the providers by the predicate
        let providers = providers.into_iter().filter(predicate).collect::<Vec<_>>();
        if providers.is_empty() {
            return Err(anyhow::anyhow!("no providers matched the given predicate"));
        }

        // Spawn futures for each provider
        let mut futures = FuturesUnordered::new();
        for provider in providers {
            let closure = closure.clone();
            futures.push(AbortOnDropHandle::new(tokio::spawn(closure(provider))));
        }

        let mut logs = vec![format!("No providers returned a successful result.\n")];
        // Return the first successful result
        while let Some(result) = futures.next().await {
            // Unwrap the inner (join) result
            let result = match result {
                Ok(res) => res,
                Err(err) => {
                    tracing::debug!("Failed to join on provider: {err:#}.");
                    logs.push(format!("Failed to join on provider: {err:#}."));
                    continue;
                },
            };

            // If a provider fails, print why
            let result = match result {
                Ok(res) => res,
                Err(err) => {
                    tracing::debug!("Failed to fetch data: {err:#}.");
                    logs.push(format!("Failed to fetch data: {err:#}."));
                    continue;
                },
            };

            return Ok(result);
        }

        Err(anyhow::anyhow!(logs.join("\n")))
    }
}

macro_rules! clone {
    ( ($( $x:ident ),*) $y:expr ) => {
        {
            $(let $x = $x.clone();)*
            $y
        }
    };
}

/// A catchup implementation that parallelizes requests to a local and remote provider.
/// It returns the result of the first provider to complete.
#[async_trait]
impl StateCatchup for ParallelStateCatchup {
    async fn try_fetch_leaf(
        &self,
        retry: usize,
        coordinator: EpochMembershipCoordinator<SeqTypes>,
        height: u64,
    ) -> anyhow::Result<Leaf2> {
        // Try fetching the leaf on the local providers first
        let local_result = self
            .on_local_providers(clone! {(coordinator) move |provider| {
                clone!{(coordinator) async move {
                    provider
                        .try_fetch_leaf(retry, coordinator, height)
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones
        self.on_remote_providers(clone! {(coordinator) move |provider| {
            clone!{(coordinator) async move {
                provider
                    .try_fetch_leaf(retry, coordinator, height)
                    .await
            }}
        }})
        .await
    }

    async fn try_fetch_accounts(
        &self,
        retry: usize,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        fee_merkle_tree_root: FeeMerkleCommitment,
        accounts: &[FeeAccount],
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        // Try to get the accounts on local providers first
        let accounts_vec = accounts.to_vec();
        let local_result = self
            .on_local_providers(clone! {(instance, accounts_vec) move |provider| {
                clone! {(instance, accounts_vec) async move {
                    provider
                        .try_fetch_accounts(
                            retry,
                            &instance,
                            height,
                            view,
                            fee_merkle_tree_root,
                            &accounts_vec,
                        )
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones
        self.on_remote_providers(clone! {(instance, accounts_vec) move |provider| {
            clone!{(instance, accounts_vec) async move {
                provider
                .try_fetch_accounts(
                    retry,
                    &instance,
                    height,
                    view,
                    fee_merkle_tree_root,
                    &accounts_vec,
                ).await
            }}
        }})
        .await
    }

    async fn try_remember_blocks_merkle_tree(
        &self,
        retry: usize,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        // Try to remember the blocks merkle tree on local providers first
        let local_result = self
            .on_local_providers(clone! {(mt, instance) move |provider| {
                let mut mt = mt.clone();
                clone! {(instance) async move {
                    // Perform the call
                    provider
                        .try_remember_blocks_merkle_tree(
                            retry,
                            &instance,
                            height,
                            view,
                            &mut mt,
                        )
                        .await?;

                    // Return the merkle tree so we can modify it
                    Ok(mt)
                }}
            }})
            .await;

        // Check if we were successful locally
        if let Ok(modified_mt) = local_result {
            // Set the merkle tree to the output of the successful local call
            *mt = modified_mt;

            return Ok(());
        }

        // If that fails, try the remote ones
        let remote_result = self
            .on_remote_providers(clone! {(mt, instance) move |provider| {
                let mut mt = mt.clone();
                clone!{(instance) async move {
                    // Perform the call
                    provider
                    .try_remember_blocks_merkle_tree(
                        retry,
                        &instance,
                        height,
                        view,
                        &mut mt,
                    )
                    .await?;

                    // Return the merkle tree
                    Ok(mt)
                }}
            }})
            .await?;

        // Update the original, local merkle tree
        *mt = remote_result;

        Ok(())
    }

    async fn try_fetch_chain_config(
        &self,
        retry: usize,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        // Try fetching the chain config on the local providers first
        let local_result = self
            .on_local_providers(move |provider| async move {
                provider.try_fetch_chain_config(retry, commitment).await
            })
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones
        self.on_remote_providers(move |provider| async move {
            provider.try_fetch_chain_config(retry, commitment).await
        })
        .await
    }

    async fn try_fetch_reward_merkle_tree_v2(
        &self,
        retry: usize,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV2,
        accounts: Arc<Vec<RewardAccountV2>>,
    ) -> anyhow::Result<PermittedRewardMerkleTreeV2> {
        let local_result = self
            .on_local_providers(clone! {(accounts) move |provider| {
                clone! {(accounts) async move {
                    provider
                        .try_fetch_reward_merkle_tree_v2(
                            retry,
                            height,
                            view,
                            reward_merkle_tree_root,
                            accounts,
                        )
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones
        self.on_remote_providers(clone! {(accounts) move |provider| {
            clone!{(accounts) async move {
                provider
                .try_fetch_reward_merkle_tree_v2(
                    retry,
                    height,
                    view,
                    reward_merkle_tree_root,
                    accounts
                ).await
            }}
        }})
        .await
    }

    async fn try_fetch_reward_accounts_v1(
        &self,
        retry: usize,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV1,
        accounts: &[RewardAccountV1],
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        // Try to get the accounts on local providers first
        let accounts_vec = accounts.to_vec();
        let local_result = self
            .on_local_providers(clone! {(instance, accounts_vec) move |provider| {
                clone! {(instance, accounts_vec) async move {
                    provider
                        .try_fetch_reward_accounts_v1(
                            retry,
                            &instance,
                            height,
                            view,
                            reward_merkle_tree_root,
                            &accounts_vec,
                        )
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones
        self.on_remote_providers(clone! {(instance, accounts_vec) move |provider| {
            clone!{(instance, accounts_vec) async move {
                provider
                .try_fetch_reward_accounts_v1(
                    retry,
                    &instance,
                    height,
                    view,
                    reward_merkle_tree_root,
                    &accounts_vec,
                ).await
            }}
        }})
        .await
    }

    async fn try_fetch_state_cert(
        &self,
        retry: usize,
        epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        // Try fetching the state cert on the local providers first
        let local_result = self
            .on_local_providers(move |provider| async move {
                provider.try_fetch_state_cert(retry, epoch).await
            })
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones
        self.on_remote_providers(move |provider| async move {
            provider.try_fetch_state_cert(retry, epoch).await
        })
        .await
    }

    fn backoff(&self) -> &BackoffParams {
        &self.backoff
    }

    fn name(&self) -> String {
        format!(
            "[{}]",
            self.providers
                .lock()
                .iter()
                .map(|p| p.name())
                .collect::<Vec<_>>()
                .join(", ")
        )
    }

    async fn fetch_accounts(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        fee_merkle_tree_root: FeeMerkleCommitment,
        accounts: Vec<FeeAccount>,
    ) -> anyhow::Result<Vec<FeeAccountProof>> {
        // Try to get the accounts on local providers first
        let accounts_vec = accounts.to_vec();
        let local_result = self
            .on_local_providers(clone! {(instance, accounts_vec) move |provider| {
                clone! {(instance, accounts_vec) async move {
                    provider
                        .try_fetch_accounts(
                            0,
                            &instance,
                            height,
                            view,
                            fee_merkle_tree_root,
                            &accounts_vec,
                        )
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones (with retry)
        self.on_remote_providers(clone! {(instance, accounts_vec) move |provider| {
            clone!{(instance, accounts_vec) async move {
                provider
                .fetch_accounts(
                    &instance,
                    height,
                    view,
                    fee_merkle_tree_root,
                    accounts_vec,
                ).await
            }}
        }})
        .await
    }

    async fn fetch_leaf(
        &self,
        coordinator: EpochMembershipCoordinator<SeqTypes>,
        height: u64,
    ) -> anyhow::Result<Leaf2> {
        // Try fetching the leaf on the local providers first
        let local_result = self
            .on_local_providers(clone! {(coordinator) move |provider| {
                clone!{(coordinator) async move {
                    provider
                        .try_fetch_leaf(0, coordinator, height)
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones (with retry)
        self.on_remote_providers(clone! {(coordinator) move |provider| {
        clone!{(coordinator) async move {
            provider
                .fetch_leaf(coordinator, height)
                .await
        }}
        }})
        .await
    }

    async fn fetch_chain_config(
        &self,
        commitment: Commitment<ChainConfig>,
    ) -> anyhow::Result<ChainConfig> {
        // Try fetching the chain config on the local providers first
        let local_result = self
            .on_local_providers(move |provider| async move {
                provider.try_fetch_chain_config(0, commitment).await
            })
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones (with retry)
        self.on_remote_providers(move |provider| async move {
            provider.fetch_chain_config(commitment).await
        })
        .await
    }

    async fn fetch_reward_accounts_v1(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        reward_merkle_tree_root: RewardMerkleCommitmentV1,
        accounts: Vec<RewardAccountV1>,
    ) -> anyhow::Result<Vec<RewardAccountProofV1>> {
        // Try to get the accounts on local providers first
        let accounts_vec = accounts.to_vec();
        let local_result = self
            .on_local_providers(clone! {(instance, accounts_vec) move |provider| {
                clone! {(instance, accounts_vec) async move {
                    provider
                        .try_fetch_reward_accounts_v1(
                            0,
                            &instance,
                            height,
                            view,
                            reward_merkle_tree_root,
                            &accounts_vec,
                        )
                        .await
                }}
            }})
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones (with retry)
        self.on_remote_providers(clone! {(instance, accounts_vec) move |provider| {
            clone!{(instance, accounts_vec) async move {
                provider
                .fetch_reward_accounts_v1(
                    &instance,
                    height,
                    view,
                    reward_merkle_tree_root,
                    accounts_vec,
                ).await
            }}
        }})
        .await
    }

    async fn fetch_state_cert(
        &self,
        epoch: u64,
    ) -> anyhow::Result<LightClientStateUpdateCertificateV2<SeqTypes>> {
        let local_result = self
            .on_local_providers(move |provider| async move {
                provider.try_fetch_state_cert(0, epoch).await
            })
            .await;

        // Check if we were successful locally
        match &local_result {
            Ok(_) => return local_result,
            Err(err) => tracing::debug!("{err:#}"),
        }

        // If that fails, try the remote ones (with retry)
        self.on_remote_providers(
            move |provider| async move { provider.fetch_state_cert(epoch).await },
        )
        .await
    }

    async fn remember_blocks_merkle_tree(
        &self,
        instance: &NodeState,
        height: u64,
        view: ViewNumber,
        mt: &mut BlockMerkleTree,
    ) -> anyhow::Result<()> {
        // Try to remember the blocks merkle tree on local providers first
        let local_result = self
            .on_local_providers(clone! {(mt, instance) move |provider| {
                let mut mt = mt.clone();
                clone! {(instance) async move {
                    // Perform the call
                    provider
                        .try_remember_blocks_merkle_tree(
                            0,
                            &instance,
                            height,
                            view,
                            &mut mt,
                        )
                        .await?;

                    // Return the merkle tree so we can modify it
                    Ok(mt)
                }}
            }})
            .await;

        // Check if we were successful locally
        if let Ok(modified_mt) = local_result {
            // Set the merkle tree to the one with the
            // successful call
            *mt = modified_mt;

            return Ok(());
        }

        // If that fails, try the remote ones (with retry)
        let remote_result = self
            .on_remote_providers(clone! {(mt, instance) move |provider| {
                let mut mt = mt.clone();
                clone!{(instance) async move {
                    // Perform the call
                    provider
                    .remember_blocks_merkle_tree(
                        &instance,
                        height,
                        view,
                        &mut mt,
                    )
                    .await?;

                    // Return the merkle tree
                    Ok(mt)
                }}
            }})
            .await?;

        // Update the original, local merkle tree
        *mt = remote_result;

        Ok(())
    }

    fn is_local(&self) -> bool {
        self.providers.lock().iter().all(|p| p.is_local())
    }
}

/// Add accounts to the in-memory consensus state.
/// We use this during catchup after receiving verified accounts.
pub async fn add_fee_accounts_to_state<I: hotshot::traits::NodeImplementation<SeqTypes>>(
    consensus_handle: &ConsensusHandle<SeqTypes, I>,
    view: &ViewNumber,
    accounts: &[FeeAccount],
    tree: &FeeMerkleTree,
    leaf: Leaf2,
) -> anyhow::Result<()>
where
    I::Storage: NewProtocolStorage<SeqTypes>,
{
    let (existing_state, delta) = consensus_handle.state_and_delta(*view).await;
    let (state, delta) = match existing_state {
        Some(existing) => {
            let mut state = (*existing).clone();
            for account in accounts {
                if let Some((proof, _)) = FeeAccountProof::prove(tree, (*account).into()) {
                    if let Err(err) = proof.remember(&mut state.fee_merkle_tree) {
                        tracing::warn!(
                            ?view,
                            %account,
                            "cannot update fetched account state: {err:#}"
                        );
                    }
                } else {
                    tracing::warn!(?view, %account, "cannot update fetched account state because account is not in the merkle tree");
                };
            }
            (Arc::new(state), delta)
        },
        None => {
            // 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();
            (Arc::new(state), None)
        },
    };

    consensus_handle
        .update_leaf(leaf, state, delta)
        .await
        .with_context(|| "failed to update leaf")?;

    Ok(())
}

/// Add accounts to the in-memory consensus state.
/// We use this during catchup after receiving verified accounts.
pub async fn add_v2_reward_accounts_to_state<I: hotshot::traits::NodeImplementation<SeqTypes>>(
    consensus_handle: &ConsensusHandle<SeqTypes, I>,
    view: &ViewNumber,
    accounts: &[RewardAccountV2],
    tree: &RewardMerkleTreeV2,
    leaf: Leaf2,
) -> anyhow::Result<()>
where
    I::Storage: NewProtocolStorage<SeqTypes>,
{
    let (existing_state, delta) = consensus_handle.state_and_delta(*view).await;
    let (state, delta) = match existing_state {
        Some(existing) => {
            let mut state = (*existing).clone();
            for account in accounts {
                if let Some((proof, _)) = RewardAccountProofV2::prove(tree, (*account).into()) {
                    if let Err(err) = proof.remember(&mut state.reward_merkle_tree_v2) {
                        tracing::warn!(
                            ?view,
                            %account,
                            "cannot update fetched account state: {err:#}"
                        );
                    }
                } else {
                    tracing::warn!(?view, %account, "cannot update fetched account state because account is not in the merkle tree");
                };
            }
            (Arc::new(state), delta)
        },
        None => {
            // 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.reward_merkle_tree_v2 = tree.clone();
            (Arc::new(state), None)
        },
    };

    consensus_handle
        .update_leaf(leaf, state, delta)
        .await
        .with_context(|| "failed to update leaf")?;

    Ok(())
}

/// Add accounts to the in-memory consensus state.
/// We use this during catchup after receiving verified accounts.
pub async fn add_v1_reward_accounts_to_state<I: hotshot::traits::NodeImplementation<SeqTypes>>(
    consensus_handle: &ConsensusHandle<SeqTypes, I>,
    view: &ViewNumber,
    accounts: &[RewardAccountV1],
    tree: &RewardMerkleTreeV1,
    leaf: Leaf2,
) -> anyhow::Result<()>
where
    I::Storage: NewProtocolStorage<SeqTypes>,
{
    let (existing_state, delta) = consensus_handle.state_and_delta(*view).await;
    let (state, delta) = match existing_state {
        Some(existing) => {
            let mut state = (*existing).clone();
            for account in accounts {
                if let Some((proof, _)) = RewardAccountProofV1::prove(tree, (*account).into()) {
                    if let Err(err) = proof.remember(&mut state.reward_merkle_tree_v1) {
                        tracing::warn!(
                            ?view,
                            %account,
                            "cannot update fetched account state: {err:#}"
                        );
                    }
                } else {
                    tracing::warn!(?view, %account, "cannot update fetched account state because account is not in the merkle tree");
                };
            }
            (Arc::new(state), delta)
        },
        None => {
            // 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.reward_merkle_tree_v1 = tree.clone();
            (Arc::new(state), None)
        },
    };

    consensus_handle
        .update_leaf(leaf, state, delta)
        .await
        .with_context(|| "failed to update leaf")?;

    Ok(())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_peer_priority() {
        let good_peer = PeerScore {
            requests: 1000,
            failures: 2,
        };
        let bad_peer = PeerScore {
            requests: 10,
            failures: 1,
        };
        assert!(good_peer > bad_peer);

        let mut peers: PriorityQueue<_, _> = [(0, good_peer), (1, bad_peer)].into_iter().collect();
        assert_eq!(peers.pop(), Some((0, good_peer)));
        assert_eq!(peers.pop(), Some((1, bad_peer)));
    }
}