Merge pull request #180 from ekmartin/memtable_factory

Add support for customizing the memtable factory

src/db.rs

@@ -888,6 +888,15 @@ impl DB {
         DBIterator::new(self, &opts, mode)
     }
 
+    /// Opens an interator with `set_total_order_seek` enabled.
+    /// This must be used to iterate across prefixes when `set_memtable_factory` has been called
+    /// with a Hash-based implementation.
+    pub fn full_iterator(&self, mode: IteratorMode) -> DBIterator {
+        let mut opts = ReadOptions::default();
+        opts.set_total_order_seek(true);
+        DBIterator::new(self, &opts, mode)
+    }
+
     pub fn prefix_iterator<'a>(&self, prefix: &'a [u8]) -> DBIterator {
         let mut opts = ReadOptions::default();
         opts.set_prefix_same_as_start(true);
@@ -903,6 +912,16 @@ impl DB {
         DBIterator::new_cf(self, cf_handle, &opts, mode)
     }
 
+    pub fn full_iterator_cf(
+        &self,
+        cf_handle: ColumnFamily,
+        mode: IteratorMode,
+    ) -> Result<DBIterator, Error> {
+        let mut opts = ReadOptions::default();
+        opts.set_total_order_seek(true);
+        DBIterator::new_cf(self, cf_handle, &opts, mode)
+    }
+
     pub fn prefix_iterator_cf<'a>(
         &self,
         cf_handle: ColumnFamily,

src/db_options.rs

@@ -19,8 +19,8 @@ use std::path::Path;
 use libc::{self, c_int, c_uchar, c_uint, c_void, size_t, uint64_t};
 
 use ffi;
-use {BlockBasedOptions, DBCompactionStyle, DBCompressionType, DBRecoveryMode,
-	 Options, WriteOptions};
+use {BlockBasedOptions, DBCompactionStyle, DBCompressionType, DBRecoveryMode, MemtableFactory,
+     Options, WriteOptions};
 use compaction_filter::{self, CompactionFilterCallback, CompactionFilterFn, filter_callback};
 use comparator::{self, ComparatorCallback, CompareFn};
 use merge_operator::{self, MergeFn, MergeOperatorCallback, full_merge_callback,
@@ -904,6 +904,47 @@ impl Options {
         unsafe { ffi::rocksdb_options_set_disable_auto_compactions(self.inner, disable as c_int) }
     }
 
+    /// Defines the underlying memtable implementation.
+    /// See https://github.com/facebook/rocksdb/wiki/MemTable for more information.
+    /// Defaults to using a skiplist.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// use rocksdb::{Options, MemtableFactory};
+    /// let mut opts = Options::default();
+    /// let factory = MemtableFactory::HashSkipList {
+    ///     bucket_count: 1_000_000,
+    ///     height: 4,
+    ///     branching_factor: 4,
+    /// };
+    ///
+    /// opts.set_allow_concurrent_memtable_write(false);
+    /// opts.set_memtable_factory(factory);
+    /// ```
+    pub fn set_memtable_factory(&mut self, factory: MemtableFactory) {
+        match factory {
+            MemtableFactory::Vector => unsafe {
+                ffi::rocksdb_options_set_memtable_vector_rep(self.inner);
+            },
+            MemtableFactory::HashSkipList {
+                bucket_count,
+                height,
+                branching_factor,
+            } => unsafe {
+                ffi::rocksdb_options_set_hash_skip_list_rep(
+                    self.inner,
+                    bucket_count,
+                    height,
+                    branching_factor,
+                );
+            },
+            MemtableFactory::HashLinkList { bucket_count } => unsafe {
+                ffi::rocksdb_options_set_hash_link_list_rep(self.inner, bucket_count);
+            },
+        };
+    }
+
     pub fn set_block_based_table_factory(&mut self, factory: &BlockBasedOptions) {
         unsafe {
             ffi::rocksdb_options_set_block_based_table_factory(self.inner, factory.inner);
@@ -1065,6 +1106,7 @@ impl Default for WriteOptions {
 
 #[cfg(test)]
 mod tests {
+    use MemtableFactory;
     use Options;
 
     #[test]
@@ -1077,4 +1119,16 @@ mod tests {
         let opts = Options::default();
         assert!(opts.get_statistics().is_none());
     }
+
+    #[test]
+    fn test_set_memtable_factory() {
+        let mut opts = Options::default();
+        opts.set_memtable_factory(MemtableFactory::Vector);
+        opts.set_memtable_factory(MemtableFactory::HashLinkList { bucket_count: 100 });
+        opts.set_memtable_factory(MemtableFactory::HashSkipList {
+            bucket_count: 100,
+            height: 4,
+            branching_factor: 4,
+        });
+    }
 }

src/lib.rs

@@ -136,6 +136,14 @@ pub struct BlockBasedOptions {
     inner: *mut ffi::rocksdb_block_based_table_options_t,
 }
 
+/// Defines the underlying memtable implementation.
+/// See https://github.com/facebook/rocksdb/wiki/MemTable for more information.
+pub enum MemtableFactory {
+    Vector,
+    HashSkipList { bucket_count: usize, height: i32, branching_factor: i32 },
+    HashLinkList { bucket_count: usize }
+}
+
 /// Database-wide options around performance and behavior.
 ///
 /// Please read [the official tuning guide](https://github.com/facebook/rocksdb/wiki/RocksDB-Tuning-Guide), and most importantly, measure performance under realistic workloads with realistic hardware.

tests/test_iterator.rs

@@ -14,7 +14,7 @@
 //
 extern crate rocksdb;
 
-use rocksdb::{DB, Direction, IteratorMode, Options};
+use rocksdb::{DB, Direction, IteratorMode, MemtableFactory, Options};
 
 fn cba(input: &Box<[u8]>) -> Box<[u8]> {
     input.iter().cloned().collect::<Vec<_>>().into_boxed_slice()
@@ -171,7 +171,7 @@ pub fn test_prefix_iterator() {
         let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(3);
 
         let mut opts = Options::default();
-		opts.create_if_missing(true);
+        opts.create_if_missing(true);
         opts.set_prefix_extractor(prefix_extractor);
 
         let db = DB::open(&opts, path).unwrap();
@@ -193,4 +193,54 @@ pub fn test_prefix_iterator() {
             assert_eq!(b_iterator.collect::<Vec<_>>(), expected)
         }
     }
+    let opts = Options::default();
+    assert!(DB::destroy(&opts, path).is_ok());
+}
+
+#[test]
+pub fn test_full_iterator() {
+    let path = "_rust_rocksdb_fulliteratortest";
+    {
+        let a1: Box<[u8]> = key(b"aaa1");
+        let a2: Box<[u8]> = key(b"aaa2");
+        let b1: Box<[u8]> = key(b"bbb1");
+        let b2: Box<[u8]> = key(b"bbb2");
+
+        let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(3);
+        let factory = MemtableFactory::HashSkipList {
+            bucket_count: 1_000_000,
+            height: 4,
+            branching_factor: 4,
+        };
+
+        let mut opts = Options::default();
+        opts.create_if_missing(true);
+        opts.set_prefix_extractor(prefix_extractor);
+        opts.set_allow_concurrent_memtable_write(false);
+        opts.set_memtable_factory(factory);
+
+        let db = DB::open(&opts, path).unwrap();
+
+        assert!(db.put(&*a1, &*a1).is_ok());
+        assert!(db.put(&*a2, &*a2).is_ok());
+        assert!(db.put(&*b1, &*b1).is_ok());
+        assert!(db.put(&*b2, &*b2).is_ok());
+
+        // A normal iterator won't work here since we're using a HashSkipList for our memtable
+        // implementation (which buckets keys based on their prefix):
+        let bad_iterator = db.iterator(IteratorMode::Start);
+        assert_eq!(bad_iterator.collect::<Vec<_>>(), vec![]);
+
+        let expected = vec![
+            (cba(&a1), cba(&a1)),
+            (cba(&a2), cba(&a2)),
+            (cba(&b1), cba(&b1)),
+            (cba(&b2), cba(&b2)),
+        ];
+
+        let a_iterator = db.full_iterator(IteratorMode::Start);
+        assert_eq!(a_iterator.collect::<Vec<_>>(), expected)
+    }
+    let opts = Options::default();
+    assert!(DB::destroy(&opts, path).is_ok());
 }