Split up add_epoch_root to allow stake table storage to happen outsid…

crates/hotshot/hotshot/src/lib.rs

@@ -1211,16 +1211,18 @@ async fn load_start_epoch_info<TYPES: NodeType>(
     for epoch_info in start_epoch_info {
         if let Some(block_header) = &epoch_info.block_header {
             tracing::info!("Calling add_epoch_root for epoch {:?}", epoch_info.epoch);
-            let write_callback = {
+            let add_epoch_root_worker = {
                 let membership_reader = membership.read().await;
-                membership_reader
-                    .add_epoch_root(epoch_info.epoch, block_header.clone())
-                    .await
+                membership_reader.add_epoch_root(epoch_info.epoch, block_header.clone())
             };
 
-            if let Some(write_callback) = write_callback {
-                let mut membership_writer = membership.write().await;
-                write_callback(&mut *membership_writer);
+            if let Some(worker) = add_epoch_root_worker {
+                let add_epoch_root_updater = worker().await;
+
+                if let Some(updater) = add_epoch_root_updater {
+                    let mut membership_writer = membership.write().await;
+                    updater(&mut *membership_writer);
+                }
             }
         }
     }

crates/hotshot/hotshot/src/traits/election/dummy_catchup_membership.rs

@@ -1,4 +1,4 @@
-use std::{collections::HashSet, sync::Arc, time::Duration};
+use std::{collections::HashSet, pin::Pin, sync::Arc, time::Duration};
 
 use alloy::primitives::U256;
 use anyhow::Ok;
@@ -199,15 +199,27 @@ where
         self.inner.set_first_epoch(epoch, initial_drb_result);
     }
 
-    #[allow(refining_impl_trait)]
-    async fn add_epoch_root(
+    fn add_epoch_root(
         &self,
         epoch: TYPES::Epoch,
         _block_header: TYPES::BlockHeader,
-    ) -> Option<Box<dyn FnOnce(&mut Self) + Send>> {
-        Some(Box::new(move |mem: &mut Self| {
-            tracing::error!("Adding epoch root for {epoch}");
-            mem.epochs.insert(epoch);
+    ) -> Option<
+        Box<
+            dyn FnOnce() -> Pin<
+                    Box<
+                        dyn std::future::Future<Output = Option<Box<dyn FnOnce(&mut Self) + Send>>>
+                            + Send,
+                    >,
+                > + Send,
+        >,
+    > {
+        Some(Box::new(move || {
+            Box::pin(async move {
+                Some(Box::new(move |mem: &mut Self| {
+                    tracing::error!("Adding epoch root for {epoch}");
+                    mem.epochs.insert(epoch);
+                }) as Box<dyn FnOnce(&mut Self) + Send>)
+            })
         }))
     }
 

crates/hotshot/task-impls/src/helpers.rs

@@ -222,16 +222,20 @@ async fn decide_epoch_root<TYPES: NodeType, I: NodeImplementation<TYPES>>(
         tracing::info!("Time taken to store epoch root: {:?}", start.elapsed());
 
         start = Instant::now();
-        let write_callback = {
+
+        let add_epoch_root_worker = {
             tracing::debug!("Calling add_epoch_root for epoch {next_epoch_number}");
             let membership_reader = membership.read().await;
-            membership_reader
-                .add_epoch_root(next_epoch_number, decided_leaf.block_header().clone())
-                .await
+            membership_reader.add_epoch_root(next_epoch_number, decided_leaf.block_header().clone())
         };
-        if let Some(write_callback) = write_callback {
-            let mut membership_writer = membership.write().await;
-            write_callback(&mut *membership_writer);
+
+        if let Some(worker) = add_epoch_root_worker {
+            let add_epoch_root_updater = worker().await;
+
+            if let Some(updater) = add_epoch_root_updater {
+                let mut membership_writer = membership.write().await;
+                updater(&mut *membership_writer);
+            }
         }
         tracing::info!("Time taken to add epoch root: {:?}", start.elapsed());
 

crates/hotshot/types/src/epoch_membership.rs

@@ -373,16 +373,18 @@ where
             ));
         };
 
-        let add_epoch_root_updater = {
-            let membership_read = self.membership.read().await;
-            membership_read
-                .add_epoch_root(epoch, root_leaf.block_header().clone())
-                .await
+        let add_epoch_root_worker = {
+            let membership_reader = self.membership.read().await;
+            membership_reader.add_epoch_root(epoch, root_leaf.block_header().clone())
         };
 
-        if let Some(updater) = add_epoch_root_updater {
-            let mut membership_write = self.membership.write().await;
-            updater(&mut *(membership_write));
+        if let Some(worker) = add_epoch_root_worker {
+            let add_epoch_root_updater = worker().await;
+
+            if let Some(updater) = add_epoch_root_updater {
+                let mut membership_writer = self.membership.write().await;
+                updater(&mut *membership_writer);
+            };
         };
 
         Ok(root_leaf)

crates/hotshot/types/src/traits/election.rs

@@ -5,7 +5,7 @@
 // along with the HotShot repository. If not, see <https://mit-license.org/>.
 
 //! The election trait, used to decide which node is the leader and determine if a vote is valid.
-use std::{collections::BTreeSet, fmt::Debug, sync::Arc};
+use std::{collections::BTreeSet, fmt::Debug, pin::Pin, sync::Arc};
 
 use alloy::primitives::U256;
 use async_lock::RwLock;
@@ -164,17 +164,44 @@ pub trait Membership<TYPES: NodeType>: Debug + Send + Sync {
 
     #[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.
+    /// This function results in a chain of callbacks being returned. It is invoked
+    /// synchronously under a read lock. Return an Option-wrapped callback to run an
+    /// asynchronous handler to execute. This handler returns another Option-wrapped
+    /// callback, which will be executed synchronously under a write lock to save
+    /// results back to the membership.
     ///
-    /// #3967 REVIEW NOTE: this is only called if epoch is Some. Is there any reason to do otherwise?
+    /// Example usage:
+    /// ```
+    ///    // Do some synchronous work first inside of the read lock
+    ///    Some(Box::new(move || {
+    ///        Box::pin(async move {
+    ///            // We're now in the async context between the two locks.
+    ///            // Do some async work here, like storing the stake table
+    ///
+    ///            // Return a callback to be executed under the write lock
+    ///            Some(Box::new(move |committee: &mut Self| {
+    ///                // Do some synchronous work here, updating the membership
+    ///            }) as Box<dyn FnOnce(&mut Self) + Send>)
+    ///        })
+    ///    }))
+    /// ```
     fn add_epoch_root(
         &self,
         _epoch: TYPES::Epoch,
         _block_header: TYPES::BlockHeader,
-    ) -> impl std::future::Future<Output = Option<Box<dyn FnOnce(&mut Self) + Send>>> + Send {
-        async { None }
+    ) -> Option<
+        Box<
+            dyn FnOnce() -> Pin<
+                    Box<
+                        dyn std::future::Future<Output = Option<Box<dyn FnOnce(&mut Self) + Send>>>
+                            + Send,
+                    >,
+                > + Send,
+        >,
+    > {
+        None
     }
+    //) -> impl std::future::Future<Output = Option<Box<dyn FnOnce(&mut Self) + Send>>> + Send {
 
     /// Called to notify the Membership when a new DRB result has been calculated.
     /// Observes the same semantics as add_epoch_root

types/src/v0/impls/stake_table.rs

@@ -2,6 +2,7 @@ use std::{
     cmp::{max, min},
     collections::{BTreeMap, BTreeSet, HashMap, HashSet},
     future::Future,
+    pin::Pin,
     sync::Arc,
 };
 
@@ -1221,12 +1222,20 @@ impl Membership<SeqTypes> for EpochCommittees {
         max(higher_threshold, normal_threshold)
     }
 
-    #[allow(refining_impl_trait)]
-    async fn add_epoch_root(
+    fn add_epoch_root(
         &self,
         epoch: Epoch,
         block_header: Header,
-    ) -> Option<Box<dyn FnOnce(&mut Self) + Send>> {
+    ) -> Option<
+        Box<
+            dyn FnOnce() -> Pin<
+                    Box<
+                        dyn std::future::Future<Output = Option<Box<dyn FnOnce(&mut Self) + Send>>>
+                            + Send,
+                    >,
+                > + Send,
+        >,
+    > {
         if self.state.contains_key(&epoch) {
             tracing::info!(
                 "We already have a the stake table for epoch {}. Skipping L1 fetching.",
@@ -1235,20 +1244,26 @@ impl Membership<SeqTypes> for EpochCommittees {
             return None;
         }
 
-        let stake_tables = self.fetcher.fetch(epoch, block_header).await?;
-        // Store stake table in persistence
-        {
-            let persistence_lock = self.fetcher.persistence.lock().await;
-            if let Err(e) = persistence_lock
-                .store_stake(epoch, stake_tables.clone())
-                .await
-            {
-                tracing::error!(?e, "`add_epoch_root`, error storing stake table");
-            }
-        }
+        let fetcher = Arc::clone(&self.fetcher);
+
+        Some(Box::new(move || {
+            Box::pin(async move {
+                let stake_tables = fetcher.fetch(epoch, block_header).await?;
+                // Store stake table in persistence
+                {
+                    let persistence_lock = fetcher.persistence.lock().await;
+                    if let Err(e) = persistence_lock
+                        .store_stake(epoch, stake_tables.clone())
+                        .await
+                    {
+                        tracing::error!(?e, "`add_epoch_root`, error storing stake table");
+                    }
+                }
 
-        Some(Box::new(move |committee: &mut Self| {
-            committee.update_stake_table(epoch, stake_tables);
+                Some(Box::new(move |committee: &mut Self| {
+                    committee.update_stake_table(epoch, stake_tables);
+                }) as Box<dyn FnOnce(&mut Self) + Send>)
+            })
         }))
     }