3966 add_epoch_root and sync_l1 on Membership

crates/task-impls/src/helpers.rs

@@ -163,6 +163,56 @@ pub(crate) async fn fetch_proposal<TYPES: NodeType, V: Versions>(
     Ok((leaf, view))
 }
 
+/// Handles calling add_epoch_root and sync_l1 on Membership if necessary.
+async fn decide_from_proposal_add_epoch_root<TYPES: NodeType>(
+    proposal: &QuorumProposal2<TYPES>,
+    leaf_views: &[LeafInfo<TYPES>],
+    epoch_height: u64,
+    membership: &Arc<RwLock<TYPES::Membership>>,
+) {
+    if leaf_views.is_empty() {
+        return;
+    }
+
+    let decided_block_number = leaf_views
+        .last()
+        .unwrap()
+        .leaf
+        .block_header()
+        .block_number();
+
+    // Skip if this is not the expected block.
+    if epoch_height != 0 && (decided_block_number + 3) % epoch_height == 0 {
+        let next_epoch_number =
+            TYPES::Epoch::new(epoch_from_block_number(decided_block_number, epoch_height) + 1);
+
+        let write_callback = {
+            let membership_reader = membership.read().await;
+            membership_reader
+                .add_epoch_root(next_epoch_number, proposal.block_header.clone())
+                .await
+        };
+
+        if let Some(write_callback) = write_callback {
+            let mut membership_writer = membership.write().await;
+            write_callback(&mut *membership_writer);
+        } else {
+            // If we didn't get a write callback out of add_epoch_root, then don't bother locking and calling sync_l1
+            return;
+        }
+
+        let write_callback = {
+            let membership_reader = membership.read().await;
+            membership_reader.sync_l1().await
+        };
+
+        if let Some(write_callback) = write_callback {
+            let mut membership_writer = membership.write().await;
+            write_callback(&mut *membership_writer);
+        }
+    }
+}
+
 /// Helper type to give names and to the output values of the leaf chain traversal operation.
 #[derive(Debug)]
 pub struct LeafChainTraversalOutcome<TYPES: NodeType> {
@@ -202,7 +252,7 @@ impl<TYPES: NodeType + Default> Default for LeafChainTraversalOutcome<TYPES> {
     }
 }
 
-/// calculate the new decided leaf chain based on the rules of hostuff 2
+/// calculate the new decided leaf chain based on the rules of HotStuff 2
 ///
 /// # Panics
 /// Can't actually panic
@@ -211,6 +261,8 @@ pub async fn decide_from_proposal_2<TYPES: NodeType>(
     consensus: OuterConsensus<TYPES>,
     existing_upgrade_cert: Arc<RwLock<Option<UpgradeCertificate<TYPES>>>>,
     public_key: &TYPES::SignatureKey,
+    with_epochs: bool,
+    membership: &Arc<RwLock<TYPES::Membership>>,
 ) -> LeafChainTraversalOutcome<TYPES> {
     let mut res = LeafChainTraversalOutcome::default();
     let consensus_reader = consensus.read().await;
@@ -282,6 +334,14 @@ pub async fn decide_from_proposal_2<TYPES: NodeType>(
         res.included_txns = Some(txns);
     }
 
+    if with_epochs && res.new_decided_view_number.is_some() {
+        let epoch_height = consensus_reader.epoch_height;
+        drop(consensus_reader);
+
+        decide_from_proposal_add_epoch_root(proposal, &res.leaf_views, epoch_height, membership)
+            .await;
+    }
+
     res
 }
 
@@ -317,6 +377,8 @@ pub async fn decide_from_proposal<TYPES: NodeType>(
     consensus: OuterConsensus<TYPES>,
     existing_upgrade_cert: Arc<RwLock<Option<UpgradeCertificate<TYPES>>>>,
     public_key: &TYPES::SignatureKey,
+    with_epochs: bool,
+    membership: &Arc<RwLock<TYPES::Membership>>,
 ) -> LeafChainTraversalOutcome<TYPES> {
     let consensus_reader = consensus.read().await;
     let existing_upgrade_cert_reader = existing_upgrade_cert.read().await;
@@ -428,6 +490,14 @@ pub async fn decide_from_proposal<TYPES: NodeType>(
         tracing::debug!("Leaf ascension failed; error={e}");
     }
 
+    if with_epochs && res.new_decided_view_number.is_some() {
+        let epoch_height = consensus_reader.epoch_height;
+        drop(consensus_reader);
+
+        decide_from_proposal_add_epoch_root(proposal, &res.leaf_views, epoch_height, membership)
+            .await;
+    }
+
     res
 }
 

crates/task-impls/src/quorum_vote/handlers.rs

@@ -167,6 +167,8 @@ pub(crate) async fn handle_quorum_proposal_validated<
             OuterConsensus::new(Arc::clone(&task_state.consensus.inner_consensus)),
             Arc::clone(&task_state.upgrade_lock.decided_upgrade_certificate),
             &task_state.public_key,
+            version >= V::Epochs::VERSION,
+            &task_state.membership,
         )
         .await
     } else {
@@ -175,6 +177,8 @@ pub(crate) async fn handle_quorum_proposal_validated<
             OuterConsensus::new(Arc::clone(&task_state.consensus.inner_consensus)),
             Arc::clone(&task_state.upgrade_lock.decided_upgrade_certificate),
             &task_state.public_key,
+            version >= V::Epochs::VERSION,
+            &task_state.membership,
         )
         .await
     };

crates/types/src/traits/election.rs

@@ -7,11 +7,13 @@
 //! The election trait, used to decide which node is the leader and determine if a vote is valid.
 use std::{collections::BTreeSet, fmt::Debug, num::NonZeroU64};
 
+use async_trait::async_trait;
 use utils::anytrace::Result;
 
 use super::node_implementation::NodeType;
 use crate::{traits::signature_key::SignatureKey, PeerConfig};
 
+#[async_trait]
 /// A protocol for determining membership in and participating in a committee.
 pub trait Membership<TYPES: NodeType>: Debug + Send + Sync {
     /// The error type returned by methods like `lookup_leader`.
@@ -125,4 +127,23 @@ pub trait Membership<TYPES: NodeType>: Debug + Send + Sync {
 
     /// Returns the threshold required to upgrade the network protocol
     fn upgrade_threshold(&self, epoch: TYPES::Epoch) -> NonZeroU64;
+
+    #[allow(clippy::type_complexity)]
+    /// Handles notifications that a new epoch root has been created
+    /// Is called under a read lock to the Membership. Return a callback
+    /// with Some to have that callback invoked under a write lock.
+    async fn add_epoch_root(
+        &self,
+        _epoch: TYPES::Epoch,
+        _block_header: TYPES::BlockHeader,
+    ) -> Option<Box<dyn FnOnce(&mut Self) + Send>> {
+        None
+    }
+
+    #[allow(clippy::type_complexity)]
+    /// Called after add_epoch_root runs and any callback has been invoked.
+    /// Causes a read lock to be reacquired for this functionality.
+    async fn sync_l1(&self) -> Option<Box<dyn FnOnce(&mut Self) + Send>> {
+        None
+    }
 }