From 8acabab60f0b9edfef989278d7a2fc64d1b8f921 Mon Sep 17 00:00:00 2001 From: Martin Ek Date: Thu, 5 Apr 2018 16:51:06 +0200 Subject: [PATCH 1/3] Add memtable factory customization --- src/db_options.rs | 58 +++++++++++++++++++++++++++++++++++++++++++++-- src/lib.rs | 8 +++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/src/db_options.rs b/src/db_options.rs index 716b6f7..fcfe408 100644 --- a/src/db_options.rs +++ b/src/db_options.rs @@ -18,8 +18,8 @@ use std::mem; 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, @@ -883,6 +883,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); @@ -1024,6 +1065,7 @@ impl Default for WriteOptions { #[cfg(test)] mod tests { + use MemtableFactory; use Options; #[test] @@ -1036,4 +1078,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, + }); + } } diff --git a/src/lib.rs b/src/lib.rs index f982702..d7133d0 100644 --- a/src/lib.rs +++ b/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. From 37470d341cad23b9b56c02be0ad677a69f50ca47 Mon Sep 17 00:00:00 2001 From: Martin Ek Date: Fri, 6 Apr 2018 14:56:36 +0200 Subject: [PATCH 2/3] Add db.full_iterator() --- src/db.rs | 19 ++++++++++++++++ tests/test_iterator.rs | 50 +++++++++++++++++++++++++++++++++++++++++- 2 files changed, 68 insertions(+), 1 deletion(-) diff --git a/src/db.rs b/src/db.rs index a26fae9..78fd8c7 100644 --- a/src/db.rs +++ b/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 { + 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, diff --git a/tests/test_iterator.rs b/tests/test_iterator.rs index 3cb368f..5f8b1bb 100644 --- a/tests/test_iterator.rs +++ b/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::>().into_boxed_slice() @@ -194,3 +194,51 @@ pub fn test_prefix_iterator() { } } } + +#[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![]); + + 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::>(), expected) + } + let opts = Options::default(); + assert!(DB::destroy(&opts, path).is_ok()); +} From dbcca72ce6d495341496edbc04e69e96bc35215f Mon Sep 17 00:00:00 2001 From: Martin Ek Date: Fri, 6 Apr 2018 14:59:16 +0200 Subject: [PATCH 3/3] Clean-up after test_prefix_iterator --- tests/test_iterator.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_iterator.rs b/tests/test_iterator.rs index 5f8b1bb..ef7a7bd 100644 --- a/tests/test_iterator.rs +++ b/tests/test_iterator.rs @@ -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,6 +193,8 @@ pub fn test_prefix_iterator() { assert_eq!(b_iterator.collect::>(), expected) } } + let opts = Options::default(); + assert!(DB::destroy(&opts, path).is_ok()); } #[test]