diff --git a/src/setup.rs b/src/setup.rs index 2ec7446e..ad628864 100644 --- a/src/setup.rs +++ b/src/setup.rs @@ -359,12 +359,16 @@ pub fn init_tracing(daemon: bool) { let upload_cap = XET_UPLOAD_CONCURRENCY_CAP.to_string(); for (k, v) in [ ("HF_XET_CLIENT_AC_INITIAL_DOWNLOAD_CONCURRENCY", "16"), + // Let the adaptive controller scale far enough to saturate fast links + // under many concurrent readers (xet-core's default cap is 64). + ("HF_XET_CLIENT_AC_MAX_DOWNLOAD_CONCURRENCY", "124"), ("HF_XET_CLIENT_AC_MIN_BYTES_REQUIRED_FOR_ADJUSTMENT", "4194304"), ("HF_XET_RECONSTRUCTION_MIN_RECONSTRUCTION_FETCH_SIZE", "8388608"), ("HF_XET_RECONSTRUCTION_MIN_PREFETCH_BUFFER", "8388608"), ("HF_XET_RECONSTRUCTION_TARGET_BLOCK_COMPLETION_TIME", "30"), ("HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_SIZE", "134217728"), - ("HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_LIMIT", "268435456"), + // Also the budget split between per-stream read buffers (see xet.rs). + ("HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_LIMIT", "1073741824"), // Per-read inactivity timeout for CAS/CDN transfers (resets on every byte // received, so slow-but-progressing reads are fine). This governs the // DOWNLOAD/reconstruction path (term fetches and whole-file downloads); diff --git a/src/xet.rs b/src/xet.rs index a3019f8f..1d1e7d79 100644 --- a/src/xet.rs +++ b/src/xet.rs @@ -1,16 +1,18 @@ use std::path::{Path, PathBuf}; -use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; use bytes::Bytes; use xet_client::cas_client::Client; use xet_client::cas_types::FileRange; use xet_client::chunk_cache::ChunkCache; use xet_core_structures::merklehash::MerkleHash; +use xet_core_structures::xorb_object::constants::MAX_XORB_BYTES; use xet_data::file_reconstruction::{DownloadStream, FileReconstructor}; use xet_data::processing::configurations::TranslatorConfig; use xet_data::processing::{FileDownloadSession, FileUploadSession, Sha256Policy, SingleFileCleaner, XetFileInfo}; use xet_runtime::core::XetContext; +use xet_runtime::utils::adjustable_semaphore::AdjustableSemaphore; use crate::error::{Error, Result}; @@ -48,6 +50,110 @@ pub trait DownloadStreamOps: Send { async fn next(&mut self) -> Result>; } +// ── Per-stream download buffers ─────────────────────────────────────── + +/// Ceiling of a single stream's buffer: pipelining depth of a lone reader. +const STREAM_BUFFER_MAX: u64 = 256 * 1_048_576; + +/// Floor of a single stream's buffer: one full xorb, the largest possible +/// term. xet-core clamps a term acquire against the semaphore total only +/// when it is issued (`AdjustableSemaphore::to_physical_acquire`) and does +/// not re-clamp pending acquires on shrink, so shrinking below the largest +/// term would leave an in-flight acquire that can never be satisfied. The +/// floor can become memory-driven once that is fixed upstream. +fn stream_buffer_min() -> u64 { + *MAX_XORB_BYTES as u64 +} + +/// Gives every read stream a private download buffer. +/// +/// Streams are consumed by FUSE `read()` calls that each pin a worker thread. +/// With xet-core's single global buffer (FIFO), a stream that already holds +/// buffer cannot release it until its consumer gets a worker thread, which +/// may itself be blocked waiting for buffer on another stream. Under many +/// concurrent cold readers this starves new streams at their first byte and +/// wedges the mount (#234). Private buffers mean streams never wait on each +/// other. Each is sized `budget / active_streams`, clamped to +/// `[stream_buffer_min(), STREAM_BUFFER_MAX]`, so memory in flight is bounded +/// by `max(budget, streams * stream_buffer_min())`. +struct StreamBufferPool { + budget: u64, + inner: Mutex, +} + +#[derive(Default)] +struct StreamBuffers { + active: Vec>, + /// Share currently applied to every active buffer. + share: u64, +} + +impl StreamBufferPool { + fn new(budget: u64) -> Arc { + Arc::new(Self { + budget, + inner: Mutex::new(StreamBuffers::default()), + }) + } + + /// Create a buffer for a new stream, already sized to the new fair + /// share, and resize the other active buffers to match. The guard + /// unregisters the buffer on drop. + fn register(self: &Arc) -> StreamBufferGuard { + let mut inner = self.inner.lock().expect("stream buffers poisoned"); + let share = self.share_for(inner.active.len() + 1); + let buffer = AdjustableSemaphore::new(share, (stream_buffer_min(), STREAM_BUFFER_MAX)); + inner.active.push(buffer.clone()); + Self::apply_share(&mut inner, share); + StreamBufferGuard { + pool: self.clone(), + buffer, + } + } + + fn unregister(&self, buffer: &Arc) { + let mut inner = self.inner.lock().expect("stream buffers poisoned"); + if let Some(index) = inner.active.iter().position(|other| Arc::ptr_eq(other, buffer)) { + inner.active.swap_remove(index); + } + let share = self.share_for(inner.active.len()); + Self::apply_share(&mut inner, share); + } + + /// Resize every active buffer to `share`; a no-op when it is unchanged. + /// Shrinks apply lazily: permits already held are reclaimed as they return. + fn apply_share(inner: &mut StreamBuffers, share: u64) { + if inner.share == share { + return; + } + inner.share = share; + for buffer in &inner.active { + // Each call is a no-op when the target is on the other side of + // the current total; the increment's virtual permit releases the + // added capacity on drop. + drop(buffer.increment_permits_to_target(share)); + buffer.decrement_permits_to_target(share); + } + } + + fn share_for(&self, active_streams: usize) -> u64 { + (self.budget / active_streams.max(1) as u64).clamp(stream_buffer_min(), STREAM_BUFFER_MAX) + } +} + +/// Keeps a stream's buffer registered in its pool; dropping it (with the +/// stream) hands the freed share back to the remaining streams. +struct StreamBufferGuard { + pool: Arc, + buffer: Arc, +} + +impl Drop for StreamBufferGuard { + fn drop(&mut self) { + self.pool.unregister(&self.buffer); + } +} + // ── XetSessions ─────────────────────────────────────────────────────── /// Core xet-core sessions for CAS downloads and uploads. @@ -61,6 +167,7 @@ pub struct XetSessions { /// Chunk cache attached to unbounded streams; bounded range downloads skip it /// to avoid pulling whole xorbs for small range requests. chunk_cache: Option>, + stream_buffers: Arc, } impl XetSessions { @@ -71,34 +178,43 @@ impl XetSessions { cas_client: Arc, chunk_cache: Option>, ) -> Arc { + // The same knob xet-core uses for its global buffer, so the existing + // HF_XET_RECONSTRUCTION_DOWNLOAD_BUFFER_LIMIT override keeps working. + let stream_buffers = StreamBufferPool::new(ctx.config.reconstruction.download_buffer_limit.as_u64()); Arc::new(Self { ctx, session, upload_config, cas_client, chunk_cache, + stream_buffers, }) } /// Start a streaming download for a byte range. /// When `end` is `Some`, only bytes `[offset, end)` are fetched (bounded range). /// When `end` is `None`, fetches from `offset` to end of file (unbounded stream). - pub fn download_stream(&self, file_info: &XetFileInfo, offset: u64, end: Option) -> Result { + fn download_stream(&self, file_info: &XetFileInfo, offset: u64, end: Option) -> Result { let hash = file_info .merkle_hash() .map_err(|e| Error::Xet(format!("invalid hash: {e}")))?; let is_unbounded = end.is_none(); let file_size = file_info.file_size().unwrap_or(u64::MAX); let end = end.unwrap_or(file_size); - let mut reconstructor = - FileReconstructor::new(&self.ctx, &self.cas_client, hash).with_byte_range(FileRange::new(offset, end)); + let buffer = self.stream_buffers.register(); + let mut reconstructor = FileReconstructor::new(&self.ctx, &self.cas_client, hash) + .with_byte_range(FileRange::new(offset, end)) + .with_buffer_semaphore(buffer.buffer.clone()); // Attach chunk cache only to the unbounded stream path: the xorb disk // cache pulls full xorbs (~64MB) even for small range requests, which // is wasteful for random reads. Sequential reads (unbounded) benefit. if is_unbounded && let Some(cache) = self.chunk_cache.as_ref() { reconstructor = reconstructor.with_chunk_cache(cache.clone()); } - Ok(reconstructor.reconstruct_to_stream()) + Ok(DownloadStreamWrapper { + stream: reconstructor.reconstruct_to_stream(), + _buffer: buffer, + }) } } @@ -146,8 +262,7 @@ impl XetOps for XetSessions { offset: u64, end: Option, ) -> Result> { - let stream = self.download_stream(file_info, offset, end)?; - Ok(Box::new(DownloadStreamWrapper(stream))) + Ok(Box::new(self.download_stream(file_info, offset, end)?)) } async fn warm_reconstruction_cache(&self, xet_hash: &str) { @@ -159,12 +274,17 @@ impl XetOps for XetSessions { // ── DownloadStreamWrapper ───────────────────────────────────────────── -struct DownloadStreamWrapper(DownloadStream); +struct DownloadStreamWrapper { + stream: DownloadStream, + /// Declared after `stream` so the stream (and its reconstruction task) is + /// cancelled before the buffer share is handed back. + _buffer: StreamBufferGuard, +} #[async_trait::async_trait] impl DownloadStreamOps for DownloadStreamWrapper { async fn next(&mut self) -> Result> { - Ok(self.0.next().await?) + Ok(self.stream.next().await?) } } @@ -346,3 +466,58 @@ impl StreamingWriterOps for StreamingWriter { self.bytes_written == 0 } } + +#[cfg(test)] +mod stream_buffer_tests { + use super::*; + + const MIB: u64 = 1_048_576; + const BUDGET: u64 = 1024 * MIB; + + fn total(guard: &StreamBufferGuard) -> u64 { + guard.buffer.total_permits() + } + + #[test] + fn lone_stream_gets_the_ceiling() { + let pool = StreamBufferPool::new(BUDGET); + let stream = pool.register(); + assert_eq!(total(&stream), 256 * MIB); + } + + #[test] + fn shares_shrink_to_the_floor_and_grow_back() { + let pool = StreamBufferPool::new(BUDGET); + let mut streams: Vec<_> = (0..4).map(|_| pool.register()).collect(); + // 1 GiB / 4 = 256 MiB: still at the ceiling. + assert!(streams.iter().all(|stream| total(stream) == 256 * MIB)); + + streams.extend((0..4).map(|_| pool.register())); + // 1 GiB / 8 = 128 MiB for everyone, including the early streams. + assert!(streams.iter().all(|stream| total(stream) == 128 * MIB)); + + streams.extend((0..56).map(|_| pool.register())); + // 1 GiB / 64 = 16 MiB, clamped up to the floor. + assert!(streams.iter().all(|stream| total(stream) == 64 * MIB)); + + streams.truncate(2); + // 1 GiB / 2 = 512 MiB, clamped down to the ceiling. + assert!(streams.iter().all(|stream| total(stream) == 256 * MIB)); + assert_eq!(pool.inner.lock().unwrap().active.len(), 2); + } + + #[tokio::test] + async fn shrink_applies_once_held_permits_return() { + let pool = StreamBufferPool::new(BUDGET); + let first = pool.register(); + let held = first.buffer.acquire_many(200 * MIB).await.unwrap(); + + let _others: Vec<_> = (0..7).map(|_| pool.register()).collect(); + // Target is 128 MiB but 200 MiB are out: the total is updated now, + // the part that cannot be reclaimed yet stays pending. + assert_eq!(total(&first), 128 * MIB); + assert!(first.buffer.available_permits() <= 56 * MIB); + drop(held); + assert_eq!(first.buffer.available_permits(), 128 * MIB); + } +} diff --git a/tests/common/mod.rs b/tests/common/mod.rs index aaeb8e05..37b73c28 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -4,7 +4,7 @@ pub mod bench; pub mod fs_tests; use std::path::Path; -use std::process::{Child, Command}; +use std::process::{Child, Command, Stdio}; use std::sync::Arc; use std::time::{Duration, SystemTime, UNIX_EPOCH}; @@ -281,6 +281,46 @@ pub async fn upload_file(config: Arc, staged_path: &Path) -> X /// Spawn hf-mount-fuse as a child process, wait until the mountpoint is live. /// `extra_args` are appended to the command (e.g. `&["--read-only"]`). pub fn mount_bucket(bucket_id: &str, mount_point: &str, cache_dir: &str, extra_args: &[&str]) -> Child { + spawn_mount_bucket( + bucket_id, + mount_point, + cache_dir, + extra_args, + Stdio::inherit(), + Stdio::inherit(), + ) +} + +/// Like `mount_bucket`, but with the daemon's output written to `log_path` +/// so a test can assert on what it logged. Tracing goes to stdout, so both +/// stdout and stderr are captured. +pub fn mount_bucket_logged( + bucket_id: &str, + mount_point: &str, + cache_dir: &str, + extra_args: &[&str], + log_path: &str, +) -> Child { + let log = std::fs::File::create(log_path).expect("create daemon log file"); + let log_err = log.try_clone().expect("clone daemon log file"); + spawn_mount_bucket( + bucket_id, + mount_point, + cache_dir, + extra_args, + log.into(), + log_err.into(), + ) +} + +fn spawn_mount_bucket( + bucket_id: &str, + mount_point: &str, + cache_dir: &str, + extra_args: &[&str], + stdout: Stdio, + stderr: Stdio, +) -> Child { let token = std::env::var("HF_TOKEN").unwrap(); let binary = std::env::current_exe() @@ -314,6 +354,8 @@ pub fn mount_bucket(bucket_id: &str, mount_point: &str, cache_dir: &str, extra_a ]) .args(extra_args) .args(["bucket", bucket_id, mount_point]) + .stdout(stdout) + .stderr(stderr) .spawn() .expect("Failed to spawn hf-mount-fuse"); @@ -595,3 +637,75 @@ pub fn generate_pattern(size: usize) -> Vec { pub fn verify_pattern(data: &[u8], offset: usize) -> bool { data.iter().enumerate().all(|(i, &b)| b == ((offset + i) % 251) as u8) } + +/// Deterministic pseudo-random byte generator (xorshift64). Each seeded file +/// has distinct content so nothing dedups across files, and readers can +/// regenerate the expected bytes chunk by chunk instead of holding them. +pub struct PseudoRandomBytes(u64); + +impl PseudoRandomBytes { + pub fn new(seed: u64) -> Self { + Self(seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1) + } + + pub fn fill(&mut self, buf: &mut [u8]) { + for chunk in buf.chunks_mut(8) { + let mut x = self.0; + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + self.0 = x; + let bytes = x.to_le_bytes(); + chunk.copy_from_slice(&bytes[..chunk.len()]); + } + } +} + +/// Seed a bucket with `count` files of `size` bytes each under `par/`, all +/// uploaded in one session. Returns the relative paths; file `i` regenerates +/// from `PseudoRandomBytes::new(i)`. +pub async fn seed_parallel_read_files( + hub: &Arc, + tmp_dir_tag: &str, + count: usize, + size: usize, +) -> Vec { + let write_config = build_write_config(hub).await; + let tmp_dir = std::env::temp_dir().join(format!("hf-mount-{}-{}", tmp_dir_tag, std::process::id())); + std::fs::create_dir_all(&tmp_dir).ok(); + + let mut buf = vec![0u8; size]; + let staged: Vec<_> = (0..count) + .map(|i| { + let path = tmp_dir.join(format!("f_{i:02}.bin")); + PseudoRandomBytes::new(i as u64).fill(&mut buf); + std::fs::write(&path, &buf).unwrap(); + path + }) + .collect(); + + let upload_session = FileUploadSession::new(write_config) + .await + .expect("FileUploadSession::new failed"); + let infos = upload_session + .upload_files(staged.iter().map(|p| (p.clone(), Sha256Policy::Skip))) + .await + .expect("upload_files failed"); + upload_session.finalize().await.expect("finalize failed"); + assert_eq!(infos.len(), count); + + let rel_paths: Vec = (0..count).map(|i| format!("par/f_{i:02}.bin")).collect(); + let ops: Vec<_> = rel_paths + .iter() + .zip(&infos) + .map(|(rel, info)| hf_mount::hub_api::BatchOp::AddFile { + path: rel.clone(), + xet_hash: info.hash().to_string(), + mtime: 0, + content_type: None, + }) + .collect(); + hub.batch_operations(&ops).await.expect("batch add failed"); + std::fs::remove_dir_all(&tmp_dir).ok(); + rel_paths +} diff --git a/tests/fuse_ops.rs b/tests/fuse_ops.rs index 44e95b84..6a0cfebf 100644 --- a/tests/fuse_ops.rs +++ b/tests/fuse_ops.rs @@ -249,3 +249,106 @@ async fn test_fuse_revalidation() { panic!("FUSE revalidation test failed: {}", e); } } + +/// Many cold sequential readers at once, more readers than FUSE worker +/// threads, files larger than the initial prefetch window so every stream +/// outlives its first read. Regression test for #234: streams must not +/// starve each other for download buffer, and no read may time out or fail. +#[tokio::test] +async fn test_fuse_parallel_cold_reads() { + const FILE_COUNT: usize = 24; + const CHUNK: usize = 1 << 20; + const FILE_SIZE: usize = 64 * CHUNK; + + let guard = match common::setup_bucket("fuse-par-read").await { + Some(g) => g, + None => return, + }; + let rel_paths = common::seed_parallel_read_files(&guard.hub, "fuse-pr", FILE_COUNT, FILE_SIZE).await; + + let mount_point = format!("/tmp/hf-mount-pr-mnt-{}", std::process::id()); + let cache_dir = format!("/tmp/hf-mount-pr-cache-{}", std::process::id()); + let log_path = format!("/tmp/hf-mount-pr-log-{}.txt", std::process::id()); + + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let child = common::mount_bucket_logged( + &guard.bucket_id, + &mount_point, + &cache_dir, + // Few workers and a short fetch timeout so a wedge shows up as + // EIO within seconds instead of minutes. + &["--read-only", "--max-threads", "4", "--read-fetch-timeout-ms", "5000"], + &log_path, + ); + + let started = std::time::Instant::now(); + let readers: Vec<_> = rel_paths + .iter() + .enumerate() + .map(|(i, rel)| { + let path = format!("{}/{}", mount_point, rel); + std::thread::spawn(move || -> Result<(), String> { + use std::io::Read; + let mut file = std::fs::File::open(&path).map_err(|e| format!("{path}: open: {e}"))?; + let mut expected_gen = common::PseudoRandomBytes::new(i as u64); + let mut got = vec![0u8; CHUNK]; + let mut expected = vec![0u8; CHUNK]; + for chunk in 0..FILE_SIZE / CHUNK { + file.read_exact(&mut got) + .map_err(|e| format!("{path}: read at {}: {e}", chunk * CHUNK))?; + expected_gen.fill(&mut expected); + if got != expected { + return Err(format!("{path}: content mismatch at {}", chunk * CHUNK)); + } + } + if file.read(&mut got).map_err(|e| format!("{path}: read at eof: {e}"))? != 0 { + return Err(format!("{path}: longer than {FILE_SIZE}")); + } + Ok(()) + }) + }) + .collect(); + + let mut errors = Vec::new(); + for reader in readers { + match reader.join() { + Ok(Ok(())) => {} + Ok(Err(e)) => errors.push(e), + Err(_) => errors.push("reader thread panicked".to_string()), + } + } + eprintln!( + "{} readers x {} MiB done in {:?}", + FILE_COUNT, + FILE_SIZE / CHUNK, + started.elapsed() + ); + + common::unmount(&mount_point, child, 30); + + // A single timeout is retried and would otherwise be invisible. + let log_text = std::fs::read_to_string(&log_path).unwrap_or_default(); + let timeouts = log_text.matches("stream read timed out").count(); + if !errors.is_empty() || timeouts > 0 { + let log_tail: Vec<&str> = log_text.lines().rev().take(20).collect(); + return Err(format!( + "{} read errors, {} stream timeouts\nerrors: {:#?}\ndaemon log tail:\n{}", + errors.len(), + timeouts, + errors, + log_tail.into_iter().rev().collect::>().join("\n"), + )); + } + Ok(()) + })); + + std::fs::remove_dir_all(&mount_point).ok(); + std::fs::remove_dir_all(&cache_dir).ok(); + std::fs::remove_file(&log_path).ok(); + + match result { + Ok(Ok(())) => {} + Ok(Err(e)) => panic!("parallel cold read test failed: {}", e), + Err(e) => std::panic::resume_unwind(e), + } +}