[LivePreload] Preload id tracker (#10222)

* impl live_preload for read-only id trackers

* remove `raw_fs` argument

* fixup! use AtomicRefCell

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

@@ -25,7 +25,7 @@ impl<S: UniversalRead> ReadOnlyDiskIdTracker<S> {
         }
     }
 
-    fn deleted_open_options() -> OpenOptions {
+    pub(super) fn deleted_open_options() -> OpenOptions {
         OpenOptions {
             writeable: false,
             need_sequential: false,

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

@@ -3,7 +3,7 @@
 use common::bitvec::BitVec;
 use common::stored_bitslice::StoredBitSlice;
 use common::types::PointOffsetType;
-use common::universal_io::{UniversalRead, UniversalReadFs};
+use common::universal_io::{CachedReadFs, OkUnchanged, UniversalRead, UniversalReadFs};
 
 use super::ReadOnlyDiskIdTracker;
 use crate::common::operation_error::OperationResult;
@@ -11,6 +11,17 @@ use crate::id_tracker::immutable_id_tracker::deleted_path;
 use crate::id_tracker::mutable_id_tracker::read_only::LiveReloadResult;
 
 impl<S: UniversalRead> ReadOnlyDiskIdTracker<S> {
+    /// Stage the fresh deleted-bitslice handle [`live_reload`](Self::live_reload) swaps in.
+    pub fn live_preload(&self, fs: &impl CachedReadFs<File = S>) -> OperationResult<()> {
+        // The reload reads the whole bitslice
+        fs.reschedule_prefetch(
+            &deleted_path(&self.path),
+            Some(Self::deleted_open_options()),
+            None,
+        )?;
+        Ok(())
+    }
+
     /// Re-read the on-disk deleted bitslice and report points deleted since the
     /// last reload. Mappings are immutable, so nothing is ever inserted.
     ///
@@ -27,12 +38,20 @@ impl<S: UniversalRead> ReadOnlyDiskIdTracker<S> {
         &mut self,
         fs: &impl UniversalReadFs<File = S>,
     ) -> OperationResult<LiveReloadResult> {
-        let fresh = StoredBitSlice::<S>::open(
+        let Some(fresh) = StoredBitSlice::<S>::open(
             fs,
             deleted_path(&self.path),
             Self::open_options(),
             Default::default(),
-        )?;
+        )
+        .ok_unchanged()?
+        else {
+            return Ok(LiveReloadResult {
+                inserted: Vec::new(),
+                deleted: Vec::new(),
+            });
+        };
+
         let new: BitVec = fresh.read_all()?.into_owned();
         self.deleted_file = fresh;
 

lib/segment/src/id_tracker/disk_id_tracker/tests.rs

@@ -229,8 +229,7 @@ fn detect_and_load_selects_disk_format() {
     let _disk =
         DiskIdTracker::<MmapFile>::new(&MmapFs, disk_dir.path(), &versions, mappings).unwrap();
     let loaded =
-        ReadOnlyIdTrackerEnum::<MmapFile>::detect_and_load(&MmapFs, &MmapFs, disk_dir.path(), None)
-            .unwrap();
+        ReadOnlyIdTrackerEnum::<MmapFile>::detect_and_load(&MmapFs, disk_dir.path(), None).unwrap();
     assert_eq!(loaded.name(), "read-only disk id tracker");
     assert_read_parity(&immutable, &loaded);
 
@@ -240,19 +239,14 @@ fn detect_and_load_selects_disk_format() {
     let _imm = ImmutableIdTracker::<MmapFile>::new(&MmapFs, imm_dir.path(), &versions2, mappings2)
         .unwrap();
     let loaded =
-        ReadOnlyIdTrackerEnum::<MmapFile>::detect_and_load(&MmapFs, &MmapFs, imm_dir.path(), None)
-            .unwrap();
+        ReadOnlyIdTrackerEnum::<MmapFile>::detect_and_load(&MmapFs, imm_dir.path(), None).unwrap();
     assert_eq!(loaded.name(), "read-only immutable id tracker");
 
     // An empty segment (no mapping files) falls back to the appendable reader.
     let empty_dir = Builder::new().prefix("empty").tempdir().unwrap();
-    let loaded = ReadOnlyIdTrackerEnum::<MmapFile>::detect_and_load(
-        &MmapFs,
-        &MmapFs,
-        empty_dir.path(),
-        None,
-    )
-    .unwrap();
+    let loaded =
+        ReadOnlyIdTrackerEnum::<MmapFile>::detect_and_load(&MmapFs, empty_dir.path(), None)
+            .unwrap();
     assert_eq!(loaded.name(), "read-only appendable id tracker");
 }
 

lib/segment/src/id_tracker/id_tracker_base/read_only_tracker_enum.rs

@@ -33,24 +33,13 @@ impl<S: UniversalRead> ReadOnlyIdTrackerEnum<S> {
     }
 
     /// Detect the persisted id-tracker format and load it, by *attempting* each
-    /// format's open rather than probing file names one by one.
-    ///
-    /// This avoids the separate `exists` round-trips that a name-based detector
-    /// would issue — costly on object storage (S3/GCS/Azure), where each is a
-    /// remote request. Each candidate's open reads its own defining file, so a
-    /// not-found there simply means "not this format" and we fall through.
+    /// format's open.
     ///
     /// Order: disk-resident (the serverless/object-storage format) first, then
     /// the in-RAM immutable format, then the appendable/mutable format (whose
     /// open tolerates absent files, i.e. a fresh or empty segment).
-    ///
-    /// The attempts are sequential for now; they are independent and can be
-    /// issued concurrently later (the slow-path being remote opens).
-    /// `raw_fs` is the canonical backend for the appendable tracker's
-    /// bootstrap opens, which bypass any prefetch pool.
     pub fn detect_and_load(
         fs: &impl UniversalReadFs<File = S>,
-        raw_fs: &S::Fs,
         segment_path: &Path,
         deferred_internal_id: Option<PointOffsetType>,
     ) -> OperationResult<Self> {
@@ -61,7 +50,7 @@ impl<S: UniversalRead> ReadOnlyIdTrackerEnum<S> {
             return Ok(Self::Immutable(tracker));
         }
         Ok(Self::Appendable(ReadOnlyAppendableIdTracker::open(
-            raw_fs,
+            fs,
             segment_path,
             deferred_internal_id,
         )?))
@@ -69,9 +58,11 @@ impl<S: UniversalRead> ReadOnlyIdTrackerEnum<S> {
 
     /// Stage everything the next [`Self::live_reload`] needs. Shared access.
     pub fn live_preload(&self, fs: &impl CachedReadFs<File = S>) -> OperationResult<()> {
-        // todo(uio): dispatch per variant as the trackers gain live_preload
-        let _ = fs;
-        Ok(())
+        match self {
+            Self::Appendable(id_tracker) => id_tracker.live_preload(fs),
+            Self::Immutable(id_tracker) => id_tracker.live_preload(fs),
+            Self::DiskResident(id_tracker) => id_tracker.live_preload(fs),
+        }
     }
 
     /// Reload externally-applied changes, dispatching to the active variant.

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

@@ -1,6 +1,6 @@
 use common::stored_bitslice::StoredBitSlice;
 use common::types::PointOffsetType;
-use common::universal_io::{UniversalRead, UniversalReadFs};
+use common::universal_io::{CachedReadFs, OkUnchanged, UniversalRead, UniversalReadFs};
 
 use super::ReadOnlyImmutableIdTracker;
 use crate::common::operation_error::OperationResult;
@@ -8,6 +8,12 @@ use crate::id_tracker::immutable_id_tracker::deleted_storage::deleted_path;
 use crate::id_tracker::mutable_id_tracker::read_only::LiveReloadResult;
 
 impl<S: UniversalRead> ReadOnlyImmutableIdTracker<S> {
+    /// Stage the fresh deleted-bitslice handle [`live_reload`](Self::live_reload) swaps in.
+    pub fn live_preload(&self, fs: &impl CachedReadFs<File = S>) -> OperationResult<()> {
+        fs.reschedule_prefetch(&deleted_path(&self.path), Some(Self::open_options()), None)?;
+        Ok(())
+    }
+
     /// Re-read the on-disk `deleted` bitslice and apply points deleted since the last reload.
     ///
     /// The bitslice is a fixed-size bitmap whose bits the writer flips in
@@ -22,12 +28,19 @@ impl<S: UniversalRead> ReadOnlyImmutableIdTracker<S> {
         &mut self,
         fs: &impl UniversalReadFs<File = S>,
     ) -> OperationResult<LiveReloadResult> {
-        let fresh = StoredBitSlice::<S>::open(
+        let Some(fresh) = StoredBitSlice::<S>::open(
             fs,
             deleted_path(&self.path),
             Self::open_options(),
             Default::default(),
-        )?;
+        )
+        .ok_unchanged()?
+        else {
+            return Ok(LiveReloadResult {
+                inserted: Vec::new(),
+                deleted: Vec::new(),
+            });
+        };
 
         // `mappings` already reflects every previously reported deletion, so
         // it is the diff baseline: a set bit not yet dropped there is new.

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

@@ -15,7 +15,7 @@ use crate::id_tracker::point_mappings::PointMappings;
 use crate::types::PointIdType;
 
 impl<S: UniversalRead> ReadOnlyAppendableIdTracker<S> {
-    fn open_options() -> OpenOptions {
+    pub(super) fn open_options() -> OpenOptions {
         OpenOptions {
             writeable: false,
             need_sequential: false,
@@ -55,7 +55,7 @@ impl<S: UniversalRead> ReadOnlyAppendableIdTracker<S> {
     /// mappings log and versions file are consumed, applying only committed points (a partial
     /// trailing entry is simply not consumed and picked up on a later reload).
     pub fn open(
-        fs: &S::Fs,
+        fs: &impl UniversalReadFs<File = S>,
         segment_path: impl Into<PathBuf>,
         deferred_internal_id: Option<PointOffsetType>,
     ) -> OperationResult<Self> {

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

@@ -2,7 +2,7 @@ use std::io::Cursor;
 
 use common::generic_consts::Sequential;
 use common::types::PointOffsetType;
-use common::universal_io::{OkNotFound, ReadRange, UniversalRead, UniversalReadFs};
+use common::universal_io::{CachedReadFs, OkNotFound, ReadRange, UniversalRead, UniversalReadFs};
 
 use super::ReadOnlyAppendableIdTracker;
 use crate::common::operation_error::OperationResult;
@@ -59,6 +59,27 @@ impl LiveReloadResult {
 }
 
 impl<S: UniversalRead> ReadOnlyAppendableIdTracker<S> {
+    /// Stage what the next [`live_reload`](Self::live_reload) does per file: a
+    /// reopen for held handles, a prefetch for files it opens lazily. Absence
+    /// is tolerated the same way the reload tolerates it.
+    pub fn live_preload(&self, fs: &impl CachedReadFs<File = S>) -> OperationResult<()> {
+        let options = Self::open_options();
+        for (file, path) in [
+            (&self.versions_file, versions_path(&self.segment_path)),
+            (&self.mappings_file, mappings_path(&self.segment_path)),
+        ] {
+            match file {
+                Some(file) => file
+                    .schedule_reopen(|p| fs.cached_file_info(p))
+                    .ok_not_found()?,
+                None => fs
+                    .schedule_prefetch(&path, Some(options), None)
+                    .ok_not_found()?,
+            };
+        }
+        Ok(())
+    }
+
     /// Consume mapping and version changes appended to storage since the last reload.
     ///
     /// File handles are refreshed via [`UniversalRead::reopen`] so data appended by the writer

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

@@ -8,7 +8,6 @@ use common::types::PointOffsetType;
 use common::universal_io::{
     CachedFs, CachedReadFs, OkNotFound, Populate, UniversalReadFs, read_json_via,
 };
-use parking_lot::Mutex;
 use uuid::Uuid;
 
 use super::{ReadOnlySegment, ReadOnlyVectorData};
@@ -243,7 +242,6 @@ impl<S: UniversalReadExt + 'static> ReadOnlySegment<S> {
         // per-file `exists` round-trips — important for object-storage backends).
         let id_tracker = Arc::new(AtomicRefCell::new(ReadOnlyIdTrackerEnum::detect_and_load(
             &fs,
-            raw_fs,
             segment_path,
             deferred_internal_id,
         )?));
@@ -332,7 +330,7 @@ impl<S: UniversalReadExt + 'static> ReadOnlySegment<S> {
             payload_index,
             payload_storage,
             pending_reload: AtomicRefCell::new(Default::default()),
-            reload_fs: Mutex::new(fs),
+            reload_fs: AtomicRefCell::new(fs),
             segment_type,
             segment_config: config,
         })

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

@@ -27,7 +27,7 @@ impl<S: UniversalReadExt + 'static> ReadOnlySegment<S> {
             segment_config: _,
         } = self;
 
-        let mut reload_fs = reload_fs.lock();
+        let mut reload_fs = reload_fs.borrow_mut();
         // perf: one LIST per segment per refresh; could be a single shard-prefix
         // LIST partitioned into the per-segment snapshots.
         reload_fs.cache_file_info()?;

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

@@ -5,7 +5,6 @@ use std::sync::Arc;
 
 use atomic_refcell::{AtomicRef, AtomicRefCell};
 use common::universal_io::CachedFs;
-use parking_lot::Mutex;
 use uuid::Uuid;
 
 use crate::id_tracker::mutable_id_tracker::read_only::LiveReloadResult;
@@ -54,7 +53,7 @@ pub struct ReadOnlySegment<S: UniversalReadExt + 'static> {
     /// re-snapshots its listing and stages into its prefetch pool, [`live_reload`](ReadOnlySegment::live_reload)
     /// consumes what was staged. Retaining it gives unchanged-detection a previous
     /// listing to compare against — the open-time one for the first reload.
-    pub(crate) reload_fs: Mutex<CachedFs<S::Fs>>,
+    pub(crate) reload_fs: AtomicRefCell<CachedFs<S::Fs>>,
 
     /// Shows what kind of indexes and storages are used in this segment
     pub segment_type: SegmentType,

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

@@ -72,21 +72,17 @@ impl<S: UniversalRead + 'static> LookupSegment<S> {
         fs: &S::Fs,
         segment_path: &Path,
         deferred_internal_id: Option<PointOffsetType>,
-    ) -> OperationResult<Self>
-    where
-        S::Fs: UniversalReadFs<File = S>,
-    {
+    ) -> OperationResult<Self> {
         let cached_fs = build_cached_fs(fs, segment_path)?;
         let config = Self::preopen(&cached_fs, segment_path)?;
-        Self::open_via(&cached_fs, fs, segment_path, config, deferred_internal_id)
+        Self::open_via(&cached_fs, segment_path, config, deferred_internal_id)
     }
 
     /// Open the segment's components: the id tracker, the payload storage and
     /// one storage per named vector — nothing else.
     ///
     /// `fs` opens the component files (in production the [`CachedFs`] that
-    /// [`open`](Self::open) primed); `raw_fs` is the canonical backend, kept
-    /// by components that re-open files after this call. `config` is the one
+    /// [`open`](Self::open) primed). `config` is the one
     /// [`preopen`](Self::preopen) already parsed, so the state file is not
     /// read twice.
     ///
@@ -100,7 +96,6 @@ impl<S: UniversalRead + 'static> LookupSegment<S> {
     /// cutoff is ignored.
     pub fn open_via(
         fs: &impl UniversalReadFs<File = S>,
-        raw_fs: &S::Fs,
         segment_path: &Path,
         config: SegmentConfig,
         deferred_internal_id: Option<PointOffsetType>,
@@ -122,13 +117,11 @@ impl<S: UniversalRead + 'static> LookupSegment<S> {
 
         let appendable = config.is_appendable();
 
-        // Detect the persisted format by attempting each format's open (no
-        // per-file `exists` round-trips — important for object-storage
-        // backends). The deferred threshold applies to the appendable tracker
-        // only, mirroring `ReadOnlySegment::open_via`.
+        // Detect the persisted format by attempting each format's open. The
+        // deferred threshold applies to the appendable tracker only, mirroring
+        // `ReadOnlySegment::open_via`.
         let id_tracker = Arc::new(AtomicRefCell::new(ReadOnlyIdTrackerEnum::detect_and_load(
             fs,
-            raw_fs,
             segment_path,
             deferred_internal_id.filter(|_| appendable),
         )?));