Introduce the AggregatorRuntime

Author: Alenar

mithril-aggregator/src/lib.rs

@@ -2,6 +2,7 @@ mod dependency;
 mod entities;
 mod http_server;
 mod multi_signer;
+mod runtime;
 mod snapshot_store;
 mod snapshotter;
 
@@ -13,4 +14,5 @@ pub use crate::multi_signer::{
 };
 pub use crate::snapshot_store::SnapshotStoreHTTPClient;
 pub use dependency::DependencyManager;
+pub use runtime::AggregatorRuntime;
 pub use snapshotter::Snapshotter;

mithril-aggregator/src/main.rs

@@ -3,9 +3,9 @@
 use clap::Parser;
 
 use mithril_aggregator::{
-    key_decode_hex, Config, DependencyManager, MultiSigner, MultiSignerImpl, ProtocolParameters,
-    ProtocolPartyId, ProtocolSignerVerificationKey, ProtocolStake, Server, SnapshotStoreHTTPClient,
-    Snapshotter,
+    key_decode_hex, AggregatorRuntime, Config, DependencyManager, MultiSigner, MultiSignerImpl,
+    ProtocolParameters, ProtocolPartyId, ProtocolSignerVerificationKey, ProtocolStake, Server,
+    SnapshotStoreHTTPClient,
 };
 use mithril_common::fake_data;
 use slog::{Drain, Level, Logger};
@@ -98,12 +98,9 @@ async fn main() {
 
     // Start snapshot uploader
     let handle = tokio::spawn(async move {
-        let snapshotter = Snapshotter::new(
-            args.snapshot_interval * 1000,
-            args.db_directory.clone(),
-            slog_scope::logger(),
-        );
-        snapshotter.run().await
+        let runtime =
+            AggregatorRuntime::new(args.snapshot_interval * 1000, args.db_directory.clone());
+        runtime.run().await
     });
 
     // Start REST server

mithril-aggregator/src/runtime.rs

@@ -0,0 +1,48 @@
+use crate::Snapshotter;
+use mithril_common::immutable_digester::ImmutableDigester;
+use slog_scope::{error, info};
+use tokio::time::{sleep, Duration};
+
+/// AggregatorRuntime
+pub struct AggregatorRuntime {
+    /// Interval between each snapshot, in seconds
+    interval: u32,
+
+    /// DB directory to snapshot
+    db_directory: String,
+}
+
+impl AggregatorRuntime {
+    /// AggregatorRuntime factory
+    pub fn new(interval: u32, db_directory: String) -> Self {
+        Self {
+            interval,
+            db_directory,
+        }
+    }
+
+    /// Run snapshotter loop
+    pub async fn run(&self) {
+        info!("Starting Snapshotter");
+        let snapshotter = Snapshotter::new(self.db_directory.clone());
+        let digester = ImmutableDigester::new(self.db_directory.clone(), slog_scope::logger());
+
+        loop {
+            info!("Snapshotting");
+
+            match digester.compute_digest() {
+                Ok(digest_result) => {
+                    if let Err(e) = snapshotter.snapshot(digest_result.digest).await {
+                        error!("{:?}", e)
+                    }
+                }
+                Err(e) => {
+                    error!("{:?}", e)
+                }
+            };
+
+            info!("Sleeping for {}", self.interval);
+            sleep(Duration::from_millis(self.interval.into())).await;
+        }
+    }
+}

mithril-aggregator/src/snapshotter.rs

@@ -1,5 +1,4 @@
 use mithril_common::entities::Snapshot;
-use mithril_common::immutable_digester::ImmutableDigester;
 
 use chrono::prelude::*;
 use cloud_storage::bucket::Entity;
@@ -8,31 +7,24 @@ use cloud_storage::object_access_control::NewObjectAccessControl;
 use cloud_storage::Client;
 use flate2::write::GzEncoder;
 use flate2::Compression;
-use slog::{error, info, Logger};
+use slog_scope::info;
 use std::env;
 use std::fs::File;
 use std::io;
 use std::io::{Seek, SeekFrom};
 use std::path::Path;
 use thiserror::Error;
-use tokio::time::{sleep, Duration};
 use tokio_util::codec::BytesCodec;
 use tokio_util::codec::FramedRead;
 
 /// Snapshotter
 pub struct Snapshotter {
-    /// Interval between each snapshot, in seconds
-    interval: u32,
-
     /// DB directory to snapshot
     db_directory: String,
-
-    /// The logger where the logs should be written
-    logger: Logger,
 }
 
 #[derive(Error, Debug)]
-enum SnapshotError {
+pub enum SnapshotError {
     #[error("Create archive error: ")]
     CreateArchiveError(#[from] io::Error),
 
@@ -42,40 +34,13 @@ enum SnapshotError {
 
 impl Snapshotter {
     /// Snapshotter factory
-    pub fn new(interval: u32, db_directory: String, logger: Logger) -> Self {
-        Self {
-            interval,
-            db_directory,
-            logger,
-        }
-    }
-
-    /// Run snapshotter loop
-    pub async fn run(&self) {
-        info!(self.logger, "Starting Snapshotter");
-        loop {
-            info!(self.logger, "Snapshotting");
-
-            let digester = ImmutableDigester::new(self.db_directory.clone(), self.logger.clone());
-            match digester.compute_digest() {
-                Ok(digest_result) => {
-                    if let Err(e) = self.snapshot(digest_result.digest).await {
-                        error!(self.logger, "{:?}", e)
-                    }
-                }
-                Err(e) => {
-                    error!(self.logger, "{:?}", e)
-                }
-            };
-
-            info!(self.logger, "Sleeping for {}", self.interval);
-            sleep(Duration::from_millis(self.interval.into())).await;
-        }
+    pub fn new(db_directory: String) -> Self {
+        Self { db_directory }
     }
 
-    async fn snapshot(&self, immutable_digest: String) -> Result<(), SnapshotError> {
+    pub async fn snapshot(&self, immutable_digest: String) -> Result<(), SnapshotError> {
         let archive_name = "testnet.tar.gz";
-        info!(self.logger, "snapshot hash: {}", immutable_digest);
+        info!("snapshot hash: {}", immutable_digest);
 
         let size = self.create_archive(archive_name)?;
 
@@ -93,11 +58,7 @@ impl Snapshotter {
             )],
         }];
 
-        info!(
-            self.logger,
-            "snapshot: {}",
-            serde_json::to_string(&snapshots).unwrap()
-        );
+        info!("snapshot: {}", serde_json::to_string(&snapshots).unwrap());
         serde_json::to_writer(&File::create("snapshots.json").unwrap(), &snapshots).unwrap();
 
         self.upload_file(archive_name).await?;
@@ -113,7 +74,6 @@ impl Snapshotter {
         let mut tar = tar::Builder::new(enc);
 
         info!(
-            self.logger,
             "compressing {} into {}",
             &self.db_directory,
             &path.to_str().unwrap()
@@ -142,7 +102,7 @@ impl Snapshotter {
             ));
         };
 
-        info!(self.logger, "uploading {}", filename);
+        info!("uploading {}", filename);
         let client = Client::default();
         let file = tokio::fs::File::open(filename).await.unwrap();
         let stream = FramedRead::new(file, BytesCodec::new());
@@ -158,7 +118,7 @@ impl Snapshotter {
             .await
             .map_err(|e| SnapshotError::UploadFileError(e.to_string()))?;
 
-        info!(self.logger, "uploaded {}", filename);
+        info!("uploaded {}", filename);
 
         // ensure the uploaded file as public read access
         // when a file is uploaded to gcloud storage its permissions are overwritten so
@@ -169,8 +129,8 @@ impl Snapshotter {
         };
 
         info!(
-            self.logger,
-            "updating acl for {}: {:?}", filename, new_bucket_access_control
+            "updating acl for {}: {:?}",
+            filename, new_bucket_access_control
         );
 
         client
@@ -179,7 +139,7 @@ impl Snapshotter {
             .await
             .map_err(|e| SnapshotError::UploadFileError(e.to_string()))?;
 
-        info!(self.logger, "updated acl for {} ", filename);
+        info!("updated acl for {} ", filename);
 
         Ok(())
     }