Merge pull request #1904 from input-output-hk/jpraynaud/1903-handle-multiple-merkle-tree-storage-backend

Optimize memory usage of signer for Cardano transactions

Author: jpraynaud

CHANGELOG.md

@@ -9,9 +9,13 @@ As a minor extension, we have adopted a slightly different versioning convention
 
 ## Mithril Distribution [XXXX] - UNRELEASED
 
-- Support for Mithril nodes footprint support in Prometheus monitoring in infrastructure
-- Add support for custom HTTP headers in Mithril client WASM library
-- Support `file://` URLs for snapshot locations
+- Support for Mithril nodes footprint support in Prometheus monitoring in infrastructure.
+
+- Add support for custom HTTP headers in Mithril client WASM library.
+
+- Support `file://` URLs for snapshot locations in Mithril client.
+
+- Support for Mithril signer memory optimization when signing Cardano transactions with multiple Merkle tree storage backends.
 
 - **UNSTABLE** Cardano stake distribution certification:
 
@@ -102,13 +106,9 @@ As a minor extension, we have adopted a slightly different versioning convention
 - **UNSTABLE** Cardano transactions certification:
 
   - Support computation of the Cardano Transactions signature and proving with the pre-computed Block Range Merkle Roots retrieved from the database.
-
   - Prune Cardano Transactions from the signer database after the Block Range Merkle Roots have been computed.
-
   - Implement a Chain Reader which retrieves blocks from the Cardano chain with Pallas through the `chainsync` mini-protocol.
-
   - Implement a Resource Pool and use it for caching Block Range Merkle maps used by the Cardano transactions prover and improving the throughput.
-
   - Change the beacon of the Cardano Transactions to a block number instead of an immutable file number.
 
 - Crates versions:

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "mithril-persistence"
-version = "0.2.23"
+version = "0.2.24"
 description = "Common types, interfaces, and utilities to persist data for Mithril nodes."
 authors = { workspace = true }
 edition = { workspace = true }

internal/mithril-persistence/Cargo.toml

@@ -4,7 +4,7 @@ use std::sync::Arc;
 use anyhow::Context;
 use async_trait::async_trait;
 
-use mithril_common::crypto_helper::MKTreeNode;
+use mithril_common::crypto_helper::{MKTreeNode, MKTreeStorer};
 use mithril_common::entities::{
     BlockHash, BlockNumber, BlockRange, CardanoTransaction, ChainPoint, SlotNumber, TransactionHash,
 };
@@ -296,7 +296,7 @@ impl CardanoTransactionRepository {
 }
 
 #[async_trait]
-impl BlockRangeRootRetriever for CardanoTransactionRepository {
+impl<S: MKTreeStorer> BlockRangeRootRetriever<S> for CardanoTransactionRepository {
     async fn retrieve_block_range_roots<'a>(
         &'a self,
         up_to_beacon: BlockNumber,

internal/mithril-persistence/src/database/repository/cardano_transaction_repository.rs

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

mithril-aggregator/Cargo.toml

@@ -21,7 +21,8 @@ use mithril_common::{
     chain_observer::{CardanoCliRunner, ChainObserver, ChainObserverBuilder, FakeObserver},
     chain_reader::{ChainBlockReader, PallasChainReader},
     crypto_helper::{
-        ProtocolGenesisSigner, ProtocolGenesisVerificationKey, ProtocolGenesisVerifier,
+        MKTreeStoreInMemory, ProtocolGenesisSigner, ProtocolGenesisVerificationKey,
+        ProtocolGenesisVerifier,
     },
     digesters::{
         cache::{ImmutableFileDigestCacheProvider, JsonImmutableFileDigestCacheProviderBuilder},
@@ -1085,7 +1086,9 @@ impl DependenciesBuilder {
         ));
         let transactions_importer = self.get_transactions_importer().await?;
         let block_range_root_retriever = self.get_transaction_repository().await?;
-        let cardano_transactions_builder = Arc::new(CardanoTransactionsSignableBuilder::new(
+        let cardano_transactions_builder = Arc::new(CardanoTransactionsSignableBuilder::<
+            MKTreeStoreInMemory,
+        >::new(
             transactions_importer,
             block_range_root_retriever,
             self.get_logger()?,
@@ -1487,7 +1490,7 @@ impl DependenciesBuilder {
         let transaction_retriever = self.get_transaction_repository().await?;
         let block_range_root_retriever = self.get_transaction_repository().await?;
         let logger = self.get_logger()?;
-        let prover_service = MithrilProverService::new(
+        let prover_service = MithrilProverService::<MKTreeStoreInMemory>::new(
             transaction_retriever,
             block_range_root_retriever,
             mk_map_pool_size,

mithril-aggregator/src/dependency_injection/builder.rs

@@ -8,7 +8,7 @@ use slog::{debug, Logger};
 use tokio::{runtime::Handle, task};
 
 use mithril_common::cardano_block_scanner::{BlockScanner, ChainScannedBlocks};
-use mithril_common::crypto_helper::{MKTree, MKTreeNode};
+use mithril_common::crypto_helper::{MKTree, MKTreeNode, MKTreeStoreInMemory};
 use mithril_common::entities::{
     BlockNumber, BlockRange, CardanoTransaction, ChainPoint, SlotNumber,
 };
@@ -152,7 +152,7 @@ impl CardanoTransactionsImporter {
                 continue;
             }
 
-            let merkle_root = MKTree::new(&transactions)?.compute_root()?;
+            let merkle_root = MKTree::<MKTreeStoreInMemory>::new(&transactions)?.compute_root()?;
             block_ranges_with_merkle_root.push((block_range, merkle_root));
 
             if block_ranges_with_merkle_root.len() >= 100 {
@@ -256,7 +256,10 @@ mod tests {
             .iter()
             .flat_map(|br| br.clone().into_transactions())
             .collect();
-        MKTree::new(&tx).unwrap().compute_root().unwrap()
+        MKTree::<MKTreeStoreInMemory>::new(&tx)
+            .unwrap()
+            .compute_root()
+            .unwrap()
     }
 
     #[tokio::test]

mithril-aggregator/src/services/cardano_transactions_importer.rs

@@ -8,7 +8,7 @@ use std::{
 };
 
 use mithril_common::{
-    crypto_helper::{MKMap, MKMapNode, MKTree},
+    crypto_helper::{MKMap, MKMapNode, MKTree, MKTreeStorer},
     entities::{
         BlockNumber, BlockRange, CardanoTransaction, CardanoTransactionsSetProof, TransactionHash,
     },
@@ -51,18 +51,18 @@ pub trait TransactionsRetriever: Sync + Send {
 }
 
 /// Mithril prover
-pub struct MithrilProverService {
+pub struct MithrilProverService<S: MKTreeStorer> {
     transaction_retriever: Arc<dyn TransactionsRetriever>,
-    block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
-    mk_map_pool: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange>>>,
+    block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
+    mk_map_pool: ResourcePool<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>,
     logger: Logger,
 }
 
-impl MithrilProverService {
+impl<S: MKTreeStorer> MithrilProverService<S> {
     /// Create a new Mithril prover
     pub fn new(
         transaction_retriever: Arc<dyn TransactionsRetriever>,
-        block_range_root_retriever: Arc<dyn BlockRangeRootRetriever>,
+        block_range_root_retriever: Arc<dyn BlockRangeRootRetriever<S>>,
         mk_map_pool_size: usize,
         logger: Logger,
     ) -> Self {
@@ -113,7 +113,7 @@ impl MithrilProverService {
 }
 
 #[async_trait]
-impl ProverService for MithrilProverService {
+impl<S: MKTreeStorer> ProverService for MithrilProverService<S> {
     async fn compute_transactions_proofs(
         &self,
         up_to: BlockNumber,
@@ -126,7 +126,7 @@ impl ProverService for MithrilProverService {
             .await?;
 
         // 2 - Compute block ranges sub Merkle trees
-        let mk_trees: StdResult<Vec<(BlockRange, MKTree)>> = block_range_transactions
+        let mk_trees: StdResult<Vec<(BlockRange, MKTree<S>)>> = block_range_transactions
             .into_iter()
             .map(|(block_range, transactions)| {
                 let mk_tree = MKTree::new(&transactions)?;
@@ -183,7 +183,7 @@ impl ProverService for MithrilProverService {
                 );
                 mk_map_cache.clone()
             })
-            .collect::<Vec<MKMap<_, _>>>();
+            .collect::<Vec<MKMap<_, _, _>>>();
         debug!(self.logger, "Prover is draining the Merkle map pool");
         let discriminant_new = self.mk_map_pool.discriminant()? + 1;
         self.mk_map_pool.set_discriminant(discriminant_new)?;
@@ -211,7 +211,9 @@ impl ProverService for MithrilProverService {
 #[cfg(test)]
 mod tests {
     use anyhow::anyhow;
-    use mithril_common::crypto_helper::{MKMap, MKMapNode, MKTreeNode};
+    use mithril_common::crypto_helper::{
+        MKMap, MKMapNode, MKTreeNode, MKTreeStoreInMemory, MKTreeStorer,
+    };
     use mithril_common::entities::CardanoTransaction;
     use mithril_common::test_utils::CardanoTransactionsBuilder;
     use mockall::mock;
@@ -220,10 +222,10 @@ mod tests {
     use super::*;
 
     mock! {
-        pub BlockRangeRootRetrieverImpl { }
+        pub BlockRangeRootRetrieverImpl<S: MKTreeStorer> { }
 
         #[async_trait]
-        impl BlockRangeRootRetriever for BlockRangeRootRetrieverImpl {
+        impl<S: MKTreeStorer> BlockRangeRootRetriever<S> for BlockRangeRootRetrieverImpl<S> {
             async fn retrieve_block_range_roots<'a>(
                 &'a self,
                 up_to_beacon: BlockNumber,
@@ -232,11 +234,13 @@ mod tests {
             async fn compute_merkle_map_from_block_range_roots(
                 &self,
                 up_to_beacon: BlockNumber,
-            ) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange>>>;
+            ) -> StdResult<MKMap<BlockRange, MKMapNode<BlockRange, S>, S>>;
         }
     }
 
     mod test_data {
+        use mithril_common::crypto_helper::MKTreeStoreInMemory;
+
         use super::*;
 
         pub fn filter_transactions_for_indices(
@@ -287,15 +291,16 @@ mod tests {
 
         pub fn compute_mk_map_from_block_ranges_map(
             block_ranges_map: BTreeMap<BlockRange, Vec<CardanoTransaction>>,
-        ) -> MKMap<BlockRange, MKMapNode<BlockRange>> {
+        ) -> MKMap<BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>, MKTreeStoreInMemory>
+        {
             MKMap::new_from_iter(
                 block_ranges_map
                     .into_iter()
                     .map(|(block_range, transactions)| {
                         (
                             block_range,
                             MKMapNode::TreeNode(
-                                MKTree::new(&transactions)
+                                MKTree::<MKTreeStoreInMemory>::new(&transactions)
                                     .unwrap()
                                     .compute_root()
                                     .unwrap()
@@ -348,13 +353,13 @@ mod tests {
         }
     }
 
-    fn build_prover<F, G>(
+    fn build_prover<F, G, S: MKTreeStorer + 'static>(
         transaction_retriever_mock_config: F,
         block_range_root_retriever_mock_config: G,
-    ) -> MithrilProverService
+    ) -> MithrilProverService<S>
     where
         F: FnOnce(&mut MockTransactionsRetriever),
-        G: FnOnce(&mut MockBlockRangeRootRetrieverImpl),
+        G: FnOnce(&mut MockBlockRangeRootRetrieverImpl<S>),
     {
         let mut transaction_retriever = MockTransactionsRetriever::new();
         transaction_retriever_mock_config(&mut transaction_retriever);
@@ -581,7 +586,7 @@ mod tests {
         let transactions_to_prove =
             test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
         let test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
-        let prover = build_prover(
+        let prover = build_prover::<_, _, MKTreeStoreInMemory>(
             |transaction_retriever_mock| {
                 transaction_retriever_mock
                     .expect_get_by_hashes()
@@ -610,7 +615,7 @@ mod tests {
         let transactions_to_prove =
             test_data::filter_transactions_for_indices(&[1, 2, 4], &transactions);
         let test_data = test_data::build_test_data(&transactions_to_prove, &transactions);
-        let prover = build_prover(
+        let prover = build_prover::<_, _, MKTreeStoreInMemory>(
             |transaction_retriever_mock| {
                 let transactions_to_prove = transactions_to_prove.clone();
                 transaction_retriever_mock

mithril-aggregator/src/services/prover.rs

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

mithril-common/Cargo.toml

@@ -1,6 +1,6 @@
 use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion};
 use mithril_common::{
-    crypto_helper::{MKMap, MKMapNode, MKMapValue, MKTree, MKTreeNode},
+    crypto_helper::{MKMap, MKMapNode, MKMapValue, MKTree, MKTreeNode, MKTreeStoreInMemory},
     entities::BlockRange,
 };
 
@@ -29,7 +29,7 @@ fn generate_block_ranges_nodes_iterator(
     total_transactions_per_block: u64,
     block_range_length: u64,
     max_uncompressed_block_ranges: u64,
-) -> impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)> {
+) -> impl Iterator<Item = (BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>)> {
     let total_transactions_per_block_range = total_transactions_per_block * block_range_length;
     let total_block_ranges = total_transactions / total_transactions_per_block_range;
     assert!(
@@ -56,8 +56,10 @@ fn generate_block_ranges_nodes_iterator(
 }
 
 fn generate_merkle_map_compressed(
-    block_ranges_nodes_iterator: impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)>,
-) -> MKMap<BlockRange, MKMapNode<BlockRange>> {
+    block_ranges_nodes_iterator: impl Iterator<
+        Item = (BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>),
+    >,
+) -> MKMap<BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>, MKTreeStoreInMemory> {
     let mut mk_map = MKMap::new(&[]).unwrap();
     for (block_range, mk_map_node) in block_ranges_nodes_iterator {
         mk_map
@@ -68,10 +70,19 @@ fn generate_merkle_map_compressed(
 }
 
 fn generate_merkle_map(
-    block_ranges_nodes_iterator: impl Iterator<Item = (BlockRange, MKMapNode<BlockRange>)>,
-    mk_map_compressed: &MKMap<BlockRange, MKMapNode<BlockRange>>,
+    block_ranges_nodes_iterator: impl Iterator<
+        Item = (BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>),
+    >,
+    mk_map_compressed: &MKMap<
+        BlockRange,
+        MKMapNode<BlockRange, MKTreeStoreInMemory>,
+        MKTreeStoreInMemory,
+    >,
     total_proofs: u64,
-) -> (MKMap<BlockRange, MKMapNode<BlockRange>>, Vec<MKTreeNode>) {
+) -> (
+    MKMap<BlockRange, MKMapNode<BlockRange, MKTreeStoreInMemory>, MKTreeStoreInMemory>,
+    Vec<MKTreeNode>,
+) {
     let mut leaves_to_prove_all: Vec<MKTreeNode> = vec![];
     let mut mk_map = mk_map_compressed.clone();
     for (mk_map_key_to_prove, mk_map_node_to_prove) in &block_ranges_nodes_iterator