support slicetransform (#63)

librocksdb_sys/src/lib.rs

@@ -17,7 +17,7 @@ extern crate libc;
 #[cfg(test)]
 extern crate tempdir;
 
-use libc::{c_char, c_uchar, c_int, c_void, size_t, uint64_t};
+use libc::{c_char, c_uchar, c_int, c_void, size_t, uint64_t, c_double};
 use std::ffi::CStr;
 
 pub enum DBOptions {}
@@ -37,6 +37,7 @@ pub enum DBFlushOptions {}
 pub enum DBCompactionFilter {}
 pub enum DBBackupEngine {}
 pub enum DBRestoreOptions {}
+pub enum DBSliceTransform {}
 
 pub fn new_bloom_filter(bits: c_int) -> *mut DBFilterPolicy {
     unsafe { crocksdb_filterpolicy_create_bloom(bits) }
@@ -190,6 +191,10 @@ extern "C" {
     pub fn crocksdb_options_set_stats_dump_period_sec(options: *mut DBOptions, v: usize);
     pub fn crocksdb_options_set_num_levels(options: *mut DBOptions, v: c_int);
     pub fn crocksdb_options_set_db_log_dir(options: *mut DBOptions, path: *const c_char);
+    pub fn crocksdb_options_set_prefix_extractor(options: *mut DBOptions,
+                                                 prefix_extractor: *mut DBSliceTransform);
+    pub fn crocksdb_options_set_memtable_prefix_bloom_size_ratio(options: *mut DBOptions,
+                                                                 ratio: c_double);
     pub fn crocksdb_filterpolicy_create_bloom_full(bits_per_key: c_int) -> *mut DBFilterPolicy;
     pub fn crocksdb_filterpolicy_create_bloom(bits_per_key: c_int) -> *mut DBFilterPolicy;
     pub fn crocksdb_open(options: *mut DBOptions,
@@ -505,6 +510,23 @@ extern "C" {
                                                                 wal_path: *const c_char,
                                                                 ropts: *const DBRestoreOptions,
                                                                 err: *mut *mut c_char);
+    // SliceTransform
+    pub fn crocksdb_slicetransform_create(state: *mut c_void,
+                                          destructor: extern "C" fn(*mut c_void),
+                                          transform: extern "C" fn(*mut c_void,
+                                                                   *const u8,
+                                                                   size_t,
+                                                                   *mut size_t)
+                                                                   -> *const u8,
+                                          in_domain: extern "C" fn(*mut c_void,
+                                                                   *const u8,
+                                                                   size_t)
+                                                                   -> u8,
+                                          in_range: extern "C" fn(*mut c_void, *const u8, size_t)
+                                                                  -> u8,
+                                          name: extern "C" fn(*mut c_void) -> *const c_char)
+                                          -> *mut DBSliceTransform;
+    pub fn crocksdb_slicetransform_destroy(transform: *mut DBSliceTransform);
 }
 
 #[cfg(test)]

src/lib.rs

@@ -25,6 +25,7 @@ pub mod rocksdb_options;
 pub mod merge_operator;
 pub mod comparator;
 mod compaction_filter;
+mod slice_transform;
 
 pub use compaction_filter::CompactionFilter;
 pub use librocksdb_sys::{DBCompactionStyle, DBCompressionType, DBRecoveryMode, new_bloom_filter,
@@ -33,3 +34,4 @@ pub use merge_operator::MergeOperands;
 pub use rocksdb::{DB, DBIterator, DBVector, Kv, SeekKey, Writable, WriteBatch, CFHandle, Range,
                   BackupEngine};
 pub use rocksdb_options::{BlockBasedOptions, Options, ReadOptions, WriteOptions, RestoreOptions};
+pub use slice_transform::SliceTransform;

src/rocksdb_options.rs

@@ -22,6 +22,7 @@ use crocksdb_ffi::{self, DBOptions, DBWriteOptions, DBBlockBasedTableOptions, DB
 use libc::{self, c_int, size_t, c_void};
 use merge_operator::{self, MergeOperatorCallback, full_merge_callback, partial_merge_callback};
 use merge_operator::MergeFn;
+use slice_transform::{SliceTransform, new_slice_transform};
 use std::ffi::{CStr, CString};
 use std::mem;
 
@@ -86,6 +87,12 @@ impl BlockBasedOptions {
                                                                                        v as u8);
         }
     }
+
+    pub fn set_whole_key_filtering(&mut self, v: bool) {
+        unsafe {
+            crocksdb_ffi::crocksdb_block_based_options_set_whole_key_filtering(self.inner, v);
+        }
+    }
 }
 
 /// The UnsafeSnap must be destroyed by db, it maybe be leaked
@@ -564,6 +571,29 @@ impl Options {
             crocksdb_ffi::crocksdb_options_set_keep_log_file_num(self.inner, num as size_t);
         }
     }
+
+    pub fn set_prefix_extractor<S>(&mut self,
+                                   name: S,
+                                   transform: Box<SliceTransform>)
+                                   -> Result<(), String>
+        where S: Into<Vec<u8>>
+    {
+        unsafe {
+            let c_name = match CString::new(name) {
+                Ok(s) => s,
+                Err(e) => return Err(format!("failed to convert to cstring: {:?}", e)),
+            };
+            let transform = try!(new_slice_transform(c_name, transform));
+            crocksdb_ffi::crocksdb_options_set_prefix_extractor(self.inner, transform);
+            Ok(())
+        }
+    }
+
+    pub fn set_memtable_prefix_bloom_size_ratio(&mut self, ratio: f64) {
+        unsafe {
+            crocksdb_ffi::crocksdb_options_set_memtable_prefix_bloom_size_ratio(self.inner, ratio);
+        }
+    }
 }
 
 pub struct FlushOptions {

src/slice_transform.rs

+// Copyright 2016 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+
+use crocksdb_ffi::{self, DBSliceTransform};
+use libc::{c_void, c_char, size_t};
+use std::ffi::CString;
+use std::slice;
+
+// `SliceTranform` is a generic pluggable way of transforming one string
+// mainly used for prefix blooms.
+pub trait SliceTransform {
+    // Extract a prefix from a specified key
+    fn transform<'a>(&mut self, key: &'a [u8]) -> &'a [u8];
+
+    // Determine whether the specified key is compatible with the logic
+    // specified in the Transform method. This method is invoked for every
+    // key that is inserted into the db. If this method returns true,
+    // then Transform is called to translate the key to its prefix and
+    // that returned prefix is inserted into the bloom filter. If this
+    // method returns false, then the call to Transform is skipped and
+    // no prefix is inserted into the bloom filters.
+    fn in_domain(&mut self, key: &[u8]) -> bool;
+
+    // This is currently not used and remains here for backward compatibility.
+    fn in_range(&mut self, key: &[u8]) -> bool {
+        true
+    }
+}
+
+#[repr(C)]
+pub struct SliceTransformProxy {
+    name: CString,
+    transform: Box<SliceTransform>,
+}
+
+extern "C" fn name(transform: *mut c_void) -> *const c_char {
+    unsafe { (*(transform as *mut SliceTransformProxy)).name.as_ptr() }
+}
+
+extern "C" fn destructor(transform: *mut c_void) {
+    unsafe {
+        Box::from_raw(transform as *mut SliceTransformProxy);
+    }
+}
+
+extern "C" fn transform(transform: *mut c_void,
+                        key: *const u8,
+                        key_len: size_t,
+                        dest_len: *mut size_t)
+                        -> *const u8 {
+    unsafe {
+        let transform = &mut *(transform as *mut SliceTransformProxy);
+        let key = slice::from_raw_parts(key, key_len);
+        let prefix = transform.transform.transform(key);
+        *dest_len = prefix.len() as size_t;
+        prefix.as_ptr() as *const u8
+    }
+}
+
+extern "C" fn in_domain(transform: *mut c_void, key: *const u8, key_len: size_t) -> u8 {
+    unsafe {
+        let transform = &mut *(transform as *mut SliceTransformProxy);
+        let key = slice::from_raw_parts(key, key_len);
+        transform.transform.in_domain(key) as u8
+    }
+}
+
+extern "C" fn in_range(transform: *mut c_void, key: *const u8, key_len: size_t) -> u8 {
+    unsafe {
+        let transform = &mut *(transform as *mut SliceTransformProxy);
+        let key = slice::from_raw_parts(key, key_len);
+        transform.transform.in_range(key) as u8
+    }
+}
+
+
+pub unsafe fn new_slice_transform(c_name: CString,
+                                  f: Box<SliceTransform>)
+                                  -> Result<*mut DBSliceTransform, String> {
+    let proxy = Box::into_raw(Box::new(SliceTransformProxy {
+        name: c_name,
+        transform: f,
+    }));
+    let transform = crocksdb_ffi::crocksdb_slicetransform_create(proxy as *mut c_void,
+                                                                 destructor,
+                                                                 transform,
+                                                                 in_domain,
+                                                                 in_range,
+                                                                 name);
+    Ok(transform)
+}

test/test.rs

@@ -7,3 +7,4 @@ mod test_column_family;
 mod test_compaction_filter;
 mod test_compact_range;
 mod test_rocksdb_options;
+mod test_slice_transform;

test/test_slice_transform.rs

+// Copyright 2016 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+
+
+use rocksdb::{Writable, DB, SliceTransform, Options, SeekKey, BlockBasedOptions};
+use tempdir::TempDir;
+
+struct FixedPostfixTransform {
+    pub postfix_len: usize,
+}
+
+impl SliceTransform for FixedPostfixTransform {
+    fn transform<'a>(&mut self, key: &'a [u8]) -> &'a [u8] {
+        let mid = key.len() - self.postfix_len;
+        &key[..mid]
+    }
+
+    fn in_domain(&mut self, key: &[u8]) -> bool {
+        key.len() >= self.postfix_len
+    }
+}
+
+#[test]
+fn test_slice_transform() {
+    let path = TempDir::new("_rust_rocksdb_slice_transform_test").expect("");
+    let mut opts = Options::new();
+
+    let mut block_opts = BlockBasedOptions::new();
+    block_opts.set_bloom_filter(10, false);
+    block_opts.set_whole_key_filtering(false);
+
+    opts.set_block_based_table_factory(&block_opts);
+    opts.set_memtable_prefix_bloom_size_ratio(0.25);
+
+    opts.set_prefix_extractor("test", Box::new(FixedPostfixTransform { postfix_len: 2 }))
+        .unwrap();
+    opts.create_if_missing(true);
+
+    let db = DB::open(opts, path.path().to_str().unwrap()).unwrap();
+    let samples = vec![(b"key_01".to_vec(), b"1".to_vec()),
+                       (b"key_02".to_vec(), b"2".to_vec()),
+                       (b"key_0303".to_vec(), b"3".to_vec()),
+                       (b"key_0404".to_vec(), b"4".to_vec())];
+
+    for &(ref k, ref v) in &samples {
+        db.put(k, v).unwrap();
+        assert_eq!(v.as_slice(), &*db.get(k).unwrap().unwrap());
+    }
+
+    let mut it = db.iter();
+
+    let invalid_seeks =
+        vec![b"key_".to_vec(), b"key_0".to_vec(), b"key_030".to_vec(), b"key_03000".to_vec()];
+
+    for key in invalid_seeks {
+        it.seek(SeekKey::Key(&key));
+        assert!(!it.valid());
+    }
+
+    let valid_seeks = vec![(b"key_00".to_vec(), b"key_01".to_vec()),
+                           (b"key_03".to_vec(), b"key_0303".to_vec()),
+                           (b"key_0301".to_vec(), b"key_0303".to_vec())];
+
+    for (key, expect_key) in valid_seeks {
+        it.seek(SeekKey::Key(&key));
+        assert!(it.valid());
+        assert_eq!(it.key(), &*expect_key);
+    }
+
+    // TODO: support total_order mode and add test later.
+}