perf: cache reconstruction plans and pre-warm on open

Cargo.toml

@@ -16,7 +16,7 @@ xorb_object = { git = "https://github.com/huggingface/xet-core.git", rev = "c805
 merklehash = { git = "https://github.com/huggingface/xet-core.git", rev = "c805591" }
 
 # External crates
-fuser = "0.17"
+fuser = { version = "0.17", features = ["macos-no-mount"] }
 tokio = { version = "1", features = ["rt-multi-thread", "macros", "signal"] }
 reqwest = { version = "0.12", features = ["json", "stream"] }
 futures = "0.3"

src/cached_xet_client.rs

@@ -1,61 +1,188 @@
 use std::collections::HashMap;
 use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
 
 use bytes::Bytes;
 use cas_client::adaptive_concurrency::ConnectionPermit;
 use cas_client::{Client, ProgressCallback, URLProvider};
-use xorb_object::SerializedXorbObject;
-use cas_types::{BatchQueryReconstructionResponse, FileRange, QueryReconstructionResponse};
+use cas_types::{
+    BatchQueryReconstructionResponse, FileRange, HexMerkleHash, QueryReconstructionResponse, XorbReconstructionTerm,
+};
 use mdb_shard::file_structs::MDBFileInfo;
 use merklehash::MerkleHash;
-use tracing::debug;
+use xorb_object::SerializedXorbObject;
 
 type Result<T> = std::result::Result<T, cas_client::CasClientError>;
 
 const MAX_CACHE_ENTRIES: usize = 1024;
+/// Presigned URLs returned by the CAS server expire after 1 hour.
+/// Evict cache entries just before expiry to avoid serving stale URLs.
+const CACHE_TTL: Duration = Duration::from_secs(59 * 60);
 
-type CacheKey = (MerkleHash, Option<FileRange>);
+struct CacheEntry {
+    response: QueryReconstructionResponse,
+    inserted_at: Instant,
+}
+
+impl CacheEntry {
+    fn new(response: QueryReconstructionResponse) -> Self {
+        Self {
+            response,
+            inserted_at: Instant::now(),
+        }
+    }
+
+    fn is_valid(&self, ttl: Duration) -> bool {
+        self.inserted_at.elapsed() < ttl
+    }
+}
 
 // TODO: move this into xet-core (cas_client or file_reconstruction) so all consumers benefit.
 pub struct CachedXetClient {
     inner: Arc<dyn Client>,
-    cache: Mutex<HashMap<CacheKey, QueryReconstructionResponse>>,
+    /// Only full-file reconstruction plans are cached (key = file hash, no range).
+    /// Range responses are derived on the fly via `derive_range_response`.
+    cache: Mutex<HashMap<MerkleHash, CacheEntry>>,
+    ttl: Duration,
 }
 
 impl CachedXetClient {
     pub fn new(inner: Arc<dyn Client>) -> Arc<Self> {
+        Self::new_with_ttl(inner, CACHE_TTL)
+    }
+
+    fn new_with_ttl(inner: Arc<dyn Client>, ttl: Duration) -> Arc<Self> {
         Arc::new(Self {
             inner,
             cache: Mutex::new(HashMap::new()),
+            ttl,
         })
     }
 }
 
+/// Derive a range-scoped `QueryReconstructionResponse` from a cached full-file response.
+///
+/// The full-file response lists all terms in file order with their unpacked byte lengths.
+/// We walk the terms, track cumulative byte offsets, and keep only terms that overlap
+/// `[range.start, range.end)`. The `offset_into_first_range` is the byte offset within
+/// the first overlapping term at which the requested range starts.
+///
+/// # Limitation: over-fetch vs CAS server trimming (P2)
+///
+/// When the CAS server handles a range query directly, it trims each term's `ChunkRange`
+/// at chunk granularity: it advances `range.start` past whole chunks whose data falls
+/// entirely before the requested offset, and cuts `range.end` after the last needed chunk.
+/// This produces tight presigned S3 URLs covering only the necessary compressed bytes.
+///
+/// This function cannot do the same trimming because `XorbReconstructionTerm` only carries
+/// `unpacked_length` for the whole term, not per-chunk sizes. As a result, derived responses
+/// keep the full `ChunkRange` of each overlapping term — the caller downloads and decompresses
+/// up to one extra chunk at the start and one at the end per term.
+///
+/// In practice the over-fetch is small: at most ~2 × chunk_size (~128 KB) per term per
+/// range query. For typical safetensors workloads this is negligible (<1% of total data).
+///
+/// TODO: fix P2 — add `chunk_uncompressed_sizes: Vec<u32>` (and `chunk_compressed_sizes`)
+/// to `XorbReconstructionTerm` in xet-core/xetcas so that chunk-level trimming can be
+/// replicated client-side, reducing over-fetch to zero.
+fn derive_range_response(full: &QueryReconstructionResponse, range: FileRange) -> QueryReconstructionResponse {
+    let mut cur_offset: u64 = 0;
+    let mut result_terms: Vec<XorbReconstructionTerm> = Vec::new();
+    let mut offset_into_first: u64 = 0;
+    let mut needed_hashes: std::collections::HashSet<HexMerkleHash> = std::collections::HashSet::new();
+
+    for term in &full.terms {
+        let term_start = cur_offset;
+        let term_end = cur_offset + term.unpacked_length as u64;
+
+        if term_end <= range.start {
+            cur_offset = term_end;
+            continue;
+        }
+        if term_start >= range.end {
+            break;
+        }
+
+        if result_terms.is_empty() {
+            offset_into_first = range.start.saturating_sub(term_start);
+        }
+        needed_hashes.insert(term.hash);
+        result_terms.push(term.clone());
+        cur_offset = term_end;
+    }
+
+    let fetch_info = full
+        .fetch_info
+        .iter()
+        .filter(|(k, _)| needed_hashes.contains(*k))
+        .map(|(k, v)| (*k, v.clone()))
+        .collect();
+
+    QueryReconstructionResponse {
+        offset_into_first_range: offset_into_first,
+        terms: result_terms,
+        fetch_info,
+    }
+}
+
 #[async_trait::async_trait]
 impl Client for CachedXetClient {
     async fn get_reconstruction(
         &self,
         file_id: &MerkleHash,
         bytes_range: Option<FileRange>,
     ) -> Result<Option<QueryReconstructionResponse>> {
-        let key = (*file_id, bytes_range);
-
-        {
-            let cache = self.cache.lock().expect("cache poisoned");
-            if let Some(response) = cache.get(&key) {
-                debug!("reconstruction cache hit for {file_id}");
-                return Ok(Some(response.clone()));
+        // Clone the full-file plan under a brief lock, then derive outside the lock so
+        // concurrent pread() calls from N threads don't serialize on this mutex.
+        let cached_full = {
+            let mut cache = self.cache.lock().expect("cache poisoned");
+            match cache.get(file_id) {
+                Some(entry) if entry.is_valid(self.ttl) => Some(entry.response.clone()),
+                Some(_) => {
+                    // Expired — evict so the next caller fetches fresh URLs from CAS.
+                    cache.remove(file_id);
+                    None
+                }
+                None => None,
+            }
+        }; // ← lock released here
+
+        if let Some(full) = cached_full {
+            if let Some(range) = bytes_range {
+                let resp = derive_range_response(&full, range);
+                // Mirror the CAS server's EOF contract:
+                // - Non-empty file, range past EOF → None
+                // - Empty file, range.start == 0 → Some(empty_terms)  (only valid range on empty file)
+                // - Empty file, range.start > 0  → None               (past EOF on empty file)
+                if resp.terms.is_empty() && (!full.terms.is_empty() || range.start > 0) {
+                    return Ok(None);
+                }
+                tracing::debug!(
+                    "recon: DERV file={:.8} range={:?} terms={}",
+                    file_id,
+                    bytes_range,
+                    resp.terms.len()
+                );
+                return Ok(Some(resp));
+            } else {
+                tracing::debug!("recon: HIT  file={:.8} range=None terms={}", file_id, full.terms.len());
+                return Ok(Some(full));
             }
         }
 
+        tracing::debug!("recon: CAS  file={:.8} range={:?}", file_id, bytes_range);
         let result = self.inner.get_reconstruction(file_id, bytes_range).await?;
 
-        if let Some(ref response) = result {
+        // Only cache full-file plans. Range responses are never stored: they are always
+        // derived from the full-file plan, and the full-file plan is warmed at open() time.
+        if bytes_range.is_none()
+            && let Some(response) = &result
+        {
             let mut cache = self.cache.lock().expect("cache poisoned");
             if cache.len() >= MAX_CACHE_ENTRIES {
                 cache.clear();
             }
-            cache.insert(key, response.clone());
+            cache.insert(*file_id, CacheEntry::new(response.clone()));
         }
 
         Ok(result)
@@ -134,7 +261,7 @@ mod tests {
 
     struct MockClient {
         mode: MockMode,
-        calls: Mutex<HashMap<CacheKey, usize>>,
+        calls: Mutex<HashMap<(MerkleHash, Option<FileRange>), usize>>,
         total_calls: AtomicUsize,
         barrier: Option<Arc<tokio::sync::Barrier>>,
     }
@@ -158,7 +285,7 @@ mod tests {
             }
         }
 
-        fn call_count(&self, key: CacheKey) -> usize {
+        fn call_count(&self, key: (MerkleHash, Option<FileRange>)) -> usize {
             self.calls
                 .lock()
                 .expect("calls lock poisoned")
@@ -297,23 +424,60 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn cache_key_includes_bytes_range() {
+    async fn range_derived_from_full_file_plan() {
         let inner_impl = Arc::new(MockClient::new(MockMode::ReturnSome));
         let inner: Arc<dyn Client> = inner_impl.clone();
         let client = CachedXetClient::new(inner);
         let key = hash_for(4);
         let r = Some(FileRange::new(10, 20));
 
+        // Fetch full-file plan first (caches it).
         client.get_reconstruction(&key, None).await.unwrap();
         client.get_reconstruction(&key, None).await.unwrap();
+
+        // Range queries are derived from the cached full-file plan — never hit backend.
         client.get_reconstruction(&key, r).await.unwrap();
-        client.get_reconstruction(&key, r).await.unwrap();
+        client
+            .get_reconstruction(&key, Some(FileRange::new(20, 30)))
+            .await
+            .unwrap();
 
         assert_eq!(inner_impl.call_count((key, None)), 1);
-        assert_eq!(inner_impl.call_count((key, r)), 1);
+        assert_eq!(inner_impl.call_count((key, r)), 0); // derived, never hit backend
+        assert_eq!(inner_impl.total_calls(), 1);
+    }
+
+    #[tokio::test]
+    async fn range_query_hits_backend_when_no_full_file_cached() {
+        let inner_impl = Arc::new(MockClient::new(MockMode::ReturnSome));
+        let inner: Arc<dyn Client> = inner_impl.clone();
+        let client = CachedXetClient::new(inner);
+        let key = hash_for(40);
+        let r = Some(FileRange::new(10, 20));
+
+        // Range query with no full-file plan cached → hits backend.
+        // The response is NOT cached (range responses are never stored).
+        client.get_reconstruction(&key, r).await.unwrap();
+        client.get_reconstruction(&key, r).await.unwrap();
+
+        assert_eq!(inner_impl.call_count((key, r)), 2); // no cache, each hits backend
         assert_eq!(inner_impl.total_calls(), 2);
     }
 
+    #[tokio::test]
+    async fn expired_entries_are_evicted_and_refetched() {
+        let inner_impl = Arc::new(MockClient::new(MockMode::ReturnSome));
+        let inner: Arc<dyn Client> = inner_impl.clone();
+        let client = CachedXetClient::new_with_ttl(inner, Duration::ZERO);
+        let key = hash_for(50);
+
+        // With TTL=0 every entry is immediately expired — each call hits the backend.
+        client.get_reconstruction(&key, None).await.unwrap();
+        client.get_reconstruction(&key, None).await.unwrap();
+
+        assert_eq!(inner_impl.call_count((key, None)), 2);
+    }
+
     #[tokio::test]
     async fn cache_clears_when_capacity_is_reached() {
         let inner_impl = Arc::new(MockClient::new(MockMode::ReturnSome));

src/fuse.rs

@@ -462,17 +462,8 @@ pub fn mount_fuse(
         config.mount_options.push(fuser::MountOption::RO);
     }
     config.acl = fuser::SessionACL::All;
-    #[cfg(not(target_os = "macos"))]
-    {
-        config.clone_fd = true;
-        config.n_threads = Some(max_threads);
-    }
-    #[cfg(target_os = "macos")]
-    {
-        // macFUSE only supports single-threaded session (fuser limitation).
-        let _ = max_threads;
-        config.n_threads = Some(1);
-    }
+    config.clone_fd = true;
+    config.n_threads = Some(max_threads);
 
     let session = match fuser::Session::new(adapter, mount_point, &config) {
         Ok(s) => s,

src/setup.rs

@@ -219,9 +219,9 @@ pub fn setup(is_nfs: bool) -> MountSetup {
         ))
         .expect("Failed to create CAS client");
     let cached_client = CachedXetClient::new(raw_client);
-    let download_session = FileDownloadSession::from_client(cached_client, None, Some(xorb_cache));
+    let download_session = FileDownloadSession::from_client(cached_client.clone(), None, Some(xorb_cache));
     let upload_config = if read_only { None } else { Some(cas_config) };
-    let xet_sessions = XetSessions::new(download_session, upload_config);
+    let xet_sessions = XetSessions::new(download_session, upload_config, cached_client);
 
     let advanced_writes = args.advanced_writes || (is_nfs && !read_only);
     // Repos need a staging dir for HTTP download cache (open_readonly),

src/test_mocks.rs

@@ -312,6 +312,8 @@ impl XetOps for MockXet {
         Ok(results)
     }
 
+    async fn warm_reconstruction_cache(&self, _xet_hash: &str) {}
+
     fn download_stream_boxed(&self, file_info: &XetFileInfo, offset: u64) -> Result<Box<dyn DownloadStreamOps>> {
         let prev_fail = self.range_fail_count.load(Ordering::SeqCst);
         if prev_fail > 0 {