iterator: fix error handling about valid (#411)

Signed-off-by: qupeng <[email protected]>

Author: hicqu

src/lib.rs

@@ -35,7 +35,7 @@ pub use metadata::{ColumnFamilyMetaData, LevelMetaData, SstFileMetaData};
 pub use perf_context::{get_perf_level, set_perf_level, IOStatsContext, PerfContext, PerfLevel};
 pub use rocksdb::{
     load_latest_options, run_ldb_tool, set_external_sst_file_global_seq_no, BackupEngine, CFHandle,
-    Cache, DBIterator, DBVector, Env, ExternalSstFileInfo, Kv, MapProperty, MemoryAllocator, Range,
+    Cache, DBIterator, DBVector, Env, ExternalSstFileInfo, MapProperty, MemoryAllocator, Range,
     SeekKey, SequentialFile, SstFileReader, SstFileWriter, Writable, WriteBatch, DB,
 };
 pub use rocksdb_options::{
@@ -54,6 +54,9 @@ pub use table_properties_collector::TablePropertiesCollector;
 pub use table_properties_collector_factory::TablePropertiesCollectorFactory;
 pub use titan::{TitanBlobIndex, TitanDBOptions};
 
+#[allow(deprecated)]
+pub use rocksdb::Kv;
+
 mod compaction_filter;
 pub mod comparator;
 mod event_listener;

src/perf_context.rs

@@ -422,10 +422,8 @@ mod test {
         let mut ctx = PerfContext::get();
 
         let mut iter = db.iter();
-        assert!(iter.seek(SeekKey::Start));
-        while iter.valid() {
-            iter.next();
-        }
+        assert!(iter.seek(SeekKey::Start).unwrap());
+        while iter.next().unwrap() {}
         assert_eq!(ctx.internal_key_skipped_count(), n);
         assert_eq!(ctx.internal_delete_skipped_count(), n / 2);
         assert_eq!(ctx.seek_internal_seek_time(), 0);
@@ -439,9 +437,9 @@ mod test {
         assert_eq!(get_perf_level(), PerfLevel::EnableTime);
 
         let mut iter = db.iter();
-        assert!(iter.seek(SeekKey::End));
-        while iter.valid() {
-            iter.prev();
+        assert!(iter.seek(SeekKey::End).unwrap());
+        while iter.valid().unwrap() {
+            iter.prev().unwrap();
         }
         assert_eq!(ctx.internal_key_skipped_count(), n + n / 2);
         assert_eq!(ctx.internal_delete_skipped_count(), n / 2);

src/rocksdb.rs

@@ -236,7 +236,7 @@ impl<D: Deref<Target = DB>> DBIterator<D> {
 }
 
 impl<D> DBIterator<D> {
-    pub fn seek(&mut self, key: SeekKey) -> bool {
+    pub fn seek(&mut self, key: SeekKey) -> Result<bool, String> {
         unsafe {
             match key {
                 SeekKey::Start => crocksdb_ffi::crocksdb_iter_seek_to_first(self.inner),
@@ -249,7 +249,7 @@ impl<D> DBIterator<D> {
         self.valid()
     }
 
-    pub fn seek_for_prev(&mut self, key: SeekKey) -> bool {
+    pub fn seek_for_prev(&mut self, key: SeekKey) -> Result<bool, String> {
         unsafe {
             match key {
                 SeekKey::Start => crocksdb_ffi::crocksdb_iter_seek_to_first(self.inner),
@@ -264,23 +264,24 @@ impl<D> DBIterator<D> {
         self.valid()
     }
 
-    pub fn prev(&mut self) -> bool {
+    pub fn prev(&mut self) -> Result<bool, String> {
         unsafe {
             crocksdb_ffi::crocksdb_iter_prev(self.inner);
         }
         self.valid()
     }
 
     #[allow(clippy::should_implement_trait)]
-    pub fn next(&mut self) -> bool {
+    pub fn next(&mut self) -> Result<bool, String> {
         unsafe {
             crocksdb_ffi::crocksdb_iter_next(self.inner);
         }
         self.valid()
     }
 
+    /// Get the key pointed by the iterator. Must be called when `self.valid() == Ok(true)`.
     pub fn key(&self) -> &[u8] {
-        assert!(self.valid());
+        debug_assert_eq!(self.valid(), Ok(true));
         let mut key_len: size_t = 0;
         let key_len_ptr: *mut size_t = &mut key_len;
         unsafe {
@@ -289,8 +290,9 @@ impl<D> DBIterator<D> {
         }
     }
 
+    /// Get the value pointed by the iterator. Must be called when `self.valid() == Ok(true)`.
     pub fn value(&self) -> &[u8] {
-        assert!(self.valid());
+        debug_assert_eq!(self.valid(), Ok(true));
         let mut val_len: size_t = 0;
         let val_len_ptr: *mut size_t = &mut val_len;
         unsafe {
@@ -299,37 +301,49 @@ impl<D> DBIterator<D> {
         }
     }
 
+    #[deprecated]
     pub fn kv(&self) -> Option<(Vec<u8>, Vec<u8>)> {
-        if self.valid() {
+        if self.valid().unwrap() {
             Some((self.key().to_vec(), self.value().to_vec()))
         } else {
             None
         }
     }
 
-    pub fn valid(&self) -> bool {
-        unsafe { crocksdb_ffi::crocksdb_iter_valid(self.inner) }
+    pub fn valid(&self) -> Result<bool, String> {
+        let valid = unsafe { crocksdb_ffi::crocksdb_iter_valid(self.inner) };
+        if !valid {
+            self.status()?;
+        }
+        Ok(valid)
     }
 
-    pub fn status(&self) -> Result<(), String> {
+    fn status(&self) -> Result<(), String> {
         unsafe {
             ffi_try!(crocksdb_iter_get_error(self.inner));
         }
         Ok(())
     }
 }
 
+#[deprecated]
 pub type Kv = (Vec<u8>, Vec<u8>);
 
+#[deprecated]
 impl<'b, D> Iterator for &'b mut DBIterator<D> {
+    #[allow(deprecated)]
     type Item = Kv;
 
-    fn next(&mut self) -> Option<Kv> {
-        let kv = self.kv();
-        if kv.is_some() {
-            DBIterator::next(self);
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.valid() {
+            Ok(true) => {}
+            Ok(false) => return None,
+            Err(e) => panic!("invalid iterator: {}", e),
         }
-        kv
+        let k = self.key().to_vec();
+        let v = self.value().to_vec();
+        let _ = DBIterator::next(self);
+        Some((k, v))
     }
 }
 
@@ -2778,7 +2792,7 @@ mod test {
         db.put(b"k2", b"v2222").expect("");
         db.put(b"k3", b"v3333").expect("");
         let mut iter = db.iter();
-        iter.seek(SeekKey::Start);
+        iter.seek(SeekKey::Start).unwrap();
         for (k, v) in &mut iter {
             println!(
                 "Hello {}: {}",

tests/cases/test_delete_files_in_range.rs

@@ -66,9 +66,9 @@ fn test_delete_files_in_range_with_iter() {
     db.delete_files_in_range(b"key2", b"key7", false).unwrap();
 
     let mut count = 0;
-    assert!(iter.seek(SeekKey::Start));
-    while iter.valid() {
-        iter.next();
+    assert!(iter.seek(SeekKey::Start).unwrap());
+    while iter.valid().unwrap() {
+        iter.next().unwrap();
         count = count + 1;
     }
 
@@ -89,11 +89,11 @@ fn test_delete_files_in_range_with_snap() {
     db.delete_files_in_range(b"key2", b"key7", false).unwrap();
 
     let mut iter = snap.iter();
-    assert!(iter.seek(SeekKey::Start));
+    assert!(iter.seek(SeekKey::Start).unwrap());
 
     let mut count = 0;
-    while iter.valid() {
-        iter.next();
+    while iter.valid().unwrap() {
+        iter.next().unwrap();
         count = count + 1;
     }
 
@@ -159,12 +159,12 @@ fn test_delete_files_in_range_with_delete_range() {
     db.compact_range(None, None);
 
     let mut it = db.iter();
-    it.seek(SeekKey::Start);
-    assert!(it.valid());
+    it.seek(SeekKey::Start).unwrap();
+    assert!(it.valid().unwrap());
     assert_eq!(it.key(), b"4");
-    assert!(it.next());
+    assert!(it.next().unwrap());
     assert_eq!(it.key(), b"5");
-    assert!(!it.next());
+    assert!(!it.next().unwrap());
 }
 
 #[test]
@@ -186,19 +186,18 @@ fn test_delete_files_in_ranges() {
 
     // Check that ["key0", "key5"] have been deleted, but ["key6", "key8"] still exist.
     let mut iter = db.iter();
-    iter.seek(SeekKey::Start);
+    iter.seek(SeekKey::Start).unwrap();
     for i in 6..9 {
-        assert!(iter.valid());
+        assert!(iter.valid().unwrap());
         let k = format!("key{}", i);
         assert_eq!(iter.key(), k.as_bytes());
-        iter.next();
+        iter.next().unwrap();
     }
-    assert!(!iter.valid());
+    assert!(!iter.valid().unwrap());
 
     // Delete the last file.
     let ranges = vec![Range::new(b"key6", b"key8")];
     db.delete_files_in_ranges_cf(cf, &ranges, true).unwrap();
     let mut iter = db.iter();
-    iter.seek(SeekKey::Start);
-    assert!(!iter.valid());
+    assert!(!iter.seek(SeekKey::Start).unwrap());
 }

tests/cases/test_delete_range.rs

@@ -44,37 +44,37 @@ fn gen_sst_from_db(opt: ColumnFamilyOptions, cf: Option<&CFHandle>, path: &str,
     };
     writer.open(path).unwrap();
     let mut iter = db.iter();
-    iter.seek(SeekKey::Start);
-    while iter.valid() {
+    iter.seek(SeekKey::Start).unwrap();
+    while iter.valid().unwrap() {
         writer.put(iter.key(), iter.value()).unwrap();
-        iter.next();
+        iter.next().unwrap();
     }
     writer.finish().unwrap();
 }
 
 fn gen_crc32_from_db(db: &DB) -> u32 {
     let mut digest = Digest::new(crc32::IEEE);
     let mut iter = db.iter();
-    iter.seek(SeekKey::Start);
-    while iter.valid() {
+    iter.seek(SeekKey::Start).unwrap();
+    while iter.valid().unwrap() {
         digest.write(iter.key());
         digest.write(iter.value());
-        iter.next();
+        iter.next().unwrap();
     }
     digest.sum32()
 }
 
 fn gen_crc32_from_db_in_range(db: &DB, start_key: &[u8], end_key: &[u8]) -> u32 {
     let mut digest = Digest::new(crc32::IEEE);
     let mut iter = db.iter();
-    iter.seek(SeekKey::Key(start_key));
-    while iter.valid() {
+    iter.seek(SeekKey::Key(start_key)).unwrap();
+    while iter.valid().unwrap() {
         if iter.key() >= end_key {
             break;
         }
         digest.write(iter.key());
         digest.write(iter.value());
-        iter.next();
+        iter.next().unwrap();
     }
     digest.sum32()
 }

tests/cases/test_ingest_external_file.rs

@@ -282,16 +282,16 @@ fn gen_sst_from_cf(opt: ColumnFamilyOptions, db: &DB, cf: &CFHandle, path: &str)
     let mut writer = SstFileWriter::new_cf(env_opt, opt, cf);
     writer.open(path).unwrap();
     let mut iter = db.iter_cf(cf);
-    iter.seek(SeekKey::Start);
-    while iter.valid() {
+    iter.seek(SeekKey::Start).unwrap();
+    while iter.valid().unwrap() {
         writer.put(iter.key(), iter.value()).unwrap();
-        iter.next();
+        iter.next().unwrap();
     }
     let info = writer.finish().unwrap();
     assert_eq!(info.file_path().to_str().unwrap(), path);
-    iter.seek(SeekKey::Start);
+    iter.seek(SeekKey::Start).unwrap();
     assert_eq!(info.smallest_key(), iter.key());
-    iter.seek(SeekKey::End);
+    iter.seek(SeekKey::End).unwrap();
     assert_eq!(info.largest_key(), iter.key());
     assert_eq!(info.sequence_number(), 0);
     assert!(info.file_size() > 0);
@@ -522,7 +522,7 @@ fn test_read_sst() {
         assert_eq!(props.num_entries(), 3);
     });
     let mut it = reader.iter();
-    it.seek(SeekKey::Start);
+    it.seek(SeekKey::Start).unwrap();
     assert_eq!(
         it.collect::<Vec<_>>(),
         vec![