From b1e9c9a9dfef506b44e008df573c4eefaeaa0dcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luis=20Coss=C3=ADo?= Date: Wed, 8 Apr 2026 13:09:55 -0400 Subject: [PATCH] [UIO] migrate `SliceBufferedUpdateWrapper` (#8518) * migrate SliceBufferedUpdateWrapper * fix compressed versions creation * @xzfc's review improvements Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com> * fix rebase * Drop usage of `TypedStorage` --------- Co-authored-by: xzfc <5121426+xzfc@users.noreply.github.com> Co-authored-by: xzfc --- .../universal_io/wrappers/buffered_update.rs} | 71 +++++++++++-------- .../common/src/universal_io/wrappers/mod.rs | 2 + lib/segment/src/common/mod.rs | 1 - .../src/id_tracker/immutable_id_tracker.rs | 62 +++++++++------- 4 files changed, 83 insertions(+), 53 deletions(-) rename lib/{segment/src/common/mmap_slice_buffered_update_wrapper.rs => common/common/src/universal_io/wrappers/buffered_update.rs} (51%) diff --git a/lib/segment/src/common/mmap_slice_buffered_update_wrapper.rs b/lib/common/common/src/universal_io/wrappers/buffered_update.rs similarity index 51% rename from lib/segment/src/common/mmap_slice_buffered_update_wrapper.rs rename to lib/common/common/src/universal_io/wrappers/buffered_update.rs index 0853765dea..f4f3e55d4d 100644 --- a/lib/segment/src/common/mmap_slice_buffered_update_wrapper.rs +++ b/lib/common/common/src/universal_io/wrappers/buffered_update.rs @@ -1,41 +1,40 @@ use std::sync::Arc; use ahash::AHashMap; -use common::is_alive_lock::IsAliveLock; -use common::mmap::MmapSlice; -use common::types::PointOffsetType; use parking_lot::{Mutex, RwLock}; -use crate::common::Flusher; +use crate::is_alive_lock::IsAliveLock; +use crate::types::PointOffsetType; +use crate::universal_io::{Flusher, UniversalIoError, UniversalWrite}; -/// A wrapper around `MmapSlice` that delays writing changes to the underlying file until they get -/// flushed manually. -/// This expects the underlying MmapSlice not to grow in size. +/// A wrapper around [`UniversalWrite`] that delays writing changes to the +/// underlying file until they get flushed manually. +/// This expects the underlying storage not to grow in size. /// /// WARN: this structure is expected to be write-only. #[derive(Debug)] -pub struct MmapSliceBufferedUpdateWrapper +pub struct SliceBufferedUpdateWrapper, T: Copy> where T: 'static, { - mmap_slice: Arc>>, - len: usize, + slice: Arc>, + len: u64, pending_updates: Arc>>, is_alive_lock: IsAliveLock, } -impl MmapSliceBufferedUpdateWrapper +impl, T: Copy> SliceBufferedUpdateWrapper where T: 'static, { - pub fn new(mmap_slice: MmapSlice) -> Self { - let len = mmap_slice.len(); - Self { - mmap_slice: Arc::new(RwLock::new(mmap_slice)), + pub fn new(slice_storage: S) -> Result { + let len = slice_storage.len()?; + Ok(Self { + slice: Arc::new(RwLock::new(slice_storage)), len, pending_updates: Arc::new(Mutex::new(AHashMap::new())), is_alive_lock: IsAliveLock::new(), - } + }) } /// Sets the item at `index` to `value` buffered. @@ -44,7 +43,7 @@ where /// Panics if the index is out of bounds. pub fn set(&self, index: PointOffsetType, value: T) { assert!( - (index as usize) < self.len, + u64::from(index) < self.len, "index {index} out of range: {}", self.len ); @@ -52,9 +51,10 @@ where } } -impl MmapSliceBufferedUpdateWrapper +impl SliceBufferedUpdateWrapper where - T: 'static + Sync + Send + Clone + PartialEq, + S: UniversalWrite + Send + Sync + 'static, + T: Sync + Send + Copy + Clone + PartialEq + 'static, { pub fn flusher(&self) -> Flusher { let updates = { @@ -66,7 +66,7 @@ where }; let pending_updates_weak = Arc::downgrade(&self.pending_updates); - let slice = Arc::downgrade(&self.mmap_slice); + let slice = Arc::downgrade(&self.slice); let is_alive_handle = self.is_alive_lock.handle(); Box::new(move || { let (Some(is_alive_guard), Some(pending_updates_arc), Some(slice)) = ( @@ -74,17 +74,32 @@ where pending_updates_weak.upgrade(), slice.upgrade(), ) else { - log::debug!( - "Aborted flushing on a dropped MmapSliceBufferedUpdateWrapper instance" - ); + log::debug!("Aborted flushing on a dropped SliceBufferedUpdateWrapper instance"); return Ok(()); }; - let mut mmap_slice_write = slice.write(); - for (&index, value) in &updates { - mmap_slice_write[index as usize] = value.clone(); - } - mmap_slice_write.flusher()()?; + let mut slice_guard = slice.write(); + + // Coalesce contiguous updates into the same write + let mut items: Vec<(PointOffsetType, T)> = + updates.iter().map(|(k, v)| (*k, *v)).collect(); + items.sort_by_key(|(index, _)| *index); + + let all_values: Vec = items.iter().map(|(_, value)| *value).collect(); + + // Batch writes + let mut remaining_values: &[T] = &all_values[..]; + let it = items.chunk_by(|(a, _), (b, _)| *a + 1 == *b).map(|chunk| { + let elem_start: PointOffsetType = chunk[0].0; + let byte_start = u64::from(elem_start) * size_of::() as u64; + let chunk_values = remaining_values + .split_off(..chunk.len()) + .expect("`chunk.len()` is sourced from the same slice as `remaining_values`"); + (byte_start, chunk_values) + }); + slice_guard.write_batch(it)?; + + slice_guard.flusher()()?; // Keep the guard till here to prevent concurrent drop/flushes // We don't touch files from here on and can drop the alive guard diff --git a/lib/common/common/src/universal_io/wrappers/mod.rs b/lib/common/common/src/universal_io/wrappers/mod.rs index fc2634364d..71a66501da 100644 --- a/lib/common/common/src/universal_io/wrappers/mod.rs +++ b/lib/common/common/src/universal_io/wrappers/mod.rs @@ -1,5 +1,7 @@ +mod buffered_update; mod read_only; mod typed; +pub use buffered_update::SliceBufferedUpdateWrapper; pub use read_only::ReadOnly; pub use typed::TypedStorage; diff --git a/lib/segment/src/common/mod.rs b/lib/segment/src/common/mod.rs index b9146b4309..b8cc5dccbd 100644 --- a/lib/segment/src/common/mod.rs +++ b/lib/segment/src/common/mod.rs @@ -3,7 +3,6 @@ pub mod error_logging; pub mod flags; pub mod macros; pub mod mmap_bitslice_buffered_update_wrapper; -pub mod mmap_slice_buffered_update_wrapper; pub mod operation_error; pub mod operation_time_statistics; pub mod reciprocal_rank_fusion; diff --git a/lib/segment/src/id_tracker/immutable_id_tracker.rs b/lib/segment/src/id_tracker/immutable_id_tracker.rs index 698cede35a..3f5900c35c 100644 --- a/lib/segment/src/id_tracker/immutable_id_tracker.rs +++ b/lib/segment/src/id_tracker/immutable_id_tracker.rs @@ -4,15 +4,16 @@ use std::path::{Path, PathBuf}; use byteorder::{ReadBytesExt, WriteBytesExt}; use common::bitvec::{BitSlice, BitSliceExt as _, BitVec}; -use common::mmap::{AdviceSetting, MmapSlice, create_and_ensure_length, open_write_mmap}; +use common::mmap::create_and_ensure_length; use common::types::PointOffsetType; -use common::universal_io::OpenOptions; +use common::universal_io::{ + MmapFile, OpenOptions, SliceBufferedUpdateWrapper, TypedStorage, UniversalRead, UniversalWrite, +}; use fs_err::File; use uuid::Uuid; use crate::common::Flusher; use crate::common::mmap_bitslice_buffered_update_wrapper::MmapBitSliceBufferedUpdateWrapper; -use crate::common::mmap_slice_buffered_update_wrapper::MmapSliceBufferedUpdateWrapper; use crate::common::operation_error::{OperationError, OperationResult}; use crate::common::stored_bitslice::MmapBitSlice; use crate::id_tracker::compressed::compressed_point_mappings::CompressedPointMappings; @@ -59,7 +60,7 @@ pub struct ImmutableIdTracker { deleted_wrapper: MmapBitSliceBufferedUpdateWrapper, internal_to_version: CompressedVersions, - internal_to_version_wrapper: MmapSliceBufferedUpdateWrapper, + internal_to_version_wrapper: SliceBufferedUpdateWrapper, mappings: CompressedPointMappings, } @@ -257,16 +258,23 @@ impl ImmutableIdTracker { let deleted_wrapper = MmapBitSliceBufferedUpdateWrapper::new(deleted_storage); - let internal_to_version_map = open_write_mmap( - &Self::version_mapping_file_path(segment_path), - AdviceSetting::Global, - true, + let internal_to_version_file = TypedStorage::::open( + Self::version_mapping_file_path(segment_path), + OpenOptions { + writeable: true, + need_sequential: false, + disk_parallel: None, + populate: Some(true), + advice: None, + prevent_caching: None, + }, )?; - let internal_to_version_mapslice: MmapSlice = - unsafe { MmapSlice::try_from(internal_to_version_map)? }; - let internal_to_version = CompressedVersions::from_slice(&internal_to_version_mapslice); + + let internal_to_version_slice = internal_to_version_file.read_whole()?; + + let internal_to_version = CompressedVersions::from_slice(&internal_to_version_slice); let internal_to_version_wrapper = - MmapSliceBufferedUpdateWrapper::new(internal_to_version_mapslice); + SliceBufferedUpdateWrapper::new(internal_to_version_file.inner)?; let reader = BufReader::new(File::open(Self::mappings_file_path(segment_path))?); let mappings = Self::load_mapping(reader, Some(deleted_bitvec))?; @@ -326,22 +334,27 @@ impl ImmutableIdTracker { let version_size = mmap_size::(min_size); create_and_ensure_length(&version_filepath, version_size)?; } - let mut internal_to_version_wrapper = unsafe { - MmapSlice::try_from(open_write_mmap( - &version_filepath, - AdviceSetting::Global, - false, - )?)? - }; - internal_to_version_wrapper[..internal_to_version.len()] - .copy_from_slice(internal_to_version); - let internal_to_version = CompressedVersions::from_slice(&internal_to_version_wrapper); + let mut internal_to_version_file = TypedStorage::::open( + &version_filepath, + OpenOptions { + writeable: true, + need_sequential: false, + disk_parallel: None, + populate: Some(false), + advice: None, + prevent_caching: None, + }, + )?; + internal_to_version_file.write(0, internal_to_version)?; + + let internal_to_version = + CompressedVersions::from_slice(&internal_to_version_file.read_whole()?); debug_assert_eq!(internal_to_version.len(), mappings.total_point_count()); let internal_to_version_wrapper = - MmapSliceBufferedUpdateWrapper::new(internal_to_version_wrapper); + SliceBufferedUpdateWrapper::new(internal_to_version_file.inner)?; // Write mappings to disk. let file = File::create(Self::mappings_file_path(path))?; @@ -457,7 +470,8 @@ impl IdTracker for ImmutableIdTracker { /// Creates a flusher function, that writes the points versions to disk. fn versions_flusher(&self) -> Flusher { - self.internal_to_version_wrapper.flusher() + let flusher = self.internal_to_version_wrapper.flusher(); + Box::new(move || flusher().map_err(OperationError::from)) } fn total_point_count(&self) -> usize {