Merge pull request #652 from input-output-hk/djo/510/optimize-snapshot-digest

Optimize snapshot digest computation by adding a cache system

Author: Alenar

Cargo.lock

@@ -2,19 +2,18 @@ use clap::{Parser, Subcommand};
 use config::{builder::DefaultState, ConfigBuilder, Map, Source, Value, ValueKind};
 use slog::Level;
 use slog_scope::debug;
-use std::error::Error;
-use std::fs;
-use std::path::PathBuf;
-use std::sync::Arc;
-use tokio::sync::RwLock;
-use tokio::time::Duration;
+use std::{error::Error, fs, path::PathBuf, sync::Arc};
+use tokio::{sync::RwLock, time::Duration};
 
 use mithril_common::{
     certificate_chain::MithrilCertificateVerifier,
     chain_observer::{CardanoCliRunner, ChainObserver},
     crypto_helper::{key_decode_hex, ProtocolGenesisSigner, ProtocolGenesisVerifier},
     database::{ApplicationNodeType, DatabaseVersionChecker},
-    digesters::{CardanoImmutableDigester, ImmutableFileSystemObserver},
+    digesters::{
+        cache::{ImmutableFileDigestCacheProvider, JsonImmutableFileDigestCacheProviderBuilder},
+        CardanoImmutableDigester, ImmutableFileSystemObserver,
+    },
     entities::{Epoch, HexEncodedGenesisSecretKey},
     store::{adapter::SQLiteAdapter, StakeStore},
     BeaconProviderImpl,
@@ -254,6 +253,16 @@ pub struct ServeCommand {
     /// Defaults to work folder
     #[clap(long)]
     pub snapshot_directory: Option<PathBuf>,
+
+    /// Disable immutables digests cache.
+    #[clap(long)]
+    disable_digests_cache: bool,
+
+    /// If set the existing immutables digests cache will be reset.
+    ///
+    /// Will be ignored if set in conjunction with `--disable-digests-cache`.
+    #[clap(long)]
+    reset_digests_cache: bool,
 }
 
 impl Source for ServeCommand {
@@ -357,6 +366,7 @@ impl ServeCommand {
         ));
         let digester = Arc::new(CardanoImmutableDigester::new(
             config.db_directory.clone(),
+            self.build_digester_cache_provider(&config).await?,
             slog_scope::logger(),
         ));
         let multi_signer = Arc::new(RwLock::new(MultiSignerImpl::new(
@@ -453,6 +463,24 @@ impl ServeCommand {
         println!("Exiting...");
         Ok(())
     }
+
+    async fn build_digester_cache_provider(
+        &self,
+        config: &Configuration,
+    ) -> Result<Option<Arc<dyn ImmutableFileDigestCacheProvider>>, Box<dyn Error>> {
+        if self.disable_digests_cache {
+            return Ok(None);
+        }
+        let cache_provider = JsonImmutableFileDigestCacheProviderBuilder::new(
+            &config.data_stores_directory,
+            &format!("immutables_digests_{}.json", config.network),
+        )
+        .should_reset_digests_cache(self.reset_digests_cache)
+        .build()
+        .await?;
+
+        Ok(Some(Arc::new(cache_provider)))
+    }
 }
 
 /// Genesis tools

mithril-aggregator/src/command_args.rs

@@ -14,6 +14,7 @@ async-trait = "0.1.52"
 clap = { version = "4.0", features = ["derive", "env"] }
 cli-table = "0.4"
 config = "0.13.1"
+directories = "4.0.1"
 flate2 = "1.0.23"
 futures = "0.3"
 hex = "0.4.3"

mithril-client/Cargo.toml

@@ -1,13 +1,16 @@
-use std::{error::Error, path::Path, sync::Arc};
-
 use clap::Parser;
 use config::{builder::DefaultState, ConfigBuilder};
+use directories::ProjectDirs;
 use mithril_common::{
     certificate_chain::MithrilCertificateVerifier,
     crypto_helper::{key_decode_hex, ProtocolGenesisVerifier},
-    digesters::CardanoImmutableDigester,
+    digesters::{
+        cache::ImmutableFileDigestCacheProvider,
+        cache::JsonImmutableFileDigestCacheProviderBuilder, CardanoImmutableDigester,
+    },
 };
-use slog_scope::debug;
+use slog_scope::{debug, warn};
+use std::{error::Error, path::Path, sync::Arc};
 
 use crate::{AggregatorHTTPClient, AggregatorHandler, Config, Runtime};
 
@@ -18,6 +21,16 @@ pub struct RestoreCommand {
     #[clap(long)]
     json: bool,
 
+    /// Disable immutables digests cache.
+    #[clap(long)]
+    disable_digests_cache: bool,
+
+    /// If set the existing immutables digests cache will be reset.
+    ///
+    /// Will be ignored if set in conjunction with `--disable-digests-cache`.
+    #[clap(long)]
+    reset_digests_cache: bool,
+
     /// Digest of the snapshot to download. Use the `list` command to get that information.
     digest: String,
 }
@@ -37,14 +50,21 @@ impl RestoreCommand {
         debug!("{:?}", config);
         let mut runtime = Runtime::new(config.network.clone());
         let aggregator_handler =
-            AggregatorHTTPClient::new(config.network.clone(), config.aggregator_endpoint);
+            AggregatorHTTPClient::new(config.network.clone(), config.aggregator_endpoint.clone());
         let certificate_verifier = Box::new(MithrilCertificateVerifier::new(slog_scope::logger()));
         let genesis_verification_key = key_decode_hex(&config.genesis_verification_key)?;
         let genesis_verifier =
             ProtocolGenesisVerifier::from_verification_key(genesis_verification_key);
         let unpacked_path = aggregator_handler.unpack_snapshot(&self.digest).await?;
+
         let digester = Box::new(CardanoImmutableDigester::new(
             Path::new(&unpacked_path).into(),
+            build_digester_cache_provider(
+                self.disable_digests_cache,
+                self.reset_digests_cache,
+                &config,
+            )
+            .await?,
             slog_scope::logger(),
         ));
         let output = runtime
@@ -73,3 +93,32 @@ docker run -v cardano-node-ipc:/ipc -v cardano-node-data:/data --mount type=bind
         Ok(())
     }
 }
+
+async fn build_digester_cache_provider(
+    disable_digests_cache: bool,
+    reset_digests_cache: bool,
+    config: &Config,
+) -> Result<Option<Arc<dyn ImmutableFileDigestCacheProvider>>, Box<dyn Error>> {
+    if disable_digests_cache {
+        return Ok(None);
+    }
+
+    match ProjectDirs::from("io", "iohk", "mithril") {
+        None => {
+            warn!("Could not get cache directory, disabling immutables digests cache");
+            Ok(None)
+        }
+        Some(project_dirs) => {
+            let cache_provider = JsonImmutableFileDigestCacheProviderBuilder::new(
+                project_dirs.cache_dir(),
+                &format!("immutables_digests_{}.json", config.network),
+            )
+            .ensure_dir_exist()
+            .should_reset_digests_cache(reset_digests_cache)
+            .build()
+            .await?;
+
+            Ok(Some(Arc::new(cache_provider)))
+        }
+    }
+}

mithril-client/src/commands/restore.rs

@@ -11,11 +11,16 @@ repository = { workspace = true }
 [lib]
 crate-type = ["lib", "cdylib", "staticlib"]
 
+[[bench]]
+name = "digester"
+harness = false
+
 [dependencies]
 async-trait = "0.1.52"
 bech32 = "0.9.1"
 blake2 = "0.10.4"
 chrono = "0.4"
+digest = "0.10.6"
 ed25519-dalek = { version = "1.0.1", features = ["serde"] }
 fixed = "1.15.0"
 glob = "0.3"
@@ -36,6 +41,7 @@ serde_json = "1.0"
 serde_yaml = "0.9.10"
 sha2 = "0.10.2"
 slog = { version = "2.7.0", features = ["max_level_trace", "release_max_level_debug"] }
+slog-scope = "4.4.0"
 sqlite = "0.28"
 thiserror = "1.0.31"
 tokio = { version = "1.17.0", features = ["full"] }
@@ -51,7 +57,10 @@ mithril-stm = { path = "../mithril-stm" }
 mithril-stm = { path = "../mithril-stm", default-features = false, features = ["num-integer-backend"] }
 
 [dev-dependencies]
+criterion   = { version = "0.4.0", features = ["html_reports", "async_tokio"] }
+slog-async = "2.7.0"
 slog-scope = "4.4.0"
+slog-term = "2.9.0"
 
 [features]
 default = ["allow_skip_signer_certification"]

mithril-common/Cargo.toml

@@ -0,0 +1,98 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use mithril_common::{
+    digesters::cache::{ImmutableFileDigestCacheProvider, JsonImmutableFileDigestCacheProvider},
+    digesters::{
+        cache::MemoryImmutableFileDigestCacheProvider, CardanoImmutableDigester, ImmutableDigester,
+    },
+    entities::{Beacon, ImmutableFileNumber},
+};
+use slog::Drain;
+use std::{
+    fs,
+    fs::File,
+    io::prelude::Write,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
+
+fn temp_dir() -> PathBuf {
+    std::env::temp_dir()
+        .join("mithril_benches")
+        .join("digester")
+}
+
+fn db_dir() -> PathBuf {
+    temp_dir().join("db").join("immutable")
+}
+
+fn create_db(dir: &Path, number_of_immutables: ImmutableFileNumber, file_size: u64) {
+    if dir.exists() {
+        fs::remove_dir_all(dir).unwrap_or_else(|e| panic!("Could not remove dir {:?}: {}", dir, e));
+    }
+    fs::create_dir_all(dir).unwrap_or_else(|e| panic!("Could not create dir {:?}: {}", dir, e));
+
+    // + 1 to simulate "in-progress" immutable trio.
+    for filename in (1..=(number_of_immutables + 1)).flat_map(|i| {
+        [
+            format!("{:05}.chunk", i),
+            format!("{:05}.primary", i),
+            format!("{:05}.secondary", i),
+        ]
+    }) {
+        let file = dir.join(Path::new(&filename));
+        let mut source_file = File::create(file).unwrap();
+
+        write!(source_file, "This is a test file named '{}'", filename).unwrap();
+        writeln!(source_file).unwrap();
+        source_file.set_len(file_size).unwrap();
+    }
+}
+
+#[inline]
+fn create_logger() -> slog::Logger {
+    let drain = slog_async::Async::new(slog::Discard).build().fuse();
+    slog::Logger::root(Arc::new(drain), slog::o!())
+}
+
+#[inline]
+async fn compute_digest(
+    cache_provider: Option<Arc<dyn ImmutableFileDigestCacheProvider>>,
+    number_of_immutables: ImmutableFileNumber,
+) {
+    let digester = CardanoImmutableDigester::new(db_dir(), cache_provider, create_logger());
+    digester
+        .compute_digest(&Beacon::new("devnet".to_string(), 1, number_of_immutables))
+        .await
+        .expect("digest computation should not fail");
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let runtime = tokio::runtime::Runtime::new().unwrap();
+    let number_of_immutable = 25;
+    create_db(&db_dir(), number_of_immutable, 65536 * 4);
+
+    c.bench_function("digester no cache", |bencher| {
+        bencher
+            .to_async(&runtime)
+            .iter(|| async { compute_digest(None, number_of_immutable).await })
+    });
+    c.bench_function("digester memory cache", |bencher| {
+        let cache = Arc::new(MemoryImmutableFileDigestCacheProvider::default());
+
+        bencher
+            .to_async(&runtime)
+            .iter(|| async { compute_digest(Some(cache.clone()), number_of_immutable).await })
+    });
+    c.bench_function("digester json cache", |bencher| {
+        let cache = Arc::new(JsonImmutableFileDigestCacheProvider::new(
+            &temp_dir().join("immutable-cache-store.json"),
+        ));
+
+        bencher
+            .to_async(&runtime)
+            .iter(|| async { compute_digest(Some(cache.clone()), number_of_immutable).await })
+    });
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);