hotshot_new_protocol/
epoch.rs

1use std::{
2    collections::{BTreeMap, BTreeSet},
3    mem::swap,
4};
5
6use hotshot_types::{
7    data::{BlockNumber, EpochNumber, Leaf2},
8    drb::DrbResult,
9    epoch_membership::EpochMembershipCoordinator,
10    traits::{block_contents::BlockHeader, node_implementation::NodeType},
11    utils::{is_epoch_root, is_transition_block},
12};
13use hotshot_utils::anytrace;
14use tokio::task::{AbortHandle, JoinSet};
15use tracing::error;
16
17pub enum EpochRootResult {
18    DrbResult(EpochNumber, DrbResult),
19}
20
21/// Epoch + error for the Err path so retries can re-kick the task.
22pub struct EpochFailure {
23    pub epoch: EpochNumber,
24    pub error: EpochManagerError,
25}
26
27/// Output of a spawned DRB task.  The leading `EpochNumber` tags the task
28/// with the epoch it was working on so [`EpochManager::next`] can update
29/// dedup state correctly even on the Err path.
30type TaskOutput = (EpochNumber, Result<EpochRootResult, EpochManagerError>);
31
32/// Manager for Epoch specific rules and actions.
33///
34/// Delegates all catchup (stake-table walk-back, epoch root fetch, DRB
35/// compute) to [`EpochMembershipCoordinator::membership_for_epoch`].  This
36/// manager exists only to issue the request and dedup concurrent callers.
37pub struct EpochManager<T: NodeType> {
38    epoch_height: BlockNumber,
39    membership_coordinator: EpochMembershipCoordinator<T>,
40    tasks: JoinSet<TaskOutput>,
41    handles: BTreeMap<EpochNumber, Vec<AbortHandle>>,
42    /// Epochs for which a `request_drb_result` task is currently in flight.
43    /// Prevents duplicate fetch/compute tasks while the first is running.
44    pending_drb_requests: BTreeSet<EpochNumber>,
45    /// Epochs whose DRB has already been computed and added to membership.
46    /// Subsequent `request_drb_result` calls for these epochs are no-ops.
47    completed_drb_requests: BTreeSet<EpochNumber>,
48}
49
50impl<T: NodeType> EpochManager<T> {
51    pub fn new<B>(epoch_height: B, membership_coordinator: EpochMembershipCoordinator<T>) -> Self
52    where
53        B: Into<BlockNumber>,
54    {
55        Self {
56            epoch_height: epoch_height.into(),
57            membership_coordinator,
58            tasks: JoinSet::new(),
59            handles: BTreeMap::new(),
60            pending_drb_requests: BTreeSet::new(),
61            completed_drb_requests: BTreeSet::new(),
62        }
63    }
64
65    pub async fn next(&mut self) -> Option<Result<EpochRootResult, EpochFailure>> {
66        loop {
67            match self.tasks.join_next().await {
68                Some(Ok((epoch, result))) => {
69                    match result {
70                        Ok(root @ EpochRootResult::DrbResult(..)) => {
71                            self.pending_drb_requests.remove(&epoch);
72                            self.completed_drb_requests.insert(epoch);
73                            return Some(Ok(root));
74                        },
75                        Err(error) => {
76                            // Clear the guard so a subsequent call can retry.
77                            self.pending_drb_requests.remove(&epoch);
78                            return Some(Err(EpochFailure { epoch, error }));
79                        },
80                    }
81                },
82                Some(Err(err)) => {
83                    if !err.is_cancelled() {
84                        error!(%err, "epoch manager task panic")
85                    }
86                },
87                None => return None,
88            }
89        }
90    }
91
92    pub fn handle_leaf_decided(&mut self, leaf: Leaf2<T>) {
93        let block_number = leaf.block_header().block_number();
94
95        // At every epoch root, trigger DRB computation for the epoch that
96        // will use this root (epoch + 2).
97        if is_epoch_root(block_number, *self.epoch_height) {
98            let Some(epoch) = leaf.epoch(*self.epoch_height) else {
99                error!("Leaf has no epoch");
100                return;
101            };
102
103            let target_epoch = epoch + 2;
104            if self.completed_drb_requests.contains(&target_epoch)
105                || self.pending_drb_requests.contains(&target_epoch)
106            {
107                return;
108            }
109            self.pending_drb_requests.insert(target_epoch);
110
111            let membership_coordinator = self.membership_coordinator.clone();
112            let root_leaf = leaf.clone();
113            let handles = self.handles.entry(target_epoch).or_default();
114            handles.push(self.tasks.spawn(async move {
115                let result = async {
116                    membership_coordinator
117                        .add_epoch_root(root_leaf.block_header().clone())
118                        .await
119                        .map_err(|e| EpochManagerError::EpochRoot(anyhow::anyhow!("{e}")))?;
120
121                    // Compute the DRB from the decided root leaf.
122                    let drb = membership_coordinator
123                        .compute_drb_result(target_epoch, root_leaf)
124                        .await
125                        .map_err(EpochManagerError::DrbCompute)?;
126                    Ok(EpochRootResult::DrbResult(target_epoch, drb))
127                }
128                .await;
129                (target_epoch, result)
130            }));
131        }
132
133        // If this is the transition block of an epoch feed the DRB result to the coordinator.
134        if is_transition_block(block_number, *self.epoch_height)
135            && let Some(epoch) = leaf.epoch(*self.epoch_height)
136            && let Some(drb) = leaf.next_drb_result
137        {
138            let target_epoch = epoch + 1;
139            if !self.completed_drb_requests.contains(&target_epoch) {
140                self.membership_coordinator.supply_drb(target_epoch, drb);
141                self.completed_drb_requests.insert(target_epoch);
142                self.pending_drb_requests.remove(&target_epoch);
143            }
144        }
145    }
146
147    pub fn gc(&mut self, epoch: EpochNumber) {
148        let mut tmp = self.handles.split_off(&epoch);
149        swap(&mut tmp, &mut self.handles);
150        for handle in tmp.into_values().flatten() {
151            handle.abort();
152        }
153        // Drop tracking entries for epochs we no longer care about.  Keeps
154        // `completed_drb_requests` bounded while the protocol runs.
155        self.pending_drb_requests = self.pending_drb_requests.split_off(&epoch);
156        self.completed_drb_requests = self.completed_drb_requests.split_off(&epoch);
157    }
158
159    pub fn request_drb_result(&mut self, epoch: EpochNumber) {
160        // Already computed — caller can read the DRB from membership.
161        if self.completed_drb_requests.contains(&epoch) {
162            return;
163        }
164        // In-flight task will deliver the result; avoid spawning a duplicate.
165        if self.pending_drb_requests.contains(&epoch) {
166            return;
167        }
168        self.pending_drb_requests.insert(epoch);
169        let membership_coordinator = self.membership_coordinator.clone();
170        let handles = self.handles.entry(epoch).or_default();
171
172        handles.push(self.tasks.spawn(async move {
173            let result = async {
174                // Kick the membership coordinator.  If the stake table is
175                // already ready, this returns it immediately; otherwise it
176                // spawns a catchup task and returns a "catchup in progress"
177                // error.  Either way, `wait_for_catchup` resolves once the
178                // stake table + DRB are both in place.
179                let membership = match membership_coordinator.membership_for_epoch(Some(epoch)) {
180                    Ok(m) => m,
181                    Err(_) => membership_coordinator
182                        .wait_for_catchup(epoch)
183                        .await
184                        .map_err(EpochManagerError::DrbLookup)?,
185                };
186                let drb = membership
187                    .get_epoch_drb()
188                    .await
189                    .map_err(EpochManagerError::DrbLookup)?;
190                Ok(EpochRootResult::DrbResult(epoch, drb))
191            }
192            .await;
193            (epoch, result)
194        }));
195    }
196}
197
198#[derive(Debug, thiserror::Error)]
199pub enum EpochManagerError {
200    #[error("failed to add epoch root: {0}")]
201    EpochRoot(#[source] anyhow::Error),
202
203    #[error("failed to compute drb: {0}")]
204    DrbCompute(#[source] anytrace::Error),
205
206    #[error("failed to get drb: {0}")]
207    DrbLookup(#[source] anytrace::Error),
208}
hotshot_new_protocol/epoch.rs

hotshot_new_protocol/
epoch.rs
