Collect telemetry for upload and download path

Get the telemetry data to the py layer. This will be sent over the
telemetry endpoint later.

Related to STO-7

Author: Ajit Banerjee

.gitignore

@@ -17,3 +17,4 @@ debug/
 .vscode
 venv
 **/*.env
+**/uv.lock

cas_client/src/caching_client.rs

@@ -36,7 +36,7 @@ impl<T: Client + Reconstructable + Send + Sync> UploadClient for CachingClient<T
 
 #[async_trait]
 impl<T: Client + Reconstructable + Send + Sync> ReconstructionClient for CachingClient<T> {
-    async fn get_file(&self, hash: &MerkleHash, writer: &mut Box<dyn Write + Send>) -> Result<()> {
+    async fn get_file(&self, hash: &MerkleHash, writer: &mut Box<dyn Write + Send>) -> Result<u64> {
         /*
         let file_info = self.reconstruct(hash, None).await?;
 
@@ -60,7 +60,7 @@ impl<T: Client + Reconstructable + Send + Sync> ReconstructionClient for Caching
         offset: u64,
         length: u64,
         writer: &mut Box<dyn Write + Send>,
-    ) -> Result<()> {
+    ) -> Result<u64> {
         todo!()
     }
 }

cas_client/src/interface.rs

@@ -42,7 +42,7 @@ pub trait UploadClient {
 #[async_trait]
 pub trait ReconstructionClient {
     /// Get a entire file by file hash.
-    async fn get_file(&self, hash: &MerkleHash, writer: &mut Box<dyn Write + Send>) -> Result<()>;
+    async fn get_file(&self, hash: &MerkleHash, writer: &mut Box<dyn Write + Send>) -> Result<u64>;
 
     /// Get a entire file by file hash at a specific bytes range.
     async fn get_file_byte_range(
@@ -51,7 +51,7 @@ pub trait ReconstructionClient {
         offset: u64,
         length: u64,
         writer: &mut Box<dyn Write + Send>,
-    ) -> Result<()>;
+    ) -> Result<u64>;
 }
 
 pub trait Client: UploadClient + ReconstructionClient {}

cas_client/src/remote_client.rs

@@ -102,13 +102,13 @@ impl UploadClient for RemoteClient {
 
 #[async_trait]
 impl ReconstructionClient for RemoteClient {
-    async fn get_file(&self, hash: &MerkleHash, writer: &mut Box<dyn Write + Send>) -> Result<()> {
+    async fn get_file(&self, hash: &MerkleHash, writer: &mut Box<dyn Write + Send>) -> Result<u64> {
         // get manifest of xorbs to download
         let manifest = self.reconstruct(hash, None).await?;
 
-        self.get_ranges(manifest, None, writer).await?;
+        let bytes_downloaded = self.get_ranges(manifest, None, writer).await?;
 
-        Ok(())
+        Ok(bytes_downloaded)
     }
 
     #[allow(unused_variables)]
@@ -118,7 +118,7 @@ impl ReconstructionClient for RemoteClient {
         offset: u64,
         length: u64,
         writer: &mut Box<dyn Write + Send>,
-    ) -> Result<()> {
+    ) -> Result<u64> {
         todo!()
     }
 }
@@ -203,7 +203,7 @@ impl RemoteClient {
         reconstruction_response: QueryReconstructionResponse,
         _byte_range: Option<(u64, u64)>,
         writer: &mut Box<dyn Write + Send>,
-    ) -> Result<usize> {
+    ) -> Result<u64> {
         let info = reconstruction_response.reconstruction;
         let total_len = info.iter().fold(0, |acc, x| acc + x.unpacked_length);
         let futs = info
@@ -215,7 +215,8 @@ impl RemoteClient {
                 .map_err(|e| CasClientError::InternalError(anyhow!("join error {e}")))??;
             writer.write_all(&piece)?;
         }
-        Ok(total_len as usize)
+        // Todo: return the bytes which were read from the cache for telemetry
+        Ok(total_len as u64)
     }
 }
 

data/src/bin/example.rs

@@ -213,7 +213,7 @@ async fn smudge(mut reader: impl Read, writer: &mut Box<dyn Write + Send>) -> Re
     let mut input = String::new();
     reader.read_to_string(&mut input)?;
 
-    let pointer_file = PointerFile::init_from_string(&input, "");
+    let mut pointer_file = PointerFile::init_from_string(&input, "");
 
     // not a pointer file, leave it as it is.
     if !pointer_file.is_valid() {
@@ -223,7 +223,7 @@ async fn smudge(mut reader: impl Read, writer: &mut Box<dyn Write + Send>) -> Re
     let translator = PointerFileTranslator::new(default_smudge_config()?).await?;
 
     translator
-        .smudge_file_from_pointer(&pointer_file, writer, None)
+        .smudge_file_from_pointer(&mut pointer_file, writer, None)
         .await?;
 
     Ok(())

data/src/clean.rs

@@ -21,6 +21,7 @@ use std::mem::take;
 use std::ops::DerefMut;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
+use std::time::Instant;
 use tokio::sync::mpsc::error::TryRecvError;
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 use tokio::sync::Mutex;
@@ -76,6 +77,9 @@ pub struct Cleaner {
 
     // Auxiliary info
     file_name: Option<PathBuf>,
+
+    // Telemetry
+    start: Instant,
 }
 
 impl Cleaner {
@@ -113,6 +117,7 @@ impl Cleaner {
             tracking_info: Mutex::new(Default::default()),
             small_file_buffer: Mutex::new(Some(Vec::with_capacity(small_file_threshold))),
             file_name: file_name.map(|f| f.to_owned()),
+            start: Instant::now(),
         });
 
         Self::run(cleaner.clone(), chunk_c).await;
@@ -239,8 +244,9 @@ impl Cleaner {
         Ok(false)
     }
 
-    async fn dedup(&self, chunks: &[ChunkYieldType]) -> Result<()> {
+    async fn dedup(&self, chunks: &[ChunkYieldType]) -> Result<u64> {
         info!("Dedup {} chunks", chunks.len());
+        let mut total_compressed_bytes = 0;
         let mut tracking_info = self.tracking_info.lock().await;
 
         let enable_global_dedup = self.enable_global_dedup_queries;
@@ -463,13 +469,14 @@ impl Cleaner {
                     tracking_info.cas_data.data.extend(bytes);
 
                     if tracking_info.cas_data.data.len() > TARGET_CAS_BLOCK_SIZE {
-                        let cas_hash = register_new_cas_block(
+                        let (cas_hash, compressed_bytes) = register_new_cas_block(
                             &mut tracking_info.cas_data,
                             &self.shard_manager,
                             &self.cas,
                             &self.cas_prefix,
                         )
                         .await?;
+                        total_compressed_bytes += compressed_bytes;
 
                         for i in take(&mut tracking_info.current_cas_file_info_indices) {
                             tracking_info.file_info[i].cas_hash = cas_hash;
@@ -483,7 +490,7 @@ impl Cleaner {
             }
         }
 
-        Ok(())
+        Ok(total_compressed_bytes)
     }
 
     async fn finish(&self) -> Result<()> {
@@ -516,7 +523,8 @@ impl Cleaner {
         Ok(())
     }
 
-    async fn summarize_dedup_info(&self) -> Result<(MerkleHash, u64)> {
+    async fn summarize_dedup_info(&self) -> Result<(MerkleHash, u64, u64)> {
+        let mut total_compressed_bytes = 0;
         let mut tracking_info = self.tracking_info.lock().await;
 
         let file_hash = file_node_hash(
@@ -577,13 +585,14 @@ impl Cleaner {
             if cas_data_accumulator.data.len() >= TARGET_CAS_BLOCK_SIZE {
                 let mut new_cas_data = take(cas_data_accumulator.deref_mut());
                 drop(cas_data_accumulator); // Release the lock.
-                register_new_cas_block(
+                let (_cas_hash, compressed_bytes) = register_new_cas_block(
                     &mut new_cas_data,
                     &self.shard_manager,
                     &self.cas,
                     &self.cas_prefix,
                 )
                 .await?;
+                total_compressed_bytes += compressed_bytes;
             } else {
                 drop(cas_data_accumulator);
             }
@@ -593,11 +602,11 @@ impl Cleaner {
 
         *tracking_info = Default::default();
 
-        Ok((file_hash, file_size))
+        Ok((file_hash, file_size, total_compressed_bytes))
     }
 
     async fn to_pointer_file(&self) -> Result<String> {
-        let (hash, filesize) = self.summarize_dedup_info().await?;
+        let (hash, filesize, compressed_size) = self.summarize_dedup_info().await?;
         let pointer_file = PointerFile::init_from_info(
             &self
                 .file_name
@@ -606,6 +615,8 @@ impl Cleaner {
                 .unwrap_or_default(),
             &hash.hex(),
             filesize,
+            compressed_size,
+            self.start.elapsed(),
         );
         Ok(pointer_file.to_string())
     }

data/src/data_processing.rs

@@ -5,7 +5,7 @@ use crate::errors::*;
 use crate::metrics::FILTER_CAS_BYTES_PRODUCED;
 use crate::remote_shard_interface::RemoteShardInterface;
 use crate::shard_interface::create_shard_manager;
-use crate::PointerFile;
+use crate::{PointerFile, PointerFileTelemetry};
 use cas_client::Client;
 use mdb_shard::cas_structs::{CASChunkSequenceEntry, CASChunkSequenceHeader, MDBCASInfo};
 use mdb_shard::file_structs::MDBFileInfo;
@@ -17,6 +17,7 @@ use std::mem::take;
 use std::ops::DerefMut;
 use std::path::Path;
 use std::sync::Arc;
+use std::time::Instant;
 use tokio::sync::Mutex;
 
 #[derive(Default, Debug)]
@@ -59,6 +60,9 @@ pub struct PointerFileTranslator {
 
     /* ----- Deduped data shared across files ----- */
     global_cas_data: Arc<Mutex<CASDataAggregator>>,
+    // Telemetry
+    /* ----- Telemetry ----- */
+    pub start: Instant,
 }
 
 // Constructors
@@ -97,6 +101,7 @@ impl PointerFileTranslator {
             remote_shards,
             cas: cas_client,
             global_cas_data: Default::default(),
+            start: Instant::now(),
         })
     }
 }
@@ -208,7 +213,7 @@ pub(crate) async fn register_new_cas_block(
     shard_manager: &Arc<ShardFileManager>,
     cas: &Arc<dyn Client + Send + Sync>,
     cas_prefix: &str,
-) -> Result<MerkleHash> {
+) -> Result<(MerkleHash, u64)> {
     let cas_hash = cas_node_hash(&cas_data.chunks[..]);
 
     let raw_bytes_len = cas_data.data.len();
@@ -283,29 +288,37 @@ pub(crate) async fn register_new_cas_block(
     cas_data.chunks.clear();
     cas_data.pending_file_info.clear();
 
-    Ok(cas_hash)
+    Ok((cas_hash, compressed_bytes_len as u64))
 }
 
 /// Smudge operations
 impl PointerFileTranslator {
     pub async fn smudge_file_from_pointer(
         &self,
-        pointer: &PointerFile,
+        pointer_file: &mut PointerFile,
         writer: &mut Box<dyn Write + Send>,
         range: Option<(usize, usize)>,
     ) -> Result<()> {
-        self.smudge_file_from_hash(&pointer.hash()?, writer, range)
-            .await
+        let start = Instant::now();
+        let bytes_downloaded = self
+            .smudge_file_from_hash(&pointer_file.hash()?, writer, range)
+            .await?;
+
+        pointer_file.telemetry = Some(PointerFileTelemetry {
+            latency: Some(start.elapsed()),
+            network_bytes: Some(bytes_downloaded),
+        });
+
+        Ok(())
     }
 
     pub async fn smudge_file_from_hash(
         &self,
         file_id: &MerkleHash,
         writer: &mut Box<dyn Write + Send>,
         _range: Option<(usize, usize)>,
-    ) -> Result<()> {
-        self.cas.get_file(file_id, writer).await?;
-
-        Ok(())
+    ) -> Result<u64> {
+        let bytes_downloaded = self.cas.get_file(file_id, writer).await?;
+        Ok(bytes_downloaded)
     }
 }

data/src/lib.rs

@@ -17,4 +17,4 @@ mod test_utils;
 
 pub use constants::SMALL_FILE_THRESHOLD;
 pub use data_processing::PointerFileTranslator;
-pub use pointer_file::PointerFile;
+pub use pointer_file::{PointerFile, PointerFileTelemetry};