@@ -13,7 +13,7 @@ use std::fs::{File, OpenOptions};
1313use std:: io:: { BufWriter , Error , Read , Result , Seek , SeekFrom , Write } ;
1414use std:: path:: { Path , PathBuf } ;
1515use std:: sync:: atomic:: { AtomicU64 , Ordering } ;
16- use std:: sync:: { Arc , Mutex , RwLock } ;
16+ use std:: sync:: { Arc , Mutex , RwLock , RwLockReadGuard } ;
1717use tracing:: { debug, info, warn} ;
1818
1919/// Append-Only Storage Engine
@@ -205,7 +205,13 @@ impl DataStore {
205205 {
206206 key_indexer_guard. remove ( key_hash) ;
207207 } else {
208- key_indexer_guard. insert ( * key_hash, * offset) ;
208+ // Handle the Result from the new insert method
209+ if let Err ( e) = key_indexer_guard. insert ( * key_hash, * offset) {
210+ // A collision was detected on write. The entire batch operation
211+ // should fail to prevent an inconsistent state.
212+ warn ! ( "Write operation aborted due to hash collision: {}" , e) ;
213+ return Err ( std:: io:: Error :: other ( e) ) ;
214+ }
209215 }
210216 }
211217
@@ -488,47 +494,54 @@ impl DataStore {
488494 Ok ( self . tail_offset . load ( Ordering :: Acquire ) )
489495 }
490496
491- /// Internal helper that does the real work for `read`/`batch_read`.
492- ///
493- /// * `key` – raw-byte key we are searching for.
494- /// * `mmap_arc` – the current shared memory-map.
495- /// * `key_indexer` – **already locked** read-only view of the index.
497+ /// Performs the core logic of reading an entry from the store.
496498///
497- /// The function:
498- /// 1. Hashes `key` with XXH3 (same as writers do).
499- /// 2. Looks the hash up in the index; bails out early if absent.
500- /// 3. Validates that the stored offset and metadata still fit inside the
501- /// current `mmap` (guards against truncated / corrupted files).
502- /// 4. Creates and returns an `EntryHandle` that spans the payload slice in
503- /// the `mmap`.
499+ /// This private helper centralizes the logic for both `read` and `batch_read`.
500+ /// It takes all necessary context to perform a safe lookup, including the key,
501+ /// its hash, the memory map, and a read guard for the key indexer.
504502///
505- /// It deliberately **does not** take any locks itself – that must be done by
506- /// the caller so that `batch_read` can reuse the same lock for many lookups.
503+ /// # Parameters
504+ /// - `key`: The original key bytes used for tag verification.
505+ /// - `key_hash`: The pre-computed hash of the key for index lookup.
506+ /// - `mmap_arc`: A reference to the active memory map.
507+ /// - `key_indexer_guard`: A read-lock guard for the key index.
507508///
508- /// `None` is returned when:
509- /// * the key is unknown,
510- /// * the mapped file looks inconsistent (bounds checks fail), or
511- /// * the latest record for the key is a tomb-stone (one-byte NULL payload) .
509+ /// # Returns
510+ /// - `Some(EntryHandle)` if the key is found and all checks pass.
511+ /// - `None` if the key is not found, a tag mismatch occurs (collision/corruption),
512+ /// or the entry is a tombstone .
512513#[ inline]
513- pub fn read_hashed_with_ctx (
514+ fn read_entry_with_context < ' a > (
515+ & self ,
516+ key : & [ u8 ] ,
514517 key_hash : u64 ,
515518 mmap_arc : & Arc < Mmap > ,
516- key_indexer : & KeyIndexer ,
519+ key_indexer_guard : & RwLockReadGuard < ' a , KeyIndexer > ,
517520 ) -> Option < EntryHandle > {
518- let offset = * key_indexer. get ( & key_hash) ?;
519- if offset as usize + METADATA_SIZE > mmap_arc. len ( ) {
521+ let packed = * key_indexer_guard. get_packed ( & key_hash) ?;
522+ let ( tag, offset) = KeyIndexer :: unpack ( packed) ;
523+
524+ // The crucial verification check, now centralized.
525+ if tag != KeyIndexer :: tag_from_key ( key) {
526+ warn ! ( "Tag mismatch detected for key, likely a hash collision or index corruption." ) ;
520527 return None ;
521528 }
522529
523- let metadata_bytes = & mmap_arc[ offset as usize ..offset as usize + METADATA_SIZE ] ;
524- let metadata = EntryMetadata :: deserialize ( metadata_bytes) ;
530+ let offset = offset as usize ;
531+ if offset + METADATA_SIZE > mmap_arc. len ( ) {
532+ return None ;
533+ }
525534
535+ let metadata_bytes = & mmap_arc[ offset..offset + METADATA_SIZE ] ;
536+ let metadata = EntryMetadata :: deserialize ( metadata_bytes) ;
526537 let entry_start = metadata. prev_offset as usize ;
527- let entry_end = offset as usize ;
538+ let entry_end = offset;
539+
528540 if entry_start >= entry_end || entry_end > mmap_arc. len ( ) {
529541 return None ;
530542 }
531543
544+ // Check for tombstone (deleted entry)
532545 if entry_end - entry_start == 1 && mmap_arc[ entry_start..entry_end] == NULL_BYTE {
533546 return None ;
534547 }
@@ -773,33 +786,7 @@ impl DataStoreReader for DataStore {
773786 . map_err ( |_| Error :: other ( "key-index lock poisoned" ) ) ?;
774787 let mmap_arc = self . get_mmap_arc ( ) ;
775788
776- let offset = match key_indexer_guard. get ( & key_hash) {
777- Some ( off) => * off, // found → continue
778- None => return Ok ( None ) , // not found → early-return
779- } ;
780-
781- if offset as usize + METADATA_SIZE > mmap_arc. len ( ) {
782- return Ok ( None ) ;
783- }
784-
785- let metadata_bytes = & mmap_arc[ offset as usize ..offset as usize + METADATA_SIZE ] ;
786- let metadata = EntryMetadata :: deserialize ( metadata_bytes) ;
787-
788- let entry_start = metadata. prev_offset as usize ;
789- let entry_end = offset as usize ;
790- if entry_start >= entry_end || entry_end > mmap_arc. len ( ) {
791- return Ok ( None ) ;
792- }
793-
794- if entry_end - entry_start == 1 && mmap_arc[ entry_start..entry_end] == NULL_BYTE {
795- return Ok ( None ) ;
796- }
797-
798- Ok ( Some ( EntryHandle {
799- mmap_arc : mmap_arc. clone ( ) ,
800- range : entry_start..entry_end,
801- metadata,
802- } ) )
789+ Ok ( self . read_entry_with_context ( key, key_hash, & mmap_arc, & key_indexer_guard) )
803790 }
804791
805792 fn read_last_entry ( & self ) -> Result < Option < EntryHandle > > {
@@ -836,35 +823,17 @@ impl DataStoreReader for DataStore {
836823 . key_indexer
837824 . read ( )
838825 . map_err ( |_| Error :: other ( "Key-index lock poisoned during `batch_read`" ) ) ?;
826+
839827 let hashes = compute_hash_batch ( keys) ;
840828
841829 let results = hashes
842830 . into_iter ( )
843- . map ( |h| {
844- key_indexer_guard. get ( & h) . and_then ( |offset| {
845- let offset = * offset as usize ;
846- if offset + METADATA_SIZE > mmap_arc. len ( ) {
847- return None ;
848- }
849- let metadata_bytes = & mmap_arc[ offset..offset + METADATA_SIZE ] ;
850- let metadata = EntryMetadata :: deserialize ( metadata_bytes) ;
851- let entry_start = metadata. prev_offset as usize ;
852- let entry_end = offset;
853- if entry_start >= entry_end || entry_end > mmap_arc. len ( ) {
854- return None ;
855- }
856- if entry_end - entry_start == 1 && mmap_arc[ entry_start..entry_end] == NULL_BYTE
857- {
858- return None ;
859- }
860- Some ( EntryHandle {
861- mmap_arc : mmap_arc. clone ( ) ,
862- range : entry_start..entry_end,
863- metadata,
864- } )
865- } )
831+ . zip ( keys. iter ( ) )
832+ . map ( |( key_hash, & key) | {
833+ self . read_entry_with_context ( key, key_hash, & mmap_arc, & key_indexer_guard)
866834 } )
867835 . collect ( ) ;
836+
868837 Ok ( results)
869838 }
870839
0 commit comments