feat(data): introduce streaming client

Cargo.toml

@@ -39,6 +39,7 @@ anyhow = "1"
 axum = "0.8"
 async-trait = "0.1"
 base64 = "0.22"
+futures-core = "0.3"
 bincode = "1.3"
 bitflags = { version = "2.10", features = ["serde"] }
 blake3 = "1.8"

data/Cargo.toml

@@ -40,6 +40,7 @@ progress_tracking = { path = "../progress_tracking" }
 utils = { path = "../utils" }
 xet_runtime = { path = "../xet_runtime" }
 
+futures = { workspace = true }
 anyhow = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }

data/src/bin/example.rs

@@ -91,7 +91,7 @@ async fn clean(mut reader: impl Read, mut writer: impl Write, size: u64) -> Resu
         FileUploadSession::new(TranslatorConfig::local_config(std::env::current_dir()?)?.into(), None).await?;
 
     let mut size_read = 0;
-    let mut handle = translator.start_clean(None, size, None).await;
+    let mut handle = translator.start_clean(None, Some(size), None).await;
 
     loop {
         let bytes = reader.read(&mut read_buf)?;

data/src/data_client.rs

@@ -112,9 +112,10 @@ pub async fn upload_bytes_async(
 
     let semaphore = XetRuntime::current().common().file_ingestion_semaphore.clone();
     let upload_session = FileUploadSession::new(config.into(), progress_updater).await?;
-    let clean_futures = file_contents.into_iter().map(|blob| {
+    let clean_futures = file_contents.into_iter().enumerate().map(|(i, blob)| {
         let upload_session = upload_session.clone();
-        async move { clean_bytes(upload_session, blob).await.map(|(xf, _metrics)| xf) }
+        let name: Arc<str> = format!("{i}").into();
+        async move { clean_bytes(upload_session, blob, Some(name)).await.map(|(xf, _metrics)| xf) }
             .instrument(info_span!("clean_task"))
     });
     let files = run_constrained_with_semaphore(clean_futures, semaphore).await?;
@@ -239,8 +240,9 @@ pub async fn download_async(
 pub async fn clean_bytes(
     processor: Arc<FileUploadSession>,
     bytes: Vec<u8>,
+    name: Option<Arc<str>>,
 ) -> errors::Result<(XetFileInfo, DeduplicationMetrics)> {
-    let mut handle = processor.start_clean(None, bytes.len() as u64, None).await;
+    let mut handle = processor.start_clean(name, Some(bytes.len() as u64), None).await;
     handle.add_data(&bytes).await?;
     handle.finish().await
 }
@@ -261,7 +263,11 @@ pub async fn clean_file(
     let mut buffer = vec![0u8; u64::min(filesize, *xet_config().data.ingestion_block_size) as usize];
 
     let mut handle = processor
-        .start_clean(Some(filename.as_ref().to_string_lossy().into()), filesize, Sha256::from_hex(sha256.as_ref()).ok())
+        .start_clean(
+            Some(filename.as_ref().to_string_lossy().into()),
+            Some(filesize),
+            Sha256::from_hex(sha256.as_ref()).ok(),
+        )
         .await;
 
     loop {
@@ -615,4 +621,43 @@ mod tests {
         assert_eq!(file_info1.file_size(), file_info2.file_size());
         assert_eq!(file_info1.file_size(), file_info3.file_size());
     }
+
+    #[tokio::test]
+    async fn test_upload_bytes() {
+        let temp_dir = tempdir().unwrap();
+        let endpoint = format!("local://{}", temp_dir.path().display());
+
+        let contents: Vec<Vec<u8>> = vec![
+            vec![],
+            b"Hello, World!".to_vec(),
+            (0..1_000_000).map(|i| (i % 256) as u8).collect(),
+        ];
+
+        // Upload all as bytes.
+        let file_infos = upload_bytes_async(contents.clone(), Some(endpoint.clone()), None, None, None, "test".into())
+            .await
+            .unwrap();
+
+        assert_eq!(file_infos.len(), contents.len());
+        for (info, content) in file_infos.iter().zip(&contents) {
+            assert_eq!(info.file_size(), content.len() as u64);
+            assert!(!info.hash().is_empty());
+        }
+        // Different contents produce different hashes.
+        assert_ne!(file_infos[0].hash(), file_infos[2].hash());
+
+        // Download to files and verify.
+        let download_dir = tempdir().unwrap();
+        let file_infos_with_paths: Vec<_> = file_infos
+            .into_iter()
+            .enumerate()
+            .map(|(i, info)| (info, download_dir.path().join(format!("{i}")).to_str().unwrap().to_string()))
+            .collect();
+        let paths = download_async(file_infos_with_paths, Some(endpoint), None, None, None, "test".into())
+            .await
+            .unwrap();
+        for (path, expected) in paths.iter().zip(&contents) {
+            assert_eq!(std::fs::read(path).unwrap(), *expected);
+        }
+    }
 }

data/src/file_cleaner.rs

@@ -153,7 +153,7 @@ impl SingleFileCleaner {
         let (file_hash, remaining_file_data, deduplication_metrics) =
             self.dedup_manager_fut.await?.finalize(Some(metadata_ext));
 
-        let file_info = XetFileInfo::new(file_hash.hex(), deduplication_metrics.total_bytes);
+        let file_info = XetFileInfo::with_sha256(file_hash.hex(), deduplication_metrics.total_bytes, sha256.hex());
 
         // Let's check some things that should be invariants
         #[cfg(debug_assertions)]

data/src/file_upload_session.rs

@@ -156,7 +156,10 @@ impl FileUploadSession {
 
             // Get a new file id for the completion tracking.  This also registers the size against the total bytes
             // of the file.
-            let file_id = self.completion_tracker.register_new_file(file_name.clone(), file_size).await;
+            let file_id = self
+                .completion_tracker
+                .register_new_file(file_name.clone(), Some(file_size))
+                .await;
 
             // Now, spawn a task
             let ingestion_concurrency_limiter = XetRuntime::current().common().file_ingestion_semaphore.clone();
@@ -253,7 +256,7 @@ impl FileUploadSession {
     pub async fn start_clean(
         self: &Arc<Self>,
         file_name: Option<Arc<str>>,
-        size: u64,
+        size: Option<u64>,
         sha256: Option<Sha256>,
     ) -> SingleFileCleaner {
         // Get a new file id for the completion tracking
@@ -564,7 +567,7 @@ mod tests {
             .unwrap();
 
         let mut cleaner = upload_session
-            .start_clean(Some("test".into()), read_data.len() as u64, None)
+            .start_clean(Some("test".into()), Some(read_data.len() as u64), None)
             .await;
 
         // Read blocks from the source file and hand them to the cleaning handle

data/src/lib.rs

@@ -10,6 +10,7 @@ mod prometheus_metrics;
 mod remote_client_interface;
 mod sha256;
 mod shard_interface;
+pub mod streaming;
 mod xet_file;
 
 // Reexport this one for now