Fix GraphLinksConverter::serialize_to data alignment issues (#3806)

* Assert data alignment in `transmute_from_u8*` functions

* WIP: Switch `transmute_from_u8*` alignment asserts from `debug_assert` to `assert`...

...to make sure *all* tests running on CI will enforce the alignment

* Add descriptive message to alignment assertions

* Fix alignment in graph links file (#3807)

* fix alignment in graph links file

* fix alignment while reading

* Revert "fix alignment while reading"

This reverts commit e2d1cee890.

* Revert "Revert "fix alignment while reading""

This reverts commit 7a4fdc9aea.

* small refactor

* Switch `transmute_from_u8*` alignment asserts back to `debug_assert`

Additionally:
- Remove extra `0x` from slice address in assert messages

* Trim offsets padding bytes from the end of `links_mmap` slice

---------

Co-authored-by: Ivan Pleshkov <pleshkov.ivan@gmail.com>

lib/common/memory/src/mmap_ops.rs

@@ -1,6 +1,6 @@
 use std::fs::OpenOptions;
 use std::hint::black_box;
-use std::mem::size_of;
+use std::mem::{align_of, size_of};
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
 use std::{io, mem, ops, time};
@@ -96,6 +96,19 @@ where
 }
 pub fn transmute_from_u8<T>(v: &[u8]) -> &T {
     debug_assert_eq!(v.len(), size_of::<T>());
+
+    debug_assert_eq!(
+        v.as_ptr().align_offset(align_of::<T>()),
+        0,
+        "transmuting byte slice {:p} into {}: \
+         required alignment is {} bytes, \
+         byte slice misaligned by {} bytes",
+        v.as_ptr(),
+        std::any::type_name::<T>(),
+        align_of::<T>(),
+        v.as_ptr().align_offset(align_of::<T>()),
+    );
+
     unsafe { &*(v.as_ptr() as *const T) }
 }
 
@@ -105,6 +118,19 @@ pub fn transmute_to_u8<T>(v: &T) -> &[u8] {
 
 pub fn transmute_from_u8_to_slice<T>(data: &[u8]) -> &[T] {
     debug_assert_eq!(data.len() % size_of::<T>(), 0);
+
+    debug_assert_eq!(
+        data.as_ptr().align_offset(align_of::<T>()),
+        0,
+        "transmuting byte slice {:p} into slice of {}: \
+         required alignment is {} bytes, \
+         byte slice misaligned by {} bytes",
+        data.as_ptr(),
+        std::any::type_name::<T>(),
+        align_of::<T>(),
+        data.as_ptr().align_offset(align_of::<T>()),
+    );
+
     let len = data.len() / size_of::<T>();
     let ptr = data.as_ptr() as *const T;
     unsafe { std::slice::from_raw_parts(ptr, len) }
@@ -112,6 +138,19 @@ pub fn transmute_from_u8_to_slice<T>(data: &[u8]) -> &[T] {
 
 pub fn transmute_from_u8_to_mut_slice<T>(data: &mut [u8]) -> &mut [T] {
     debug_assert_eq!(data.len() % size_of::<T>(), 0);
+
+    debug_assert_eq!(
+        data.as_ptr().align_offset(align_of::<T>()),
+        0,
+        "transmuting byte slice {:p} into mutable slice of {}: \
+         required alignment is {} bytes, \
+         byte slice misaligned by {} bytes",
+        data.as_ptr(),
+        std::any::type_name::<T>(),
+        align_of::<T>(),
+        data.as_ptr().align_offset(align_of::<T>()),
+    );
+
     let len = data.len() / size_of::<T>();
     let ptr = data.as_mut_ptr() as *mut T;
     unsafe { std::slice::from_raw_parts_mut(ptr, len) }

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

@@ -1,6 +1,6 @@
 use std::cmp::max;
 use std::fs::OpenOptions;
-use std::mem::size_of;
+use std::mem::{self, size_of};
 use std::ops::Range;
 use std::path::{Path, PathBuf};
 use std::sync::Arc;
@@ -52,6 +52,7 @@ struct GraphLinksFileHeader {
     pub levels_count: u64,
     pub total_links_len: u64,
     pub total_offsets_len: u64,
+    pub offsets_padding: u64,
 }
 
 fn get_reindex_slice<'a>(
@@ -69,10 +70,18 @@ fn get_links_slice<'a>(data: &'a [u8], header: &'a GraphLinksFileHeader) -> &'a
     mmap_ops::transmute_from_u8_to_slice(links_byte_slice)
 }
 
-fn get_offsets_slice<'a>(data: &'a [u8], header: &'a GraphLinksFileHeader) -> &'a [u64] {
+fn get_offsets_iter<'a>(
+    data: &'a [u8],
+    header: &'a GraphLinksFileHeader,
+) -> impl Iterator<Item = u64> + 'a {
     let offsets_range = header.get_offsets_range();
-    let offsets_byte_slice = &data[offsets_range];
-    mmap_ops::transmute_from_u8_to_slice(offsets_byte_slice)
+    data[offsets_range]
+        .chunks_exact(mem::size_of::<u64>())
+        .map(|chunk| {
+            // unwrap is safe because we know that chunk is always 8 bytes
+            let bytes: [u8; 8] = chunk.try_into().unwrap();
+            u64::from_ne_bytes(bytes)
+        })
 }
 
 fn get_level_offsets<'a>(data: &'a [u8], header: &GraphLinksFileHeader) -> &'a [u64] {
@@ -82,8 +91,28 @@ fn get_level_offsets<'a>(data: &'a [u8], header: &GraphLinksFileHeader) -> &'a [
 }
 
 impl GraphLinksFileHeader {
+    pub fn new(
+        point_count: usize,
+        levels_count: usize,
+        total_links_len: usize,
+        total_offsets_len: usize,
+    ) -> GraphLinksFileHeader {
+        let offsets_padding = if (point_count + total_links_len) % 2 == 0 {
+            0
+        } else {
+            4
+        };
+        GraphLinksFileHeader {
+            point_count: point_count as u64,
+            levels_count: levels_count as u64,
+            total_links_len: total_links_len as u64,
+            total_offsets_len: total_offsets_len as u64,
+            offsets_padding,
+        }
+    }
+
     pub fn raw_size() -> usize {
-        size_of::<u64>() * 4
+        size_of::<u64>() * 5
     }
 
     pub fn serialize_bytes_to(&self, raw_data: &mut [u8]) {
@@ -93,6 +122,7 @@ impl GraphLinksFileHeader {
         arr[1] = self.levels_count;
         arr[2] = self.total_links_len;
         arr[3] = self.total_offsets_len;
+        arr[4] = self.offsets_padding;
     }
 
     pub fn deserialize_bytes_from(raw_data: &[u8]) -> GraphLinksFileHeader {
@@ -103,6 +133,7 @@ impl GraphLinksFileHeader {
             levels_count: arr[1],
             total_links_len: arr[2],
             total_offsets_len: arr[3],
+            offsets_padding: arr[4],
         }
     }
 
@@ -128,7 +159,7 @@ impl GraphLinksFileHeader {
     }
 
     pub fn get_offsets_range(&self) -> Range<usize> {
-        let start = self.get_links_range().end;
+        let start = self.get_links_range().end + self.offsets_padding as usize;
         start..start + self.total_offsets_len as usize * size_of::<u64>()
     }
 }
@@ -192,12 +223,12 @@ impl GraphLinksConverter {
     }
 
     fn get_header(&self) -> GraphLinksFileHeader {
-        GraphLinksFileHeader {
-            point_count: self.reindex.len() as u64,
-            levels_count: self.get_levels_count() as u64,
-            total_links_len: self.total_links_len as u64,
-            total_offsets_len: self.total_offsets_len as u64,
-        }
+        GraphLinksFileHeader::new(
+            self.reindex.len(),
+            self.get_levels_count(),
+            self.total_links_len,
+            self.total_offsets_len,
+        )
     }
 
     /// Size of compacted graph in bytes.
@@ -224,9 +255,10 @@ impl GraphLinksConverter {
             let links_range = header.get_links_range();
             let offsets_range = header.get_offsets_range();
             let union_range = links_range.start..offsets_range.end;
-            let (links_mmap, offsets_mmap) = bytes_data[union_range]
+            let (links_mmap, offsets_with_padding_mmap) = bytes_data[union_range]
                 .as_mut()
                 .split_at_mut(links_range.len());
+            let offsets_mmap = &mut offsets_with_padding_mmap[header.offsets_padding as _..];
             let links_mmap: &mut [PointOffsetType] =
                 mmap_ops::transmute_from_u8_to_mut_slice(links_mmap);
             let offsets_mmap: &mut [u64] = mmap_ops::transmute_from_u8_to_mut_slice(offsets_mmap);
@@ -404,9 +436,8 @@ impl GraphLinksRam {
         links.try_set_capacity_exact(link_slice.len())?;
         links.extend_from_slice(link_slice);
 
-        let offsets_slice = get_offsets_slice(data, &header);
-        offsets.try_set_capacity_exact(offsets_slice.len())?;
-        offsets.extend_from_slice(offsets_slice);
+        offsets.try_set_capacity_exact(header.get_offsets_range().len() / size_of::<u64>())?;
+        offsets.extend(get_offsets_iter(data, &header));
 
         let level_offsets_slice = get_level_offsets(data, &header);
         level_offsets.try_set_capacity_exact(level_offsets_slice.len())?;
@@ -499,16 +530,17 @@ impl GraphLinksMmap {
         if let Some(mmap) = &self.mmap {
             get_links_slice(mmap, &self.header)
         } else {
-            panic!("{}", "Mmap links are not loaded");
+            panic!("{}", MMAP_PANIC_MESSAGE);
         }
     }
 
-    fn get_offsets_slice(&self) -> &[u64] {
-        if let Some(mmap) = &self.mmap {
-            get_offsets_slice(mmap, &self.header)
-        } else {
-            panic!("{}", MMAP_PANIC_MESSAGE);
-        }
+    fn get_links_offset(offsets_data: &[u8], idx: usize) -> usize {
+        let begin = mem::size_of::<u64>() * idx;
+        let end = begin + mem::size_of::<u64>();
+        let bytes = &offsets_data[begin..end];
+        // unwrap is safe because we know that bytes slice is always 8 bytes
+        let bytes: [u8; 8] = bytes.try_into().unwrap();
+        u64::from_ne_bytes(bytes) as usize
     }
 
     pub fn prefault_mmap_pages(&self, path: &Path) -> Option<mmap_ops::PrefaultMmapPages> {
@@ -560,8 +592,14 @@ impl GraphLinks for GraphLinksMmap {
     }
 
     fn get_links_range(&self, idx: usize) -> Range<usize> {
-        let offsets_slice = self.get_offsets_slice();
-        offsets_slice[idx] as usize..offsets_slice[idx + 1] as usize
+        let offsets_range = self.header.get_offsets_range();
+        let mmap: &[u8] = if let Some(mmap) = &self.mmap {
+            &mmap[offsets_range]
+        } else {
+            panic!("{}", MMAP_PANIC_MESSAGE);
+        };
+
+        Self::get_links_offset(mmap, idx)..Self::get_links_offset(mmap, idx + 1)
     }
 
     fn get_level_offset(&self, level: usize) -> usize {