[UIO] CachedFs waits for scheduled files to resolve + misc (#10353)

* [CachedFs] new `schedule` and `wait_all` primitives

* [AppendableIdTracker] don't reopen if just opened

* eager NotFound in `schedule_open`

* add traces for async reads

* finish `preopen`/`preload` with `wait_all`

* lock all segments in parallel for `live_reload`

* LIST before everything

to do: we don't have whole-fetch in async mode. to prevent sequential
`len`, we won't overlap static files with LIST.

* `wait_all` returns nothing

lib/blobstore/src/blobstore/tests.rs

@@ -1932,6 +1932,7 @@ fn test_preopen_schedules_files_for_open(#[values(Mode::Mutable, Mode::AppendOnl
         Populate::No,
     )
     .unwrap();
+    cached_fs.wait_all();
 
     // Delete everything `open` reads through the backend. The append-only tracker is read
     // directly from disk, not through the backend, so it must stay.

lib/common/common/src/universal_io/cached_fs/mod.rs

@@ -1,10 +1,12 @@
 use std::collections::HashMap;
 use std::fmt::{Debug, Formatter};
 use std::path::{Path, PathBuf};
-use std::pin::Pin;
 use std::sync::Arc;
 use std::task::Poll;
 
+use futures::StreamExt;
+use futures::future::BoxFuture;
+use futures::stream::FuturesUnordered;
 use parking_lot::Mutex;
 
 use crate::mmap::AdviceSetting;
@@ -89,7 +91,7 @@ where
 }
 
 enum ScheduledFile<S: 'static> {
-    Future(Pin<Box<dyn Future<Output = UioResult<S>> + Send + 'static>>),
+    Future(BoxFuture<'static, UioResult<S>>),
     Ready(UioResult<S>),
     Unchanged,
 }
@@ -252,7 +254,18 @@ impl<Fs: UniversalReadFs> CachedReadFs for CachedFs<Fs> {
         });
 
         let mut open_extra = open_extra.unwrap_or_default();
-        if let Some(info) = self.file_info(path) {
+        if let Some(info) = self.files_info.as_ref() {
+            let Some(info) = info.get(path) else {
+                // The file was not listed, set NotFound eagerly.
+                files_prefetched.insert(
+                    path.to_path_buf(),
+                    ScheduledFile::Ready(Err(UniversalIoError::NotFound {
+                        path: path.to_path_buf(),
+                    })),
+                );
+                return;
+            };
+
             open_extra = open_extra.with_known_len(info.size);
         }
 
@@ -279,28 +292,45 @@ impl<Fs: UniversalReadFs> CachedReadFs for CachedFs<Fs> {
         open_arguments: Option<OpenOptions>,
         open_extra: Option<Fs::OpenExtra>,
     ) {
+        // Check if their file info is complete and didn't change.
+        if self
+            .previous_file_info(path)
+            .zip(self.file_info(path))
+            .is_some_and(|(previous, current)| previous.full_eq(current))
         {
-            let mut files_prefetched = self.files_prefetched.lock();
-
-            if files_prefetched.contains_key(path) {
-                return;
-            }
-
-            // Check if their file info is complete and didn't change.
-            if self
-                .previous_file_info(path)
-                .zip(self.file_info(path))
-                .is_some_and(|(previous, current)| previous.full_eq(current))
-            {
-                files_prefetched.insert(path.to_path_buf(), ScheduledFile::Unchanged);
-                return;
-            }
+            self.files_prefetched
+                .lock()
+                .entry(path.to_path_buf())
+                .or_insert(ScheduledFile::Unchanged);
+            return;
         }
 
         // Otherwise schedule normally
         self.schedule_open(path, open_arguments, open_extra)
     }
 
+    fn schedule(&self, path: PathBuf, fut: BoxFuture<'static, UioResult<Fs::File>>) {
+        self.files_prefetched
+            .lock()
+            .insert(path, ScheduledFile::Future(fut));
+    }
+
+    fn wait_all(&self) {
+        let mut lock = self.files_prefetched.lock();
+        let futs = lock
+            .extract_if(|_path, scheduled| matches!(scheduled, ScheduledFile::Future(_)))
+            .filter_map(|(path, scheduled)| match scheduled {
+                ScheduledFile::Future(fut) => Some(async move { (path, fut.await) }),
+                ScheduledFile::Ready(_) | ScheduledFile::Unchanged => None,
+            })
+            .collect::<FuturesUnordered<_>>();
+
+        let results = futures::executor::block_on(futs.collect::<Vec<_>>());
+        for (path, result) in results {
+            lock.insert(path, ScheduledFile::Ready(result));
+        }
+    }
+
     fn cached_file_info(&self, path: &Path) -> Option<FileInfo> {
         self.files_info.as_ref()?.get(path).cloned()
     }

lib/common/common/src/universal_io/simple_disk_cache/fs.rs

@@ -290,10 +290,14 @@ where
         let remote_extra = extra.remote_extra.clone().with_prevent_caching(true);
 
         let state = match populate {
-            Populate::Auto | Populate::No | Populate::Blocking => {
-                return self.open(path, options, extra);
-            }
-            Populate::PreferBackground => {
+            Populate::Auto | Populate::No => match extra.known_len {
+                Some(len) => State::ready(
+                    self.open_remote(&path, remote_extra.clone())?,
+                    LocalState::new(&local_path, len, options)?,
+                ),
+                None => State::Uninit,
+            },
+            Populate::PreferBackground | Populate::Blocking => {
                 let remote = self.open_remote(&path, remote_extra.clone())?;
                 let len = match extra.known_len {
                     Some(len) => len,

lib/common/common/src/universal_io/traits/file_ops.rs

@@ -1,6 +1,8 @@
 use std::fmt::Debug;
 use std::path::{Path, PathBuf};
 
+use futures::future::BoxFuture;
+
 use crate::universal_io::cached_fs::FileInfo;
 use crate::universal_io::traits::append::UniversalAppend;
 use crate::universal_io::traits::open_extra::OpenExtra;
@@ -195,6 +197,12 @@ pub trait CachedReadFs: UniversalReadFs {
         open_extra: Option<Self::OpenExtra>,
     );
 
+    /// Granular version of `schedule_open` for custom `populate` cases
+    fn schedule(&self, path: PathBuf, fut: BoxFuture<'static, UioResult<Self::File>>);
+
+    /// Wait for all scheduled files to resolve.
+    fn wait_all(&self);
+
     /// Return the file info from the current snapshot.
     fn cached_file_info(&self, path: &Path) -> Option<FileInfo>;
 }

lib/common/io_bridge/src/file.rs

@@ -133,11 +133,27 @@ impl<A: AsyncRead + Clone> UniversalRead for BlobFile<A> {
         _access_pattern: P,
         align: usize,
     ) -> UioResult<ACow<'_>> {
+        let started = std::time::Instant::now();
+        log::trace!(
+            target: crate::LATENCY_LOG_TARGET,
+            "scheduled async read of {}, {:?}",
+            self.path.display(),
+            range
+        );
+
         let buf = self
             .runtime
             .handle()
             .spawn(read_into_byte_buffer::<A>(self, range, align))
             .await??;
+
+        log::trace!(
+            target: crate::LATENCY_LOG_TARGET,
+            "awaited async read of {}, {:?} bytes took {}ms",
+            self.path.display(),
+            buf.len(),
+            started.elapsed().as_millis()
+        );
         Ok(ACow::Owned(buf))
     }
 

lib/edge/src/read_only/refresh.rs

@@ -169,15 +169,23 @@ impl<S: UniversalReadExt + 'static> ReadOnlyEdgeShard<S> {
             });
         });
 
+        let reloads: Vec<_> = survivors
+            .into_iter()
+            // The counter cell is not `Sync`, so fork one per reload outside the
+            // pool; forks drain into the shared accumulator on drop.
+            .map(|(uuid, segment)| (uuid, segment, hw_counter.fork()))
+            .collect();
+        let results: Vec<_> = self.search_pool.install(|| {
+            reloads
+                .into_par_iter()
+                .map(|(uuid, segment, hw)| (uuid, segment.write().live_reload(&hw)))
+                .collect()
+        });
+
         let mut not_found: Vec<(Uuid, OperationError)> = Vec::new();
         let mut first_hard_error: Option<OperationError> = None;
-
-        // TODO(uio): currently this locks each segment one at a time. Once we
-        // do `live_preload` async we'll have a clear signal of finishing
-        // prefetches to start locking. By then, we can also make this section a
-        // rayon par_iter, and take care of hw_counter not being thread-safe.
-        for (uuid, segment) in survivors {
-            match segment.write().live_reload(hw_counter) {
+        for (uuid, result) in results {
+            match result {
                 Ok(()) => {}
                 // An essential file is gone; whether that is benign (the leader removed the
                 // segment while we reloaded it) is decided against a re-read manifest below.

lib/segment/src/id_tracker/mutable_id_tracker/read_only/live_reload.rs

@@ -163,17 +163,21 @@ impl<S: UniversalRead> ReadOnlyAppendableIdTracker<S> {
     ) -> OperationResult<Vec<MappingChange>> {
         // The mappings file is absent until the writer flushes the first point; open it lazily once
         // it appears. Until then there is nothing to read.
-        if self.mappings_file.is_none() {
-            self.mappings_file = Self::try_open(fs, &mappings_path(&self.segment_path))?;
+        match self.mappings_file.as_mut() {
+            Some(file) => {
+                // Refresh the handle to observe data appended by the writer. A lazily-opened handle whose
+                // object does not exist yet (e.g. S3) reports `NotFound` here or from `len`; treat that as
+                // an empty file.
+                file.reopen().ok_not_found()?;
+            }
+            None => {
+                self.mappings_file = Self::try_open(fs, &mappings_path(&self.segment_path))?;
+            }
         }
         let Some(file) = self.mappings_file.as_mut() else {
             return Ok(Vec::new());
         };
 
-        // Refresh the handle to observe data appended by the writer. A lazily-opened handle whose
-        // object does not exist yet (e.g. S3) reports `NotFound` here or from `len`; treat that as
-        // an empty file.
-        file.reopen().ok_not_found()?;
         let Some(file_len) = file.len::<u8>().ok_not_found()? else {
             return Ok(Vec::new());
         };
@@ -217,18 +221,21 @@ impl<S: UniversalRead> ReadOnlyAppendableIdTracker<S> {
     fn reload_versions(&mut self, fs: &impl UniversalReadFs<File = S>) -> OperationResult<usize> {
         // The versions file is absent until the writer flushes the first point; open it lazily once
         // it appears. Until then no version is committed.
-        if self.versions_file.is_none() {
-            self.versions_file = Self::try_open(fs, &versions_path(&self.segment_path))?;
+        match self.versions_file.as_mut() {
+            Some(versions_file) => {
+                // Refresh the handle to observe data appended by the writer. A lazily-opened handle whose
+                // object does not exist yet (e.g. S3) reports `NotFound` here or from `len`; treat that as
+                // an empty file (no committed versions).
+                versions_file.reopen().ok_not_found()?;
+            }
+            None => {
+                self.versions_file = Self::try_open(fs, &versions_path(&self.segment_path))?;
+            }
         }
         let Some(versions_file) = self.versions_file.as_mut() else {
             return Ok(self.internal_to_version.len());
         };
 
-        // Refresh the handle to observe data appended by the writer. A lazily-opened handle whose
-        // object does not exist yet (e.g. S3) reports `NotFound` here or from `len`; treat that as
-        // an empty file (no committed versions).
-        versions_file.reopen().ok_not_found()?;
-
         // Disjoint field borrow so the read (from `versions_file`) can extend `internal_to_version`.
         let internal_to_version = &mut self.internal_to_version;
 

lib/segment/src/index/field_index/bool_index/read_only_bool_index/mod.rs

@@ -485,6 +485,7 @@ mod tests {
             )
             .unwrap()
         );
+        cached_fs.wait_all();
 
         // Everything `open` reads must now come from the prefetch pool.
         fs_err::remove_dir_all(dir.path().join(TRUES_DIRNAME)).unwrap();

lib/segment/src/index/field_index/geo_index/tests.rs

@@ -1345,6 +1345,7 @@ fn test_block_index_preopen() {
         OnDiskGeoIndex::<Storage>::preopen(&cached_fs, temp_dir.path(), Populate::PreferBackground)
             .unwrap()
     );
+    cached_fs.wait_all();
 
     // The sidecar reads of `open` must now come from the prefetch pool.
     fs_err::remove_file(&counts_sidecar).unwrap();

lib/segment/src/index/field_index/numeric_index/tests.rs

@@ -1178,6 +1178,7 @@ fn test_block_index_preopen() {
         )
         .unwrap()
     );
+    cached_fs.wait_all();
 
     // The sidecar read of `open` must now come from the prefetch pool.
     fs_err::remove_file(&block_index_path).unwrap();

lib/segment/src/index/hnsw_index/graph_layers.rs

@@ -800,6 +800,7 @@ mod tests {
         let mut cached_fs = CachedFs::new(MmapFs, dir.path()).unwrap();
         cached_fs.cache_file_info().unwrap();
         HnswGraph::<MmapFile>::preopen_universal(&cached_fs, dir.path(), residency).unwrap();
+        cached_fs.wait_all();
 
         // Everything `load_universal` reads must now come from the prefetch pool.
         for entry in fs_err::read_dir(dir.path()).unwrap() {

lib/segment/src/segment/read_only/lifecycle.rs

@@ -30,31 +30,26 @@ use crate::vector_storage::quantized::quantized_vectors::ReadOnlyQuantizedVector
 use crate::vector_storage::read_only::VectorStorageReadEnum;
 use crate::vector_storage::sparse::read_only::ReadOnlySparseVectorStorage;
 
-/// Build a per-segment [`CachedReadFs`] over `segment_path`.
-///
-/// The files whose names are known in advance (version file, segment state)
-/// are scheduled *before* the listing snapshot is taken, so on backends with
-/// background population their fetch overlaps the listing round-trip.
+/// Build a per-segment [`CachedReadFs`] over `segment_path`. Schedules statically known files.
 fn build_cached_fs<Fs: UniversalReadFs>(
     fs: &Fs,
     segment_path: &Path,
 ) -> OperationResult<CachedFs<Fs>> {
     let mut cached_fs = CachedFs::new(fs.clone(), segment_path)?;
 
-    // Absence is tolerated here: the subsequent read reports it gracefully.
-    for file_name in [VERSION_FILE, SEGMENT_STATE_FILE] {
-        cached_fs.schedule_open(&segment_path.join(file_name), None, None);
-    }
-
-    // Payload index config
-    cached_fs.schedule_open(
-        &PayloadConfig::get_config_path(&get_payload_index_path(segment_path)),
-        None,
-        None,
-    );
-
     cached_fs.cache_file_info()?;
 
+    // TODO(uio): Schedule static files in advance, after implementing
+    // `read_whole_bytes_async`, so that their fetch can overlap with the listing
+    // round-trip.
+    for path in [
+        segment_path.join(VERSION_FILE),
+        segment_path.join(SEGMENT_STATE_FILE),
+        PayloadConfig::get_config_path(&get_payload_index_path(segment_path)),
+    ] {
+        cached_fs.schedule_open(&path, None, None);
+    }
+
     Ok(cached_fs)
 }
 
@@ -211,6 +206,8 @@ impl<S: UniversalReadExt + 'static> ReadOnlySegment<S> {
         deferred_internal_id: Option<PointOffsetType>,
         load_profile: Option<&LoadProfile>,
     ) -> OperationResult<Self> {
+        fs.wait_all();
+
         if SegmentVersion::load_universal(&fs, segment_path)?.is_none() {
             // `FileNotFound`, not a service error: the version file is written last, so
             // its absence means the segment vanished mid-open (or was never completed) —

lib/segment/src/segment/read_only/live_reload.rs

@@ -39,6 +39,10 @@ impl<S: UniversalReadExt + 'static> ReadOnlySegment<S> {
         for vector_data in vector_data.values() {
             vector_data.live_preload(fs)?;
         }
+
+        // Pin the staged files before the writer can churn them: reload then
+        // consumes ready handles and never races the filesystem.
+        fs.wait_all();
         Ok(())
     }
 

lib/segment/src/segment/update_only/lookup/lifecycle.rs

@@ -32,24 +32,24 @@ use crate::vector_storage::sparse::read_only::ReadOnlySparseVectorStorage;
 /// [`preopen`]: LookupSegment::preopen
 const WRITER_POPULATE: Populate = Populate::No;
 
-/// Build the per-segment [`CachedFs`] an open runs over: the version and
-/// state files are prefetched, and the directory listing snapshot is taken so
-/// probes for optional files resolve without inner-filesystem round-trips.
-/// Mirror of the read-only segment's `build_cached_fs`, minus the payload
-/// index config the writer never opens.
+/// Build the per-segment [`CachedFs`] an open runs over. Preloads statically
+/// known files.
+///
+/// Mirror of the read-only segment's `build_cached_fs`, minus the payload index
+/// config the writer never opens.
 fn build_cached_fs<Fs: UniversalReadFs>(
     fs: &Fs,
     segment_path: &Path,
 ) -> OperationResult<CachedFs<Fs>> {
     let mut cached_fs = CachedFs::new(fs.clone(), segment_path)?;
 
+    cached_fs.cache_file_info()?;
+
     // Absence is tolerated here: the subsequent read reports it gracefully.
     for file_name in [VERSION_FILE, SEGMENT_STATE_FILE] {
         cached_fs.schedule_open(&segment_path.join(file_name), None, None);
     }
 
-    cached_fs.cache_file_info()?;
-
     Ok(cached_fs)
 }
 

lib/segment/src/vector_storage/chunked_vectors/read_only/mod.rs

@@ -122,6 +122,7 @@ mod tests {
         let mut cached_fs = CachedFs::new(MmapFs, dir.path()).unwrap();
         cached_fs.cache_file_info().unwrap();
         LiveReload::live_preload(&reader, &cached_fs).unwrap();
+        cached_fs.wait_all();
 
         for file in fs_err::read_dir(dir.path()).unwrap() {
             fs_err::remove_file(file.unwrap().path()).unwrap();

lib/segment/src/vector_storage/dense/read_only/mod.rs

@@ -189,6 +189,7 @@ mod tests {
             let mut cached_fs = CachedFs::new(MmapFs, dir.path()).unwrap();
             cached_fs.cache_file_info().unwrap();
             reader.live_preload(&cached_fs).unwrap();
+            cached_fs.wait_all();
 
             fs_err::remove_dir_all(dir.path()).unwrap();
 

lib/segment/src/vector_storage/quantized/quantized_vectors/read_only/tests.rs

@@ -357,6 +357,7 @@ fn preopen_and_unlink(
     let mut cached_fs = CachedFs::new(MmapFs, dir).unwrap();
     cached_fs.cache_file_info().unwrap();
     ReadOnlyQuantizedVectors::<MmapFile>::preopen(&cached_fs, dir, &vector_config, None).unwrap();
+    cached_fs.wait_all();
 
     for entry in fs_err::read_dir(dir).unwrap() {
         let path = entry.unwrap().path();
@@ -602,6 +603,7 @@ fn reload_chunked_preserves_scores(preload: bool) {
         let mut cached_fs = CachedFs::new(MmapFs, quant_dir.path()).unwrap();
         cached_fs.cache_file_info().unwrap();
         ro.live_preload(&cached_fs).unwrap();
+        cached_fs.wait_all();
 
         fs_err::remove_dir_all(quant_dir.path()).unwrap();
 

lib/segment/src/vector_storage/read_only/mod.rs

@@ -321,8 +321,9 @@ mod tests {
     /// Merely opening after a `preopen` proves nothing: `CachedFs` falls back
     /// to a plain inner open for any path that was never scheduled, so a
     /// scheduled-vs-opened path mismatch would still yield a correct storage.
-    /// To make the prefetch pool the *only* possible source, the storage files
-    /// are unlinked between `preopen` and `open`: the already-open handles
+    /// To make the prefetch pool the *only* possible source, `wait_all`
+    /// materializes the scheduled opens (the sequence the segment open runs)
+    /// and the storage files are then unlinked before `open`: the handles
     /// parked in the pool stay readable, while any fallback open hits
     /// `NotFound`.
     #[test]
@@ -352,6 +353,7 @@ mod tests {
         let config = dense_config(VectorStorageType::ChunkedMmap, None);
         let cached_fs = snapshot_cached_fs(dir.path());
         VectorStorageReadEnum::<MmapFile>::preopen(&cached_fs, &config, dir.path(), None).unwrap();
+        cached_fs.wait_all();
 
         // Everything `open` reads must now come from the prefetch pool.
         for dir_name in [
@@ -400,6 +402,7 @@ mod tests {
         let config = dense_config(VectorStorageType::Mmap, None);
         let cached_fs = snapshot_cached_fs(dir.path());
         VectorStorageReadEnum::<MmapFile>::preopen(&cached_fs, &config, dir.path(), None).unwrap();
+        cached_fs.wait_all();
 
         // Everything `open` reads must now come from the prefetch pool.
         for file_name in [
@@ -463,6 +466,7 @@ mod tests {
         );
         let cached_fs = snapshot_cached_fs(dir.path());
         VectorStorageReadEnum::<MmapFile>::preopen(&cached_fs, &config, dir.path(), None).unwrap();
+        cached_fs.wait_all();
 
         // Everything `open` reads must now come from the prefetch pool.
         for dir_name in [

lib/segment/src/vector_storage/sparse/read_only/mod.rs

@@ -145,6 +145,7 @@ mod tests {
         cached_fs.cache_file_info().unwrap();
         ReadOnlySparseVectorStorage::<MmapFile>::preopen(&cached_fs, dir.path(), Populate::No)
             .unwrap();
+        cached_fs.wait_all();
 
         // Everything `open` reads must now come from the prefetch pool.
         for dir_name in [STORAGE_DIRNAME, DELETED_DIRNAME] {
@@ -234,6 +235,7 @@ mod tests {
             let mut cached_fs = CachedFs::new(MmapFs, dir.path()).unwrap();
             cached_fs.cache_file_info().unwrap();
             reader.live_preload(&cached_fs).unwrap();
+            cached_fs.wait_all();
 
             fs_err::remove_dir_all(dir.path()).unwrap();
 

lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs

@@ -724,6 +724,7 @@ mod tests {
         let mut cached_fs = CachedFs::new(MmapFs, dir.path()).unwrap();
         cached_fs.cache_file_info().unwrap();
         preopen(&cached_fs, dir.path(), Populate::No).unwrap();
+        cached_fs.wait_all();
 
         // Everything `open_ro` reads must now come from the prefetch pool.
         for entry in fs_err::read_dir(dir.path()).unwrap() {