id tracker single write in flush

IvanPleshkov · IvanPleshkov · commit c1344d00fad3 · 2026-01-06T13:57:11.000+01:00
diff --git a/lib/common/io/src/counting_write.rs b/lib/common/io/src/counting_write.rs
@@ -0,0 +1,25 @@
+use std::io::{self, Write};
+
+/// A writer that counts the number of bytes written to it without actually storing them.
+#[derive(Default)]
+pub struct CountingWrite {
+    bytes: u64,
+}
+
+impl CountingWrite {
+    /// Returns the total number of bytes that have been "written" to this writer.
+    pub fn bytes_written(&self) -> u64 {
+        self.bytes
+    }
+}
+
+impl Write for CountingWrite {
+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        self.bytes += buf.len() as u64;
+        Ok(buf.len()) // pretend we wrote everything
+    }
+
+    fn flush(&mut self) -> io::Result<()> {
+        Ok(())
+    }
+}
diff --git a/lib/common/io/src/lib.rs b/lib/common/io/src/lib.rs
@@ -1,3 +1,4 @@
+pub mod counting_write;
 pub mod file_operations;
 pub mod move_files;
 pub mod safe_delete;
diff --git a/lib/segment/src/id_tracker/mutable_id_tracker.rs b/lib/segment/src/id_tracker/mutable_id_tracker.rs
@@ -3,6 +3,7 @@ use std::io::{self, BufReader, BufWriter, Read, Seek, Write};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 
+use ::io::counting_write::CountingWrite;
 use bitvec::prelude::{BitSlice, BitVec};
 use byteorder::{ReadBytesExt, WriteBytesExt};
 use common::is_alive_lock::IsAliveLock;
@@ -16,6 +17,7 @@ use uuid::Uuid;
 use super::point_mappings::FileEndianess;
 use crate::common::Flusher;
 use crate::common::operation_error::{OperationError, OperationResult};
+use crate::common::vector_utils::TrySetCapacityExact;
 use crate::id_tracker::point_mappings::PointMappings;
 use crate::id_tracker::{DELETED_POINT_VERSION, IdTracker};
 use crate::types::{PointIdType, SeqNumberType};
@@ -399,23 +401,57 @@ fn store_mapping_changes(
     mappings_path: &Path,
     changes: &Vec<MappingChange>,
 ) -> OperationResult<()> {
+    // Estimate required buffer size
+    let mut counting_write = CountingWrite::default();
+    write_mapping_changes(&mut counting_write, changes).map_err(|err| {
+        OperationError::service_error(format!(
+            "Failed to estimate ID tracker point mappings size ({}): {err}",
+            mappings_path.display(),
+        ))
+    })?;
+
+    // Create a buffer with exact required size
+    let mut writer: Vec<u8> = Vec::default();
+    writer
+        .try_set_capacity_exact(counting_write.bytes_written() as usize)
+        .map_err(|err| {
+            OperationError::service_error(format!(
+                "Failed to allocate buffer for ID tracker point mappings ({}): {err}",
+                mappings_path.display(),
+            ))
+        })?;
+
+    // Collect all changes into a single buffer first
+    write_mapping_changes(&mut writer, changes).map_err(|err| {
+        OperationError::service_error(format!(
+            "Failed to collect ID tracker point mappings ({}): {err}",
+            mappings_path.display(),
+        ))
+    })?;
+
     // Create or open file in append mode to write new changes to the end
     let file = File::options()
         .create(true)
         .append(true)
         .open(mappings_path)?;
-    let mut writer = BufWriter::new(file);
 
-    write_mapping_changes(&mut writer, changes).map_err(|err| {
+    // Write all changes in one shot
+    let mut file_writer = BufWriter::new(file);
+    file_writer.write_all(&writer).map_err(|err| {
         OperationError::service_error(format!(
             "Failed to persist ID tracker point mappings ({}): {err}",
             mappings_path.display(),
         ))
     })?;
 
+    // Drop the buffer to free memory
+    drop(writer);
+
     // Explicitly fsync file contents to ensure durability
-    writer.flush()?;
-    let file = writer.into_inner().unwrap();
+    // `BufWriter::into_inner` flushes the buffer as well
+    let file = file_writer.into_inner().map_err(|err| {
+        OperationError::service_error(format!("Failed to flush ID tracker point mappings: {err}"))
+    })?;
     file.sync_all().map_err(|err| {
         OperationError::service_error(format!("Failed to fsync ID tracker point mappings: {err}"))
     })?;
@@ -438,10 +474,6 @@ fn write_mapping_changes<W: Write>(
     for &change in changes {
         write_entry(&mut writer, change)?;
     }
-
-    // Explicitly flush writer to catch IO errors
-    writer.flush()?;
-
     Ok(())
 }
 

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+pub mod counting_write;`
`1`	`2`	`pub mod file_operations;`
`2`	`3`	`pub mod move_files;`
`3`	`4`	`pub mod safe_delete;`