Fix flaky tests for BufferedStorage

node/libs/storage/Cargo.toml

@@ -19,6 +19,7 @@ roles = { path = "../roles" }
 schema = { path = "../schema" }
 
 [dev-dependencies]
+assert_matches.workspace = true
 tempfile.workspace = true
 test-casing.workspace = true
 tokio.workspace = true

node/libs/storage/src/buffered.rs

@@ -5,6 +5,8 @@ use crate::{
     StorageError, StorageResult,
 };
 use async_trait::async_trait;
+#[cfg(test)]
+use concurrency::ctx::channel;
 use concurrency::{
     ctx,
     sync::{self, watch, Mutex},
@@ -117,13 +119,23 @@ impl BlockBuffer {
     }
 }
 
+/// Events emitted by [`BufferedStorage`].
+#[cfg(test)]
+#[derive(Debug)]
+pub(crate) enum BufferedStorageEvent {
+    /// Update was received from
+    UpdateReceived(BlockNumber),
+}
+
 /// [`BlockStore`] with an in-memory buffer for pending blocks.
 #[derive(Debug)]
 pub struct BufferedStorage<T> {
     inner: T,
     inner_subscriber: watch::Receiver<BlockNumber>,
     block_writes_sender: watch::Sender<BlockNumber>,
     buffer: Mutex<BlockBuffer>,
+    #[cfg(test)]
+    events_sender: channel::UnboundedSender<BufferedStorageEvent>,
 }
 
 impl<T: ContiguousBlockStore> BufferedStorage<T> {
@@ -137,9 +149,19 @@ impl<T: ContiguousBlockStore> BufferedStorage<T> {
             inner_subscriber,
             block_writes_sender: watch::channel(store_block_number).0,
             buffer: Mutex::new(BlockBuffer::new(store_block_number)),
+            #[cfg(test)]
+            events_sender: channel::unbounded().0,
         }
     }
 
+    #[cfg(test)]
+    pub(crate) fn set_events_sender(
+        &mut self,
+        sender: channel::UnboundedSender<BufferedStorageEvent>,
+    ) {
+        self.events_sender = sender;
+    }
+
     #[cfg(test)]
     pub(crate) fn as_ref(&self) -> &T {
         &self.inner
@@ -167,8 +189,12 @@ impl<T: ContiguousBlockStore> BufferedStorage<T> {
             if let Some(block) = next_block_for_store {
                 self.inner.schedule_next_block(ctx, &block).await?;
                 let block_number = block.block.number;
-                tracing::trace!(%block_number, "Block put in underlying storage");
+                tracing::trace!(%block_number, "Block scheduled in underlying storage");
             }
+
+            #[cfg(test)]
+            self.events_sender
+                .send(BufferedStorageEvent::UpdateReceived(store_block_number));
         }
     }
 }
@@ -242,7 +268,7 @@ impl<T: ContiguousBlockStore> WriteBlockStore for BufferedStorage<T> {
 
         if let Some(block) = next_block_for_store {
             self.inner.schedule_next_block(ctx, &block).await?;
-            tracing::trace!(block_number = %block.block.number, "Block put in underlying storage");
+            tracing::trace!(block_number = %block.block.number, "Block scheduled in underlying storage");
         }
         self.block_writes_sender.send_replace(block.block.number);
         Ok(())

node/libs/storage/src/tests.rs

@@ -1,5 +1,6 @@
 use super::*;
-use crate::types::ReplicaState;
+use crate::{buffered::BufferedStorageEvent, types::ReplicaState};
+use assert_matches::assert_matches;
 use async_trait::async_trait;
 use concurrency::{
     ctx::{self, channel},
@@ -267,7 +268,9 @@ async fn test_buffered_storage(
     initial_blocks.insert(0, genesis_block.clone());
 
     let (block_store, block_receiver) = MockContiguousStore::new(block_store);
-    let buffered_store = BufferedStorage::new(block_store);
+    let mut buffered_store = BufferedStorage::new(block_store);
+    let (events_sender, mut events_receiver) = channel::unbounded();
+    buffered_store.set_events_sender(events_sender);
 
     // Check initial values returned by the store.
     let last_initial_block = initial_blocks.last().unwrap().clone();
@@ -343,6 +346,16 @@ async fn test_buffered_storage(
             sync::changed(ctx, &mut inner_subscriber).await?;
         }
 
+        // Check events emitted by the buffered storage. This also ensures that all underlying storage
+        // updates are processed before proceeding to the following checks.
+        let expected_numbers = (initial_block_count as u64 + 1)..=last_block_number.0;
+        for expected_number in expected_numbers.map(BlockNumber) {
+            assert_matches!(
+                events_receiver.recv(ctx).await?,
+                BufferedStorageEvent::UpdateReceived(number) if number == expected_number
+            );
+        }
+
         assert_eq!(buffered_store.buffer_len().await, 0);
         Ok(())
     })