Reduce Usage of Compression Format Detection (#267)

We use a KL-divergence based predictor to select between the LZ4 and
BG-LZ4 schemes. This is done for every chunk and seems to consume about
20% CPU time. Instead, we only run this for the first Xorb, lock the
scheme and use it from then on.

- Removes compression scheme from RemoteClient construction and instead
as an argument to put (xorb)
 - FileUploadSession is used to track compression scheme
 - Other changes demanded by clippy.

Author: ylow

cas_client/src/interface.rs

@@ -5,6 +5,7 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use cas_object::CompressionScheme;
 use cas_types::{FileRange, QueryReconstructionResponse};
 use mdb_shard::shard_file_reconstructor::FileReconstructor;
 use merklehash::MerkleHash;
@@ -35,6 +36,7 @@ pub trait UploadClient {
         hash: &MerkleHash,
         data: Vec<u8>,
         chunk_and_boundaries: Vec<(MerkleHash, u32)>,
+        compression: Option<CompressionScheme>,
     ) -> Result<usize>;
 
     /// Check if a XORB already exists.

cas_client/src/local_client.rs

@@ -5,7 +5,7 @@ use std::sync::Arc;
 
 use anyhow::anyhow;
 use async_trait::async_trait;
-use cas_object::CasObject;
+use cas_object::{CasObject, CompressionScheme};
 use cas_types::{FileRange, Key};
 use file_utils::SafeFileCreator;
 use heed::types::*;
@@ -230,6 +230,7 @@ impl UploadClient for LocalClient {
         hash: &MerkleHash,
         data: Vec<u8>,
         chunk_and_boundaries: Vec<(MerkleHash, u32)>,
+        compression: Option<CompressionScheme>,
     ) -> Result<usize> {
         // no empty writes
         if chunk_and_boundaries.is_empty() || data.is_empty() {
@@ -252,13 +253,7 @@ impl UploadClient for LocalClient {
         info!("Writing XORB {hash:?} to local path {file_path:?}");
 
         let mut file = SafeFileCreator::new(&file_path)?;
-        let (_, bytes_written) = CasObject::serialize(
-            &mut file,
-            hash,
-            &data,
-            &chunk_and_boundaries,
-            Some(cas_object::CompressionScheme::None),
-        )?;
+        let (_, bytes_written) = CasObject::serialize(&mut file, hash, &data, &chunk_and_boundaries, compression)?;
         file.close()?;
 
         // attempt to set to readonly on unix.
@@ -451,7 +446,10 @@ mod tests {
 
         // Act & Assert
         let client = LocalClient::temporary().unwrap();
-        assert!(client.put("key", &hash, data, vec![(hash, chunk_boundaries)]).await.is_ok());
+        assert!(client
+            .put("key", &hash, data, vec![(hash, chunk_boundaries)], None)
+            .await
+            .is_ok());
 
         let returned_data = client.get(&hash).unwrap();
         assert_eq!(data_again, returned_data);
@@ -465,7 +463,7 @@ mod tests {
 
         // Act & Assert
         let client = LocalClient::temporary().unwrap();
-        assert!(client.put("", &c.info.cashash, data, chunk_boundaries).await.is_ok());
+        assert!(client.put("", &c.info.cashash, data, chunk_boundaries, None).await.is_ok());
 
         let returned_data = client.get(&c.info.cashash).unwrap();
         assert_eq!(data_again, returned_data);
@@ -479,14 +477,14 @@ mod tests {
         // Act & Assert
         let client = LocalClient::temporary().unwrap();
         assert!(client
-            .put("", &c.info.cashash, data.clone(), chunk_and_boundaries.clone())
+            .put("", &c.info.cashash, data.clone(), chunk_and_boundaries.clone(), None)
             .await
             .is_ok());
 
         let ranges: Vec<(u32, u32)> = vec![(0, 1), (2, 3)];
         let returned_ranges = client.get_object_range(&c.info.cashash, ranges).unwrap();
 
-        let expected = vec![
+        let expected = [
             data[0..chunk_and_boundaries[0].1 as usize].to_vec(),
             data[chunk_and_boundaries[1].1 as usize..chunk_and_boundaries[2].1 as usize].to_vec(),
         ];
@@ -504,7 +502,7 @@ mod tests {
 
         // Act
         let client = LocalClient::temporary().unwrap();
-        assert!(client.put("", &c.info.cashash, data, chunk_boundaries).await.is_ok());
+        assert!(client.put("", &c.info.cashash, data, chunk_boundaries, None).await.is_ok());
         let len = client.get_length(&c.info.cashash).unwrap();
 
         // Assert
@@ -530,13 +528,13 @@ mod tests {
         // write "hello world"
         let client = LocalClient::temporary().unwrap();
         client
-            .put("default", &hello_hash, hello.clone(), vec![(hello_hash, hello.len() as u32)])
+            .put("default", &hello_hash, hello.clone(), vec![(hello_hash, hello.len() as u32)], None)
             .await
             .unwrap();
 
         // put the same value a second time. This should be ok.
         client
-            .put("default", &hello_hash, hello.clone(), vec![(hello_hash, hello.len() as u32)])
+            .put("default", &hello_hash, hello.clone(), vec![(hello_hash, hello.len() as u32)], None)
             .await
             .unwrap();
 
@@ -555,7 +553,13 @@ mod tests {
         assert_eq!(
             CasClientError::InvalidArguments,
             client
-                .put("hellp2", &hello_hash, "hellp wod".as_bytes().to_vec(), vec![(hello_hash, hello.len() as u32)],)
+                .put(
+                    "hellp2",
+                    &hello_hash,
+                    "hellp wod".as_bytes().to_vec(),
+                    vec![(hello_hash, hello.len() as u32)],
+                    None
+                )
                 .await
                 .unwrap_err()
         );
@@ -569,6 +573,7 @@ mod tests {
                     &hello_hash,
                     "hello world again".as_bytes().to_vec(),
                     vec![(hello_hash, hello.len() as u32)],
+                    None
                 )
                 .await
                 .unwrap_err()
@@ -577,7 +582,7 @@ mod tests {
         // empty writes should fail
         assert_eq!(
             CasClientError::InvalidArguments,
-            client.put("key", &hello_hash, vec![], vec![],).await.unwrap_err()
+            client.put("key", &hello_hash, vec![], vec![], None).await.unwrap_err()
         );
 
         // compute a hash of something we do not have in the store
@@ -621,7 +626,7 @@ mod tests {
         // insert should succeed
         let client = LocalClient::temporary().unwrap();
         client
-            .put("key", &final_hash, "helloworld".as_bytes().to_vec(), vec![(hello_hash, 5), (world_hash, 10)])
+            .put("key", &final_hash, "helloworld".as_bytes().to_vec(), vec![(hello_hash, 5), (world_hash, 10)], None)
             .await
             .unwrap();
     }

cas_client/src/remote_client.rs

@@ -59,7 +59,6 @@ utils::configurable_bool_constants! {
 
 pub struct RemoteClient {
     endpoint: String,
-    compression: Option<CompressionScheme>,
     dry_run: bool,
     http_client: Arc<ClientWithMiddleware>,
     authenticated_http_client: Arc<ClientWithMiddleware>,
@@ -74,7 +73,6 @@ impl RemoteClient {
     pub fn new(
         threadpool: Arc<ThreadPool>,
         endpoint: &str,
-        compression: Option<CompressionScheme>,
         auth: &Option<AuthConfig>,
         cache_config: &Option<CacheConfig>,
         shard_cache_directory: PathBuf,
@@ -101,7 +99,6 @@ impl RemoteClient {
 
         Self {
             endpoint: endpoint.to_string(),
-            compression,
             dry_run,
             authenticated_http_client: Arc::new(
                 http_client::build_auth_http_client(auth, RetryConfig::default()).unwrap(),
@@ -126,13 +123,14 @@ impl UploadClient for RemoteClient {
         hash: &MerkleHash,
         data: Vec<u8>,
         chunk_and_boundaries: Vec<(MerkleHash, u32)>,
+        compression: Option<CompressionScheme>,
     ) -> Result<usize> {
         let key = Key {
             prefix: prefix.to_string(),
             hash: *hash,
         };
 
-        let (was_uploaded, nbytes_trans) = self.upload(&key, data, chunk_and_boundaries).await?;
+        let (was_uploaded, nbytes_trans) = self.upload(&key, data, chunk_and_boundaries, compression).await?;
 
         if !was_uploaded {
             debug!("{key:?} not inserted into CAS.");
@@ -281,13 +279,14 @@ impl RemoteClient {
         key: &Key,
         contents: Vec<u8>,
         chunk_and_boundaries: Vec<(MerkleHash, u32)>,
+        compression: Option<CompressionScheme>,
     ) -> Result<(bool, usize)> {
         let url = Url::parse(&format!("{}/xorb/{key}", self.endpoint))?;
 
         let mut writer = Cursor::new(Vec::new());
 
         let (_, nbytes_trans) =
-            CasObject::serialize(&mut writer, &key.hash, &contents, &chunk_and_boundaries, self.compression)?;
+            CasObject::serialize(&mut writer, &key.hash, &contents, &chunk_and_boundaries, compression)?;
         // free memory before the "slow" network transfer below
         drop(contents);
 
@@ -754,18 +753,14 @@ mod tests {
         let (c, _, data, chunk_boundaries) = build_cas_object(3, ChunkSize::Random(512, 10248), CompressionScheme::LZ4);
 
         let threadpool = Arc::new(ThreadPool::new().unwrap());
-        let client = RemoteClient::new(
-            threadpool.clone(),
-            CAS_ENDPOINT,
-            Some(CompressionScheme::LZ4),
-            &None,
-            &None,
-            "".into(),
-            false,
-        );
+        let client = RemoteClient::new(threadpool.clone(), CAS_ENDPOINT, &None, &None, "".into(), false);
         // Act
         let result = threadpool
-            .external_run_async_task(async move { client.put(prefix, &c.info.cashash, data, chunk_boundaries).await })
+            .external_run_async_task(async move {
+                client
+                    .put(prefix, &c.info.cashash, data, chunk_boundaries, Some(CompressionScheme::LZ4))
+                    .await
+            })
             .unwrap();
 
         // Assert
@@ -1093,7 +1088,7 @@ mod tests {
             file_range: FileRange::new(SKIP_BYTES, FILE_SIZE - SKIP_BYTES),
             expected_data: [
                 &raw_data[SKIP_BYTES as usize..(5 * CHUNK_SIZE) as usize],
-                &raw_data[(6 * CHUNK_SIZE) as usize as usize..(NUM_CHUNKS * CHUNK_SIZE) as usize - SKIP_BYTES as usize],
+                &raw_data[(6 * CHUNK_SIZE) as usize..(NUM_CHUNKS * CHUNK_SIZE) as usize - SKIP_BYTES as usize],
             ]
             .concat(),
             expect_error: false,
@@ -1132,7 +1127,7 @@ mod tests {
 
         // test reconstruct and sequential write
         let test = test_case.clone();
-        let client = RemoteClient::new(threadpool.clone(), endpoint, None, &None, &None, "".into(), false);
+        let client = RemoteClient::new(threadpool.clone(), endpoint, &None, &None, "".into(), false);
         let provider = BufferProvider::default();
         let buf = provider.buf.clone();
         let writer = OutputProvider::Buffer(provider);
@@ -1150,7 +1145,7 @@ mod tests {
 
         // test reconstruct and parallel write
         let test = test_case;
-        let client = RemoteClient::new(threadpool.clone(), endpoint, None, &None, &None, "".into(), false);
+        let client = RemoteClient::new(threadpool.clone(), endpoint, &None, &None, "".into(), false);
         let provider = BufferProvider::default();
         let buf = provider.buf.clone();
         let writer = OutputProvider::Buffer(provider);

cas_object/src/cas_chunk_format.rs

@@ -213,7 +213,6 @@ mod tests {
     use std::io::Cursor;
 
     use rand::Rng;
-    use CompressionScheme;
 
     use super::*;
 

cas_object/src/cas_object_format.rs

@@ -1856,7 +1856,7 @@ mod tests {
             + size_of::<CasObjectIdent>()
             + size_of::<u8>();
 
-        let chunks = xorb_bytes[start_pos..].chunks(10).map(|c| Ok(c)).collect::<Vec<_>>();
+        let chunks = xorb_bytes[start_pos..].chunks(10).map(Ok).collect::<Vec<_>>();
         let mut xorb_footer_async_reader = futures::stream::iter(chunks).into_async_read();
         let cas_object_result =
             CasObject::deserialize_async(&mut xorb_footer_async_reader, CAS_OBJECT_FORMAT_VERSION).await;

cas_object/src/compression_scheme.rs

@@ -23,6 +23,7 @@ pub enum CompressionScheme {
     LZ4 = 1,
     ByteGrouping4LZ4 = 2, // 4 byte groups
 }
+pub const NUM_COMPRESSION_SCHEMES: usize = 3;
 
 impl Display for CompressionScheme {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
@@ -288,47 +289,39 @@ mod tests {
             let random_u8s: Vec<_> = (0..n).map(|_| rng.gen_range(0..255)).collect();
             let random_f32s_ng1_1: Vec<_> = (0..n / size_of::<f32>())
                 .map(|_| rng.gen_range(-1.0f32..=1.0))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
             let random_f32s_0_2: Vec<_> = (0..n / size_of::<f32>())
                 .map(|_| rng.gen_range(0f32..=2.0))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
             let random_f64s_ng1_1: Vec<_> = (0..n / size_of::<f64>())
                 .map(|_| rng.gen_range(-1.0f64..=1.0))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
             let random_f64s_0_2: Vec<_> = (0..n / size_of::<f64>())
                 .map(|_| rng.gen_range(0f64..=2.0))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
 
             // f16, a.k.a binary16 format: sign (1 bit), exponent (5 bit), mantissa (10 bit)
             let random_f16s_ng1_1: Vec<_> = (0..n / size_of::<f16>())
                 .map(|_| f16::from_f32(rng.gen_range(-1.0f32..=1.0)))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
             let random_f16s_0_2: Vec<_> = (0..n / size_of::<f16>())
                 .map(|_| f16::from_f32(rng.gen_range(0f32..=2.0)))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
 
             // bf16 format: sign (1 bit), exponent (8 bit), mantissa (7 bit)
             let random_bf16s_ng1_1: Vec<_> = (0..n / size_of::<bf16>())
                 .map(|_| bf16::from_f32(rng.gen_range(-1.0f32..=1.0)))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
             let random_bf16s_0_2: Vec<_> = (0..n / size_of::<bf16>())
                 .map(|_| bf16::from_f32(rng.gen_range(0f32..=2.0)))
-                .map(|f| f.to_le_bytes())
-                .flatten()
+                .flat_map(|f| f.to_le_bytes())
                 .collect();
 
             let dataset = [