change TinyDancer to ::start() and add comments to sampler

Author: 0xNineteen

tinydancer/src/main.rs

@@ -157,8 +157,8 @@ async fn main() -> Result<()> {
                     })
                 },
             };
-            let client = TinyDancer::new(config).await;
-            client.join().await;
+
+            TinyDancer::start(config).await.unwrap();
         }
 
         Commands::Slot => {

tinydancer/src/rpc_wrapper/mod.rs

@@ -155,6 +155,7 @@ impl ClientService<TransactionServiceConfig> for TransactionService {
             tx_handle: transaction_handle,
         }
     }
+
     async fn join(self) -> std::result::Result<(), Self::ServiceError> {
         let _ = self.tx_handle.await;
         Ok(())

tinydancer/src/sampler.rs

@@ -43,7 +43,9 @@ use tokio::{
 };
 use tungstenite::{connect, Message};
 use url::Url;
+
 pub const SHRED_CF: &str = "archived_shreds";
+
 pub struct SampleService {
     sample_indices: Vec<u64>,
     // peers: Vec<(Pubkey, SocketAddr)>,
@@ -59,63 +61,78 @@ pub struct SampleServiceConfig {
 #[derive(Clone, Debug)]
 pub struct ArchiveConfig {
     pub shred_archive_duration: u64,
-
     pub archive_path: String,
 }
+
 #[async_trait]
 impl ClientService<SampleServiceConfig> for SampleService {
     type ServiceError = tokio::task::JoinError;
+
     fn new(config: SampleServiceConfig) -> Self {
         let sampler_handle = tokio::spawn(async move {
             let rpc_url = endpoint(config.cluster);
             let pub_sub = convert_to_websocket!(rpc_url);
+
             let mut threads = Vec::default();
 
             let (slot_update_tx, slot_update_rx) = crossbeam::channel::unbounded::<u64>();
             let (shred_tx, shred_rx) = crossbeam::channel::unbounded();
             let (verified_shred_tx, verified_shred_rx) = crossbeam::channel::unbounded();
-            let status_arc = Arc::clone(&config.status_sampler);
+
+            let status_arc = config.status_sampler.clone();
+
+            // waits on new slots => triggers shred_update_loop
             threads.push(tokio::spawn(slot_update_loop(
                 slot_update_tx,
                 pub_sub,
                 config.status_sampler,
             )));
+
+            // sample shreds from new slot
+            // verify each shred in shred_verify_loop
             threads.push(tokio::spawn(shred_update_loop(
                 slot_update_rx,
                 rpc_url,
                 shred_tx,
                 status_arc,
             )));
 
+            // verify shreds + store in db in shred_archiver
             threads.push(tokio::spawn(shred_verify_loop(shred_rx, verified_shred_tx)));
+
             if let Some(archive_config) = config.archive_config {
                 threads.push(tokio::spawn(shred_archiver(
                     verified_shred_rx,
                     archive_config,
                     config.instance,
                 )));
             }
+
             for thread in threads {
-                thread.await;
+                thread.await.unwrap();
             }
         });
+
         let sample_indices: Vec<u64> = Vec::default();
         Self {
             sampler_handle,
             sample_indices,
         }
     }
+
     async fn join(self) -> std::result::Result<(), Self::ServiceError> {
         self.sampler_handle.await
     }
 }
+
 pub fn gen_random_indices(max_shreds_per_slot: usize, sample_qty: usize) -> Vec<usize> {
     let mut rng = StdRng::from_entropy();
-    let vec = (0..max_shreds_per_slot)
+    let vec = (0..sample_qty)
         .map(|_| rng.gen_range(0..max_shreds_per_slot))
         .collect::<Vec<usize>>();
-    vec.as_slice()[0..sample_qty].to_vec()
+    vec
 }
+
 pub async fn request_shreds(
     slot: usize,
     indices: Vec<usize>,
@@ -205,6 +222,7 @@ async fn shred_update_loop(
         }
 
         if let Ok(slot) = slot_update_rx.recv() {
+            // get shred length
             let shred_for_one = request_shreds(slot as usize, vec![0], endpoint.clone()).await;
             // info!("res {:?}", shred_for_one);
             let shred_indices_for_slot = match shred_for_one {
@@ -251,6 +269,8 @@ async fn shred_update_loop(
                     None
                 }
             };
+
+            // get a random sample of shreds
             info!("indices of: {:?} {:?}", shred_indices_for_slot, slot);
             if let Some(shred_indices_for_slot) = shred_indices_for_slot.clone() {
                 let shreds_for_slot = request_shreds(
@@ -314,12 +334,14 @@ async fn shred_update_loop(
         }
     }
 }
+
 // use solana_ledger::shred::dispatch;
+
+// verifies the merkle proof of the shread 
 pub fn verify_sample(shred: &Shred, leader: solana_ledger::shred::Pubkey) -> bool {
     // @TODO fix error handling here
     let verify_merkle_root = match shred {
         Shred::ShredData(ShredData::Merkle(shred)) => Some(shred.verify_merkle_proof()),
-
         Shred::ShredCode(ShredCode::Merkle(shred)) => Some(shred.verify_merkle_proof()),
         _ => None,
     };
@@ -338,6 +360,7 @@ pub fn verify_sample(shred: &Shred, leader: solana_ledger::shred::Pubkey) -> boo
     .all(|s| *s);
     verified
 }
+
 pub async fn shred_verify_loop(
     shred_rx: Receiver<(Vec<Option<Shred>>, solana_ledger::shred::Pubkey)>,
     verified_shred_tx: Sender<(Shred, solana_ledger::shred::Pubkey)>,
@@ -373,6 +396,9 @@ pub async fn shred_verify_loop(
         }
     }
 }
+
+
+// store verified shreds in db
 pub async fn shred_archiver(
     verified_shred_rx: Receiver<(Shred, solana_ledger::shred::Pubkey)>,
     _archive_config: ArchiveConfig,
@@ -413,6 +439,8 @@ pub async fn shred_archiver(
         }
     }
 }
+
+
 pub async fn pull_and_verify_shreds(slot: usize, endpoint: String) -> bool {
     let shred_for_one = request_shreds(slot, vec![0], endpoint.clone()).await;
     // info!("res {:?}", shred_for_one);
@@ -522,6 +550,7 @@ pub async fn pull_and_verify_shreds(slot: usize, endpoint: String) -> bool {
         false
     }
 }
+
 pub fn put_serialized<T: serde::Serialize + std::fmt::Debug>(
     instance: &rocksdb::DB,
     cf: &ColumnFamily,

tinydancer/src/tinydancer.rs

@@ -15,6 +15,8 @@ use crate::{
     ui::{UiConfig, UiService},
 };
 use async_trait::async_trait;
+use futures::future::join_all;
+use rand::seq::index::sample;
 use tiny_logger::logs::info;
 // use log::info;
 // use log4rs;
@@ -29,13 +31,15 @@ pub trait ClientService<T> {
     fn new(config: T) -> Self;
     async fn join(self) -> std::result::Result<(), Self::ServiceError>;
 }
+
 pub struct TinyDancer {
     sample_service: SampleService,
     ui_service: Option<UiService>,
     sample_qty: u64,
     config: TinyDancerConfig,
     transaction_service: TransactionService,
 }
+
 #[derive(Clone)]
 pub struct TinyDancerConfig {
     pub rpc_endpoint: Cluster,
@@ -52,6 +56,7 @@ use std::fs::read_dir;
 use std::io;
 use std::io::ErrorKind;
 use std::path::PathBuf;
+
 // use tiny_logger::logs::info;
 pub fn get_project_root() -> io::Result<PathBuf> {
     let path = env::current_dir()?;
@@ -71,11 +76,13 @@ pub fn get_project_root() -> io::Result<PathBuf> {
         "Ran out of places to find Cargo.toml",
     ))
 }
+
 impl TinyDancer {
-    pub async fn new(config: TinyDancerConfig) -> Self {
+    pub async fn start(config: TinyDancerConfig) -> Result<()> {
         let status = ClientStatus::Initializing(String::from("Starting Up Tinydancer"));
+
         let client_status = Arc::new(Mutex::new(status));
-        let status_sampler = Arc::clone(&client_status);
+        let status_sampler = client_status.clone();
 
         let TinyDancerConfig {
             enable_ui_service,
@@ -87,33 +94,34 @@ impl TinyDancer {
         } = config.clone();
         std::env::set_var("RUST_LOG", "info");
         tiny_logger::setup_file_with_default(&log_path, "RUST_LOG");
-        let rpc_cluster = rpc_endpoint.clone();
 
         let mut opts = rocksdb::Options::default();
         opts.create_if_missing(true);
         opts.set_error_if_exists(false);
         opts.create_missing_column_families(true);
-        let instance = rocksdb::DB::open_cf(
+
+        // setup db
+        let db = rocksdb::DB::open_cf(
             &opts,
             archive_config.clone().unwrap().archive_path,
             vec![SHRED_CF],
         )
         .unwrap();
-        let instance = Arc::new(instance);
-        let sample_instance = Arc::clone(&instance);
+        let db = Arc::new(db);
 
-        let rpc_instance = Arc::clone(&instance);
         let sample_service_config = SampleServiceConfig {
-            cluster: rpc_endpoint,
+            cluster: rpc_endpoint.clone(),
             archive_config,
-            instance: sample_instance,
+            instance: db.clone(),
             status_sampler,
         };
         let sample_service = SampleService::new(sample_service_config);
+
         let transaction_service = TransactionService::new(TransactionServiceConfig {
-            cluster: rpc_cluster,
-            db_instance: rpc_instance,
+            cluster: rpc_endpoint.clone(),
+            db_instance: db.clone(),
         });
+
         let ui_service = if enable_ui_service {
             Some(UiService::new(UiConfig {
                 client_status,
@@ -130,26 +138,22 @@ impl TinyDancer {
             None
         };
 
-        Self {
-            config,
-            ui_service,
-            sample_qty,
-            sample_service,
-            transaction_service,
-        }
-    }
-    pub async fn join(self) {
-        self.sample_service
+        // run
+        sample_service
             .join()
             .await
             .expect("error in sample service thread");
-        self.transaction_service
+
+        transaction_service
             .join()
             .await
             .expect("ERROR IN SIMPLE PAYMENT SERVICE");
-        if let Some(ui_service) = self.ui_service {
+
+        if let Some(ui_service) = ui_service {
             block_on!(async { ui_service.join().await }, "Ui Service Error");
         }
+
+        Ok(())
     }
 }
 
@@ -161,6 +165,7 @@ pub enum Cluster {
     Localnet,
     Custom(String),
 }
+
 pub fn endpoint(cluster: Cluster) -> String {
     let cluster = cluster;
     match cluster {