feat(data): introduce a streaming client

Author: kszucs

Cargo.lock

@@ -37,9 +37,9 @@ debug = 1
 [workspace.dependencies]
 anyhow = "1"
 axum = "0.8"
-async-stream = "0.3"
 async-trait = "0.1"
 base64 = "0.22"
+futures-core = "0.3"
 bincode = "1.3"
 bitflags = { version = "2.9", features = ["serde"] }
 blake3 = "1.5"

Cargo.toml

@@ -42,7 +42,6 @@ xet_runtime = { path = "../xet_runtime" }
 
 futures = { workspace = true }
 anyhow = { workspace = true }
-async-stream = { workspace = true }
 async-trait = { workspace = true }
 bytes = { workspace = true }
 chrono = { workspace = true }

data/Cargo.toml

@@ -3,14 +3,11 @@ use std::io::Read;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
-use async_stream::try_stream;
 use bytes::Bytes;
 use cas_client::remote_client::PREFIX_DEFAULT;
 use cas_object::CompressionScheme;
-use cas_types::FileRange;
 use deduplication::{Chunker, DeduplicationMetrics};
 use file_reconstruction::DataOutput;
-use futures::Stream;
 use lazy_static::lazy_static;
 use mdb_shard::Sha256;
 use merklehash::MerkleHash;
@@ -136,75 +133,6 @@ pub async fn upload_bytes_async(
     Ok(files)
 }
 
-/// Downloads multiple files and returns their contents as byte streams.
-///
-/// Returns one stream per file. Each stream is lazy — the download starts only
-/// when the stream is first polled. The number of concurrent downloads is
-/// bounded by a global semaphore. `stream_buffer_size` controls how many chunks
-/// can be buffered in each stream before backpressure is applied.
-///
-/// **Error handling:** If a download fails after some chunks have already been
-/// yielded, the error surfaces only when the stream is polled past the last
-/// buffered chunk. Callers must consume the stream to completion (or until an
-/// error is returned) to detect download failures.
-#[allow(clippy::too_many_arguments)]
-#[instrument(skip_all, name = "data_client::download_bytes", fields(session_id = tracing::field::Empty, num_files=file_infos.len()))]
-pub async fn download_bytes_async(
-    file_infos: Vec<XetFileInfo>,
-    file_ranges: Option<Vec<Option<FileRange>>>,
-    endpoint: Option<String>,
-    token_info: Option<(String, u64)>,
-    token_refresher: Option<Arc<dyn TokenRefresher>>,
-    progress_updaters: Option<Vec<Arc<dyn TrackingProgressUpdater>>>,
-    user_agent: String,
-    stream_buffer_size: usize,
-) -> errors::Result<Vec<impl Stream<Item = errors::Result<Bytes>>>> {
-    if let Some(updaters) = &progress_updaters
-        && updaters.len() != file_infos.len()
-    {
-        return Err(DataProcessingError::ParameterError("updaters are not same length as file_infos".to_string()));
-    }
-    if let Some(ranges) = &file_ranges
-        && ranges.len() != file_infos.len()
-    {
-        return Err(DataProcessingError::ParameterError("file_ranges are not same length as file_infos".to_string()));
-    }
-
-    let config = default_config(
-        endpoint.unwrap_or_else(|| xet_config().data.default_cas_endpoint.clone()),
-        None,
-        token_info,
-        token_refresher,
-        user_agent,
-    )?;
-    Span::current().record("session_id", &config.session_id);
-
-    let downloader = Arc::new(FileDownloader::new(config.into()).await?);
-    let semaphore = XetRuntime::current().global_semaphore(*CONCURRENT_FILE_DOWNLOAD_LIMITER);
-    let updaters = match progress_updaters {
-        None => vec![None; file_infos.len()],
-        Some(updaters) => updaters.into_iter().map(Some).collect(),
-    };
-    let ranges = match file_ranges {
-        None => vec![None; file_infos.len()],
-        Some(ranges) => ranges,
-    };
-
-    let mut readers = Vec::with_capacity(updaters.len());
-    for ((file_info, file_range), updater) in file_infos.into_iter().zip(ranges).zip(updaters) {
-        readers.push(smudge_bytes(
-            downloader.clone(),
-            semaphore.clone(),
-            file_info,
-            file_range,
-            updater,
-            stream_buffer_size,
-        )?);
-    }
-
-    Ok(readers)
-}
-
 // The sha256, if provided and valid, will be directly used in shard upload to avoid redundant computation.
 #[instrument(skip_all, name = "data_client::upload_files",
     fields(session_id = tracing::field::Empty,
@@ -480,50 +408,9 @@ async fn smudge_file(
     Ok(file_path.to_string())
 }
 
-/// Downloads a file's content and returns it as a stream of byte chunks.
-///
-/// The download is lazy: nothing happens until the consumer starts polling the
-/// stream. On first poll, a `StreamingWriter` is created that sends resolved
-/// `Bytes` chunks directly through an async channel, avoiding the blocking
-/// thread and redundant copies of `SequentialWriter`. Backpressure is provided
-/// by the bounded channel buffer. A semaphore limits the number of concurrent
-/// downloads.
-fn smudge_bytes(
-    downloader: Arc<FileDownloader>,
-    semaphore: Arc<tokio::sync::Semaphore>,
-    file_info: XetFileInfo,
-    file_range: Option<FileRange>,
-    progress_updater: Option<Arc<dyn TrackingProgressUpdater>>,
-    stream_buffer_size: usize,
-) -> errors::Result<impl Stream<Item = errors::Result<Bytes>>> {
-    let merkle_hash = file_info.merkle_hash()?;
-    let file_hash = file_info.hash().into();
-
-    Ok(try_stream! {
-        let progress_updater = progress_updater.map(ItemProgressUpdater::new);
-        let (output, mut rx) = DataOutput::write_byte_stream(stream_buffer_size);
-
-        let handle = tokio::spawn(async move {
-            let _permit = semaphore.acquire().await.map_err(|_| {
-                DataProcessingError::InternalError("download semaphore closed".to_string())
-            })?;
-            downloader
-                .smudge_file_from_hash(&merkle_hash, file_hash, output, file_range, progress_updater)
-                .await
-        });
-
-        while let Some(chunk) = rx.recv().await {
-            yield chunk;
-        }
-
-        handle.await??;
-    })
-}
-
 #[cfg(test)]
 mod tests {
     use dirs::home_dir;
-    use futures::TryStreamExt;
     use serial_test::serial;
     use tempfile::tempdir;
     use utils::EnvVarGuard;
@@ -738,39 +625,16 @@ mod tests {
         assert_eq!(file_info1.file_size(), file_info3.file_size());
     }
 
-    fn test_contents() -> Vec<Vec<u8>> {
-        vec![
-            vec![],
-            b"Hello, World!".to_vec(),
-            (0..1_000_000).map(|i| (i % 256) as u8).collect(),
-        ]
-    }
-
-    #[tokio::test]
-    async fn test_bytes_round_trip() {
-        let temp_dir = tempdir().unwrap();
-        let endpoint = format!("local://{}", temp_dir.path().display());
-        let contents = test_contents();
-
-        let file_infos = upload_bytes_async(contents.clone(), Some(endpoint.clone()), None, None, None, "test".into())
-            .await
-            .unwrap();
-
-        let readers = download_bytes_async(file_infos, None, Some(endpoint), None, None, None, "test".into(), 64)
-            .await
-            .unwrap();
-        for (reader, expected) in readers.into_iter().zip(&contents) {
-            let chunks: Vec<Bytes> = reader.try_collect().await.unwrap();
-            assert_eq!(chunks.concat(), *expected);
-        }
-    }
-
     #[tokio::test]
     async fn test_upload_bytes() {
         let temp_dir = tempdir().unwrap();
         let endpoint = format!("local://{}", temp_dir.path().display());
 
-        let contents = test_contents();
+        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())
@@ -799,36 +663,4 @@ mod tests {
             assert_eq!(std::fs::read(path).unwrap(), *expected);
         }
     }
-
-    #[tokio::test]
-    async fn test_download_bytes() {
-        let temp_dir = tempdir().unwrap();
-        let endpoint = format!("local://{}", temp_dir.path().display());
-
-        let contents = test_contents();
-
-        // Write contents to files and upload via upload_async.
-        let upload_dir = tempdir().unwrap();
-        let file_paths: Vec<String> = contents
-            .iter()
-            .enumerate()
-            .map(|(i, content)| {
-                let path = upload_dir.path().join(format!("{i}"));
-                std::fs::write(&path, content).unwrap();
-                path.to_str().unwrap().to_string()
-            })
-            .collect();
-        let file_infos = upload_async(file_paths, None, Some(endpoint.clone()), None, None, None, "test".into())
-            .await
-            .unwrap();
-
-        // Download all as bytes and verify.
-        let readers = download_bytes_async(file_infos, None, Some(endpoint), None, None, None, "test".into(), 64)
-            .await
-            .unwrap();
-        for (reader, expected) in readers.into_iter().zip(&contents) {
-            let chunks: Vec<Bytes> = reader.try_collect().await.unwrap();
-            assert_eq!(chunks.concat(), *expected);
-        }
-    }
 }

data/src/data_client.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