Skip to main content

espresso_node/request_response/
recipient_source.rs

1use std::sync::Arc;
2
3use anyhow::{Context, Result};
4use async_trait::async_trait;
5use espresso_types::{PubKey, SeqTypes};
6use hotshot::traits::NodeImplementation;
7use hotshot_new_protocol::storage::NewProtocolStorage;
8use hotshot_types::{data::EpochNumber, epoch_membership::EpochMembershipCoordinator};
9use request_response::recipient_source::RecipientSource as RecipientSourceTrait;
10use tracing::warn;
11
12use super::request::Request;
13use crate::consensus_handle::ConsensusHandle;
14
15#[derive(Clone)]
16pub struct RecipientSource<I: NodeImplementation<SeqTypes>> {
17    /// The consensus adapter handle
18    pub consensus_handle: Arc<ConsensusHandle<SeqTypes, I>>,
19    /// A copy of the membership coordinator
20    pub memberships: EpochMembershipCoordinator<SeqTypes>,
21    /// The public key of the node
22    pub public_key: PubKey,
23}
24
25/// Implement the RecipientSourceTrait, which allows the request-response protocol to derive the
26/// intended recipients for a given request
27#[async_trait]
28impl<I: NodeImplementation<SeqTypes>> RecipientSourceTrait<Request, PubKey> for RecipientSource<I>
29where
30    I::Storage: NewProtocolStorage<SeqTypes>,
31{
32    async fn get_expected_responders(&self, _request: &Request) -> Result<Vec<PubKey>> {
33        // Get the current epoch number
34        let epoch_number = self
35            .consensus_handle
36            .current_epoch()
37            .await
38            .unwrap_or(EpochNumber::genesis());
39
40        // Attempt to get the membership for the current epoch
41        let membership = match self.memberships.stake_table_for_epoch(Some(epoch_number)) {
42            Ok(membership) => membership,
43            Err(e) => {
44                warn!(
45                    "Failed to get membership for epoch {}: {e:#}. Failing over to previous epoch",
46                    epoch_number
47                );
48                let prev_epoch = epoch_number.saturating_sub(1);
49                self.memberships
50                    .stake_table_for_epoch(Some(EpochNumber::new(prev_epoch)))
51                    .with_context(|| "failed to get stake table for epoch")?
52            },
53        };
54
55        // Sum all participants in the membership
56        Ok(membership
57            .stake_table()
58            .map(|entry| entry.stake_table_entry.stake_key)
59            .filter(|key| *key != self.public_key)
60            .collect())
61    }
62}