diff --git a/lib/blobstore/src/blobstore/tests.rs b/lib/blobstore/src/blobstore/tests.rs index db5e0cede4..ad89d17e06 100644 --- a/lib/blobstore/src/blobstore/tests.rs +++ b/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. diff --git a/lib/common/common/src/universal_io/cached_fs/mod.rs b/lib/common/common/src/universal_io/cached_fs/mod.rs index 4a2b262f4d..3460ecfecb 100644 --- a/lib/common/common/src/universal_io/cached_fs/mod.rs +++ b/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 { - Future(Pin> + Send + 'static>>), + Future(BoxFuture<'static, UioResult>), Ready(UioResult), Unchanged, } @@ -252,7 +254,18 @@ impl CachedReadFs for CachedFs { }); 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 CachedReadFs for CachedFs { open_arguments: Option, open_extra: Option, ) { + // 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>) { + 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::>(); + + let results = futures::executor::block_on(futs.collect::>()); + for (path, result) in results { + lock.insert(path, ScheduledFile::Ready(result)); + } + } + fn cached_file_info(&self, path: &Path) -> Option { self.files_info.as_ref()?.get(path).cloned() } diff --git a/lib/common/common/src/universal_io/simple_disk_cache/fs.rs b/lib/common/common/src/universal_io/simple_disk_cache/fs.rs index 734d00c7a4..63007bea92 100644 --- a/lib/common/common/src/universal_io/simple_disk_cache/fs.rs +++ b/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, diff --git a/lib/common/common/src/universal_io/traits/file_ops.rs b/lib/common/common/src/universal_io/traits/file_ops.rs index 337db3de02..5cd47e6df7 100644 --- a/lib/common/common/src/universal_io/traits/file_ops.rs +++ b/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, ); + /// Granular version of `schedule_open` for custom `populate` cases + fn schedule(&self, path: PathBuf, fut: BoxFuture<'static, UioResult>); + + /// 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; } diff --git a/lib/common/io_bridge/src/file.rs b/lib/common/io_bridge/src/file.rs index 6bc0703306..ab2ae4c006 100644 --- a/lib/common/io_bridge/src/file.rs +++ b/lib/common/io_bridge/src/file.rs @@ -133,11 +133,27 @@ impl UniversalRead for BlobFile { _access_pattern: P, align: usize, ) -> UioResult> { + 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::(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)) } diff --git a/lib/edge/src/read_only/refresh.rs b/lib/edge/src/read_only/refresh.rs index 4fb20eea7c..69d2aaa75a 100644 --- a/lib/edge/src/read_only/refresh.rs +++ b/lib/edge/src/read_only/refresh.rs @@ -169,15 +169,23 @@ impl ReadOnlyEdgeShard { }); }); + 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 = 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. diff --git a/lib/segment/src/id_tracker/mutable_id_tracker/read_only/live_reload.rs b/lib/segment/src/id_tracker/mutable_id_tracker/read_only/live_reload.rs index 2e6788c7a7..1570f24d23 100644 --- a/lib/segment/src/id_tracker/mutable_id_tracker/read_only/live_reload.rs +++ b/lib/segment/src/id_tracker/mutable_id_tracker/read_only/live_reload.rs @@ -163,17 +163,21 @@ impl ReadOnlyAppendableIdTracker { ) -> OperationResult> { // 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::().ok_not_found()? else { return Ok(Vec::new()); }; @@ -217,18 +221,21 @@ impl ReadOnlyAppendableIdTracker { fn reload_versions(&mut self, fs: &impl UniversalReadFs) -> OperationResult { // 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; diff --git a/lib/segment/src/index/field_index/bool_index/read_only_bool_index/mod.rs b/lib/segment/src/index/field_index/bool_index/read_only_bool_index/mod.rs index 550ae85f14..f13de732a4 100644 --- a/lib/segment/src/index/field_index/bool_index/read_only_bool_index/mod.rs +++ b/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(); diff --git a/lib/segment/src/index/field_index/geo_index/tests.rs b/lib/segment/src/index/field_index/geo_index/tests.rs index 2ff5c090d0..db4616333c 100644 --- a/lib/segment/src/index/field_index/geo_index/tests.rs +++ b/lib/segment/src/index/field_index/geo_index/tests.rs @@ -1345,6 +1345,7 @@ fn test_block_index_preopen() { OnDiskGeoIndex::::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(); diff --git a/lib/segment/src/index/field_index/numeric_index/tests.rs b/lib/segment/src/index/field_index/numeric_index/tests.rs index 3e9bb0dc9b..1851bab393 100644 --- a/lib/segment/src/index/field_index/numeric_index/tests.rs +++ b/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(); diff --git a/lib/segment/src/index/hnsw_index/graph_layers.rs b/lib/segment/src/index/hnsw_index/graph_layers.rs index 1d01d0b9d2..66b538286b 100644 --- a/lib/segment/src/index/hnsw_index/graph_layers.rs +++ b/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::::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() { diff --git a/lib/segment/src/segment/read_only/lifecycle.rs b/lib/segment/src/segment/read_only/lifecycle.rs index 67c18d1e07..ca7b1df521 100644 --- a/lib/segment/src/segment/read_only/lifecycle.rs +++ b/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: &Fs, segment_path: &Path, ) -> OperationResult> { 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 ReadOnlySegment { deferred_internal_id: Option, load_profile: Option<&LoadProfile>, ) -> OperationResult { + 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) — diff --git a/lib/segment/src/segment/read_only/live_reload.rs b/lib/segment/src/segment/read_only/live_reload.rs index 5523357dc8..8375a5a688 100644 --- a/lib/segment/src/segment/read_only/live_reload.rs +++ b/lib/segment/src/segment/read_only/live_reload.rs @@ -39,6 +39,10 @@ impl ReadOnlySegment { 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(()) } diff --git a/lib/segment/src/segment/update_only/lookup/lifecycle.rs b/lib/segment/src/segment/update_only/lookup/lifecycle.rs index 75616b599a..d3eaf11c76 100644 --- a/lib/segment/src/segment/update_only/lookup/lifecycle.rs +++ b/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: &Fs, segment_path: &Path, ) -> OperationResult> { 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) } diff --git a/lib/segment/src/vector_storage/chunked_vectors/read_only/mod.rs b/lib/segment/src/vector_storage/chunked_vectors/read_only/mod.rs index 4b502fafc8..a2252881b1 100644 --- a/lib/segment/src/vector_storage/chunked_vectors/read_only/mod.rs +++ b/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(); diff --git a/lib/segment/src/vector_storage/dense/read_only/mod.rs b/lib/segment/src/vector_storage/dense/read_only/mod.rs index 8b9db056a3..5abe9aaa86 100644 --- a/lib/segment/src/vector_storage/dense/read_only/mod.rs +++ b/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(); diff --git a/lib/segment/src/vector_storage/quantized/quantized_vectors/read_only/tests.rs b/lib/segment/src/vector_storage/quantized/quantized_vectors/read_only/tests.rs index 728053c946..323c00c97b 100644 --- a/lib/segment/src/vector_storage/quantized/quantized_vectors/read_only/tests.rs +++ b/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::::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(); diff --git a/lib/segment/src/vector_storage/read_only/mod.rs b/lib/segment/src/vector_storage/read_only/mod.rs index 08fd9dcfce..caf2c44b83 100644 --- a/lib/segment/src/vector_storage/read_only/mod.rs +++ b/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::::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::::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::::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 [ diff --git a/lib/segment/src/vector_storage/sparse/read_only/mod.rs b/lib/segment/src/vector_storage/sparse/read_only/mod.rs index 278c4017dc..6587cb1419 100644 --- a/lib/segment/src/vector_storage/sparse/read_only/mod.rs +++ b/lib/segment/src/vector_storage/sparse/read_only/mod.rs @@ -145,6 +145,7 @@ mod tests { cached_fs.cache_file_info().unwrap(); ReadOnlySparseVectorStorage::::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(); diff --git a/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs b/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs index ccf72179aa..ede62cf0bc 100644 --- a/lib/sparse/src/index/inverted_index/inverted_index_compressed_mmap.rs +++ b/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() {