Fix point read degradation in disjoint trees

src/level_manifest/level.rs

@@ -153,7 +153,11 @@ impl<'a> DisjointLevel<'a> {
             .segments
             .partition_point(|x| &*x.metadata.key_range.1 < key.as_ref());
 
-        level.segments.get(idx).cloned()
+        level
+            .segments
+            .get(idx)
+            .filter(|x| x.is_key_in_key_range(key))
+            .cloned()
     }
 
     pub fn range_indexes(
@@ -395,6 +399,38 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    #[allow(clippy::unwrap_used)]
+    fn level_disjoint_containing_key() -> crate::Result<()> {
+        let folder = tempfile::tempdir()?;
+
+        let tree = crate::Config::new(&folder).open()?;
+
+        for k in 'c'..'k' {
+            tree.insert([k as u8], "", 0);
+            tree.flush_active_memtable(0).expect("should flush");
+        }
+
+        let first = tree
+            .levels
+            .read()
+            .expect("lock is poisoned")
+            .levels
+            .first()
+            .expect("should exist")
+            .clone();
+
+        let dis = first.as_disjoint().unwrap();
+        assert!(dis.get_segment_containing_key("a").is_none());
+        assert!(dis.get_segment_containing_key("b").is_none());
+        for k in 'c'..'k' {
+            assert!(dis.get_segment_containing_key([k as u8]).is_some());
+        }
+        assert!(dis.get_segment_containing_key("l").is_none());
+
+        Ok(())
+    }
+
     #[test]
     fn level_not_disjoint() -> crate::Result<()> {
         let folder = tempfile::tempdir()?;

src/segment/mod.rs

@@ -344,6 +344,10 @@ impl Segment {
         Ok(Some(entry))
     }
 
+    pub fn is_key_in_key_range<K: AsRef<[u8]>>(&self, key: K) -> bool {
+        self.metadata.key_range.contains_key(key)
+    }
+
     // NOTE: Clippy false positive
     #[allow(unused)]
     /// Retrieves an item from the segment.
@@ -362,7 +366,7 @@ impl Segment {
             }
         }
 
-        if !self.metadata.key_range.contains_key(&key) {
+        if !self.is_key_in_key_range(&key) {
             return Ok(None);
         }
 

src/tree/mod.rs

@@ -654,10 +654,10 @@ impl Tree {
                             }
                             return Ok(Some(item));
                         }
-                    } else {
-                        // NOTE: Don't go to fallback, go to next level instead
-                        continue;
                     }
+
+                    // NOTE: Go to next level
+                    continue;
                 }
             }
 
@@ -890,11 +890,21 @@ impl Tree {
         };
 
         let tree_path = tree_path.as_ref();
-        log::debug!("Recovering disk segments from {tree_path:?}");
+
+        log::info!("Recovering LSM-tree at {tree_path:?}");
 
         let level_manifest_path = tree_path.join(LEVELS_MANIFEST_FILE);
 
         let segment_ids_to_recover = LevelManifest::recover_ids(&level_manifest_path)?;
+        let cnt = segment_ids_to_recover.len();
+
+        log::debug!("Recovering {cnt} disk segments from {tree_path:?}");
+
+        let progress_mod = match cnt {
+            _ if cnt <= 20 => 1,
+            _ if cnt <= 100 => 10,
+            _ => 100,
+        };
 
         let mut segments = vec![];
 
@@ -905,7 +915,7 @@ impl Tree {
             fsync_directory(&segment_base_folder)?;
         }
 
-        for dirent in std::fs::read_dir(&segment_base_folder)? {
+        for (idx, dirent) in std::fs::read_dir(&segment_base_folder)?.enumerate() {
             let dirent = dirent?;
 
             let file_name = dirent.file_name();
@@ -943,18 +953,22 @@ impl Tree {
 
                 segments.push(Arc::new(segment));
                 log::debug!("Recovered segment from {segment_file_path:?}");
+
+                if idx % progress_mod == 0 {
+                    log::debug!("Recovered {idx}/{cnt} disk segments");
+                }
             } else {
                 log::debug!("Deleting unfinished segment: {segment_file_path:?}",);
                 std::fs::remove_file(&segment_file_path)?;
             }
         }
 
         if segments.len() < segment_ids_to_recover.len() {
-            log::error!("Expected segments: {segment_ids_to_recover:?}");
+            log::error!("Recovered less segments than expected: {segment_ids_to_recover:?}");
             return Err(crate::Error::Unrecoverable);
         }
 
-        log::debug!("Recovered {} segments", segments.len());
+        log::debug!("Successfully recovered {} segments", segments.len());
 
         LevelManifest::recover(&level_manifest_path, segments)
     }