WIP even more synchronization

MystenLabs · Oct 15, 2024 · 904d92d · 904d92d
1 parent 62e731c
commit 904d92d
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 8 deletions.
diff --git a/crates/sui-core/src/checkpoints/mod.rs b/crates/sui-core/src/checkpoints/mod.rs
@@ -819,9 +819,12 @@ impl CheckpointBuilder {
         }
     }
 
-    async fn run(mut self, startup_wait: tokio::sync::oneshot::Receiver<()>) {
+    async fn run(
+        mut self,
+        startup_wait: tokio::sync::oneshot::Receiver<tokio::sync::oneshot::Sender<()>>,
+    ) {
         info!("CheckpointBuilder waiting for startup signal");
-        startup_wait.await.ok();
+        let mut notify = Some(startup_wait.await.unwrap());
         info!("Starting CheckpointBuilder");
         loop {
             // Check whether an exit signal has been received, if so we break the loop.
@@ -835,6 +838,10 @@ impl CheckpointBuilder {
 
             self.maybe_build_checkpoints().await;
 
+            if let Some(notify) = notify.take() {
+                notify.send(()).unwrap();
+            }
+
             match select(self.exit.changed().boxed(), self.notify.notified().boxed()).await {
                 Either::Left(_) => {
                     // break loop on exit signal
@@ -2146,8 +2153,8 @@ impl CheckpointService {
         max_checkpoint_size_bytes: usize,
     ) -> (
         Arc<Self>,
-        watch::Sender<()>,                /* The exit sender */
-        tokio::sync::oneshot::Sender<()>, /* builder start-up sender */
+        watch::Sender<()>, /* The exit sender */
+        tokio::sync::oneshot::Sender<tokio::sync::oneshot::Sender<()>>, /* builder start-up sender */
     ) {
         info!(
             "Starting checkpoint service with {max_transactions_per_checkpoint} max_transactions_per_checkpoint and {max_checkpoint_size_bytes} max_checkpoint_size_bytes"
@@ -2456,7 +2463,9 @@ mod tests {
             3,
             100_000,
         );
-        startup.send(()).unwrap();
+        let (tx, rx) = tokio::sync::oneshot::channel();
+        startup.send(tx).unwrap();
+        rx.await.unwrap();
 
         checkpoint_service
             .write_and_notify_checkpoint_for_testing(&epoch_store, p(0, vec![4], 0))

diff --git a/crates/sui-core/src/unit_tests/mysticeti_manager_tests.rs b/crates/sui-core/src/unit_tests/mysticeti_manager_tests.rs
@@ -45,7 +45,11 @@ pub fn checkpoint_service_for_testing(state: Arc<AuthorityState>) -> Arc<Checkpo
         3,
         100_000,
     );
-    startup.send(()).unwrap();
+    let (tx, rx) = tokio::sync::oneshot::channel();
+    startup.send(tx).unwrap();
+    tokio::task::spawn(async move {
+        rx.await.unwrap();
+    });
     checkpoint_service
 }
 

diff --git a/crates/sui-node/src/lib.rs b/crates/sui-node/src/lib.rs
@@ -1333,10 +1333,13 @@ impl SuiNode {
             .await;
 
         info!("consensus manager started");
+        let (tx, rx) = tokio::sync::oneshot::channel();
         startup_sender
-            .send(())
+            .send(tx)
             .expect("Failed to send startup signal");
 
+        rx.await.expect("Failed to receive startup signal");
+
         if epoch_store.authenticator_state_enabled() {
             Self::start_jwk_updater(
                 config,
@@ -1371,7 +1374,7 @@ impl SuiNode {
     ) -> (
         Arc<CheckpointService>,
         watch::Sender<()>,
-        tokio::sync::oneshot::Sender<()>,
+        tokio::sync::oneshot::Sender<tokio::sync::oneshot::Sender<()>>,
     ) {
         let epoch_start_timestamp_ms = epoch_store.epoch_start_state().epoch_start_timestamp_ms();
         let epoch_duration_ms = epoch_store.epoch_start_state().epoch_duration_ms();