Merge pull request #1716 from input-output-hk/ensemble/modify-block-streamer-return-type

Update `BlockStreamer` return type

Author: dlachaume

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-aggregator"
-version = "0.5.9"
+version = "0.5.10"
 description = "A Mithril Aggregator server"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-aggregator/Cargo.toml

@@ -1,3 +1,4 @@
+use anyhow::anyhow;
 use std::mem;
 use std::ops::Range;
 use std::path::{Path, PathBuf};
@@ -6,7 +7,7 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use slog::{debug, Logger};
 
-use mithril_common::cardano_block_scanner::BlockScanner;
+use mithril_common::cardano_block_scanner::{BlockScanner, ChainScannedBlocks};
 use mithril_common::crypto_helper::{MKTree, MKTreeNode};
 use mithril_common::entities::{BlockNumber, BlockRange, CardanoTransaction, ImmutableFileNumber};
 use mithril_common::signable_builder::TransactionsImporter;
@@ -105,14 +106,21 @@ impl CardanoTransactionsImporter {
         let mut streamer = self.block_scanner.scan(&self.dirpath, from, until).await?;
 
         while let Some(blocks) = streamer.poll_next().await? {
-            let parsed_transactions: Vec<CardanoTransaction> = blocks
-                .into_iter()
-                .flat_map(|b| b.into_transactions())
-                .collect();
-
-            self.transaction_store
-                .store_transactions(parsed_transactions)
-                .await?;
+            match blocks {
+                ChainScannedBlocks::RollForwards(forward_blocks) => {
+                    let parsed_transactions: Vec<CardanoTransaction> = forward_blocks
+                        .into_iter()
+                        .flat_map(|b| b.into_transactions())
+                        .collect();
+
+                    self.transaction_store
+                        .store_transactions(parsed_transactions)
+                        .await?;
+                }
+                ChainScannedBlocks::RollBackward(_) => {
+                    return Err(anyhow!("RollBackward not yet implemented"))
+                }
+            }
         }
 
         Ok(())

mithril-aggregator/src/services/cardano_transactions_importer.rs

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-common"
-version = "0.4.8"
+version = "0.4.9"
 description = "Common types, interfaces, and utilities for Mithril nodes."
 authors = { workspace = true }
 edition = { workspace = true }

mithril-common/Cargo.toml

@@ -60,6 +60,7 @@ impl BlockScanner for CardanoBlockScanner {
 
 #[cfg(test)]
 mod tests {
+    use crate::cardano_block_scanner::BlockStreamerTestExtensions;
     use crate::test_utils::{TempDir, TestLogger};
 
     use super::*;

mithril-common/src/cardano_block_scanner/block_scanner.rs

@@ -4,6 +4,7 @@ use std::path::Path;
 use async_trait::async_trait;
 use tokio::sync::RwLock;
 
+use crate::cardano_block_scanner::ChainScannedBlocks;
 use crate::cardano_block_scanner::{BlockScanner, BlockStreamer, ScannedBlock};
 use crate::entities::ImmutableFileNumber;
 use crate::StdResult;
@@ -57,13 +58,18 @@ impl DumbBlockStreamer {
 
 #[async_trait]
 impl BlockStreamer for DumbBlockStreamer {
-    async fn poll_next(&mut self) -> StdResult<Option<Vec<ScannedBlock>>> {
-        Ok(self.blocks.pop_front())
+    async fn poll_next(&mut self) -> StdResult<Option<ChainScannedBlocks>> {
+        Ok(self
+            .blocks
+            .pop_front()
+            .map(ChainScannedBlocks::RollForwards))
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use crate::cardano_block_scanner::BlockStreamerTestExtensions;
+
     use super::*;
 
     #[tokio::test]
@@ -79,7 +85,10 @@ mod tests {
         let mut streamer = DumbBlockStreamer::new(vec![expected_blocks.clone()]);
 
         let blocks = streamer.poll_next().await.unwrap();
-        assert_eq!(blocks, Some(expected_blocks));
+        assert_eq!(
+            blocks,
+            Some(ChainScannedBlocks::RollForwards(expected_blocks))
+        );
 
         let blocks = streamer.poll_next().await.unwrap();
         assert_eq!(blocks, None);
@@ -98,13 +107,22 @@ mod tests {
         let mut streamer = DumbBlockStreamer::new(expected_blocks.clone());
 
         let blocks = streamer.poll_next().await.unwrap();
-        assert_eq!(blocks, Some(expected_blocks[0].clone()));
+        assert_eq!(
+            blocks,
+            Some(ChainScannedBlocks::RollForwards(expected_blocks[0].clone()))
+        );
 
         let blocks = streamer.poll_next().await.unwrap();
-        assert_eq!(blocks, Some(expected_blocks[1].clone()));
+        assert_eq!(
+            blocks,
+            Some(ChainScannedBlocks::RollForwards(expected_blocks[1].clone()))
+        );
 
         let blocks = streamer.poll_next().await.unwrap();
-        assert_eq!(blocks, Some(expected_blocks[2].clone()));
+        assert_eq!(
+            blocks,
+            Some(ChainScannedBlocks::RollForwards(expected_blocks[2].clone()))
+        );
 
         let blocks = streamer.poll_next().await.unwrap();
         assert_eq!(blocks, None);

mithril-common/src/cardano_block_scanner/dumb_block_scanner.rs

@@ -7,6 +7,7 @@ use pallas_hardano::storage::immutable::chunk::{read_blocks, Reader};
 use pallas_traverse::MultiEraBlock;
 use slog::{debug, error, Logger};
 
+use crate::cardano_block_scanner::ChainScannedBlocks;
 use crate::cardano_block_scanner::{BlockStreamer, ScannedBlock};
 use crate::digesters::ImmutableFile;
 use crate::StdResult;
@@ -20,7 +21,7 @@ pub struct ImmutableBlockStreamer {
 
 #[async_trait]
 impl BlockStreamer for ImmutableBlockStreamer {
-    async fn poll_next(&mut self) -> StdResult<Option<Vec<ScannedBlock>>> {
+    async fn poll_next(&mut self) -> StdResult<Option<ChainScannedBlocks>> {
         match &self.remaining_immutable_files.pop_front() {
             Some(immutable_file) => {
                 debug!(
@@ -37,7 +38,7 @@ impl BlockStreamer for ImmutableBlockStreamer {
                             immutable_file.path.display()
                         )
                     })?;
-                Ok(Some(blocks))
+                Ok(Some(ChainScannedBlocks::RollForwards(blocks)))
             }
             None => Ok(None),
         }
@@ -123,12 +124,22 @@ impl ImmutableBlockStreamer {
 
 #[cfg(test)]
 mod tests {
+    use crate::cardano_block_scanner::BlockStreamerTestExtensions;
     use crate::test_utils::{TempDir, TestLogger};
 
     use super::*;
 
     #[tokio::test]
     async fn test_parse_expected_number_of_transactions() {
+        fn sum_of_transactions_len(o: Option<ChainScannedBlocks>) -> Option<usize> {
+            match o {
+                Some(ChainScannedBlocks::RollForwards(b)) => {
+                    Some(b.into_iter().map(|b| b.transactions_len()).sum())
+                }
+                _ => None,
+            }
+        }
+
         // We know the number of transactions in those prebuilt immutables
         let immutable_files = [
             ("00000.chunk", 0usize),
@@ -148,19 +159,19 @@ mod tests {
 
         let immutable_blocks = streamer.poll_next().await.unwrap();
         assert_eq!(
-            immutable_blocks.map(|b| b.into_iter().map(|b| b.transactions_len()).sum()),
+            sum_of_transactions_len(immutable_blocks),
             Some(immutable_files[0].1)
         );
 
         let immutable_blocks = streamer.poll_next().await.unwrap();
         assert_eq!(
-            immutable_blocks.map(|b| b.into_iter().map(|b| b.transactions_len()).sum()),
+            sum_of_transactions_len(immutable_blocks),
             Some(immutable_files[1].1)
         );
 
         let immutable_blocks = streamer.poll_next().await.unwrap();
         assert_eq!(
-            immutable_blocks.map(|b| b.into_iter().map(|b| b.transactions_len()).sum()),
+            sum_of_transactions_len(immutable_blocks),
             Some(immutable_files[2].1)
         );
 

mithril-common/src/cardano_block_scanner/immutable_block_streamer.rs

@@ -1,9 +1,10 @@
+use anyhow::anyhow;
 use std::path::Path;
 
 use async_trait::async_trait;
 
 use crate::cardano_block_scanner::ScannedBlock;
-use crate::entities::ImmutableFileNumber;
+use crate::entities::{ChainPoint, ImmutableFileNumber};
 use crate::StdResult;
 
 /// A scanner that can read cardano transactions in a cardano database
@@ -55,18 +56,46 @@ pub trait BlockScanner: Sync + Send {
     ) -> StdResult<Box<dyn BlockStreamer>>;
 }
 
+/// [ChainScannedBlocks] allows to scan new blocks and handle rollbacks
+#[derive(Debug, PartialEq)]
+pub enum ChainScannedBlocks {
+    /// Roll forward on the chain to the next list of [ScannedBlock]
+    RollForwards(Vec<ScannedBlock>),
+    /// Roll backward on the chain to the previous [ChainPoint]
+    RollBackward(ChainPoint),
+}
+
 /// Trait that define how blocks are streamed from a Cardano database
 #[async_trait]
 pub trait BlockStreamer: Sync + Send {
     /// Stream the next available blocks
-    async fn poll_next(&mut self) -> StdResult<Option<Vec<ScannedBlock>>>;
+    async fn poll_next(&mut self) -> StdResult<Option<ChainScannedBlocks>>;
+}
 
-    /// Stream all the available blocks, may be very memory intensive
-    async fn poll_all(&mut self) -> StdResult<Vec<ScannedBlock>> {
-        let mut blocks = Vec::new();
-        while let Some(mut next_blocks) = self.poll_next().await? {
-            blocks.append(&mut next_blocks);
+cfg_test_tools! {
+    /// Tests extensions methods for the [BlockStreamer] trait.
+    #[async_trait]
+    pub trait BlockStreamerTestExtensions{
+        /// Stream all the available blocks, may be very memory intensive
+        async fn poll_all(&mut self) -> StdResult<Vec<ScannedBlock>>;
+    }
+
+    #[async_trait]
+    impl <S: BlockStreamer + ?Sized> BlockStreamerTestExtensions for S {
+        async fn poll_all(&mut self) -> StdResult<Vec<ScannedBlock>> {
+            let mut all_blocks = Vec::new();
+            while let Some(next_blocks) = self.poll_next().await? {
+                match next_blocks {
+                    ChainScannedBlocks::RollForwards(mut forward_blocks) => {
+                        all_blocks.append(&mut forward_blocks);
+                    }
+                    ChainScannedBlocks::RollBackward(_) => {
+                        return Err(anyhow!("poll_all: RollBackward not supported"));
+                    }
+                };
+            }
+            Ok(all_blocks)
         }
-        Ok(blocks)
     }
+
 }

mithril-common/src/cardano_block_scanner/interface.rs

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-signer"
-version = "0.2.136"
+version = "0.2.137"
 description = "A Mithril Signer"
 authors = { workspace = true }
 edition = { workspace = true }

mithril-signer/Cargo.toml

@@ -6,7 +6,7 @@ use std::sync::Arc;
 use async_trait::async_trait;
 use slog::{debug, Logger};
 
-use mithril_common::cardano_block_scanner::BlockScanner;
+use mithril_common::cardano_block_scanner::{BlockScanner, ChainScannedBlocks};
 use mithril_common::crypto_helper::{MKTree, MKTreeNode};
 use mithril_common::entities::{BlockNumber, BlockRange, CardanoTransaction, ImmutableFileNumber};
 use mithril_common::signable_builder::TransactionsImporter;
@@ -105,14 +105,21 @@ impl CardanoTransactionsImporter {
         let mut streamer = self.block_scanner.scan(&self.dirpath, from, until).await?;
 
         while let Some(blocks) = streamer.poll_next().await? {
-            let parsed_transactions: Vec<CardanoTransaction> = blocks
-                .into_iter()
-                .flat_map(|b| b.into_transactions())
-                .collect();
-
-            self.transaction_store
-                .store_transactions(parsed_transactions)
-                .await?;
+            match blocks {
+                ChainScannedBlocks::RollForwards(blocks) => {
+                    let parsed_transactions: Vec<CardanoTransaction> = blocks
+                        .into_iter()
+                        .flat_map(|b| b.into_transactions())
+                        .collect();
+
+                    self.transaction_store
+                        .store_transactions(parsed_transactions)
+                        .await?;
+                }
+                ChainScannedBlocks::RollBackward(_) => {
+                    return Err(anyhow::anyhow!("RollBackward not supported"));
+                }
+            }
         }
 
         Ok(())