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