Merge branch 'master' of https://github.com/spacejam/rust-rocksdb
Jordan Terrell
7 years ago
commit
f06812a00d
@ -1 +1 @@ |
||||
Subproject commit acf935e40f9d6f4c3d13c7d310def7064c1f1c95 |
||||
Subproject commit dbd8fa09b823826dd2a30bc119dad7a6fa9a4c6d |
||||
File diff suppressed because one or more lines are too long
@ -0,0 +1,149 @@ |
||||
// Copyright 2018 Tyler Neely
|
||||
//
|
||||
// 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,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
use std::ffi::CString; |
||||
use std::mem; |
||||
use std::ptr; |
||||
use std::slice; |
||||
|
||||
use libc::{self, c_char, c_void, size_t}; |
||||
|
||||
use ffi; |
||||
|
||||
/// A SliceTranform is a generic pluggable way of transforming one string
|
||||
/// to another. Its primary use-case is in configuring rocksdb
|
||||
/// to store prefix blooms by setting prefix_extractor in
|
||||
/// ColumnFamilyOptions.
|
||||
pub struct SliceTransform { |
||||
pub inner: *mut ffi::rocksdb_slicetransform_t, |
||||
} |
||||
|
||||
// NB we intentionally don't implement a Drop that passes
|
||||
// through to rocksdb_slicetransform_destroy because
|
||||
// this is currently only used (to my knowledge)
|
||||
// by people passing it as a prefix extractor when
|
||||
// opening a DB.
|
||||
|
||||
impl SliceTransform { |
||||
pub fn create( |
||||
name: &str, |
||||
transform_fn: TransformFn, |
||||
in_domain_fn: Option<InDomainFn>, |
||||
) -> SliceTransform{ |
||||
let cb = Box::new(TransformCallback { |
||||
name: CString::new(name.as_bytes()).unwrap(), |
||||
transform_fn: transform_fn, |
||||
in_domain_fn: in_domain_fn, |
||||
}); |
||||
|
||||
let st = unsafe { |
||||
ffi::rocksdb_slicetransform_create( |
||||
mem::transmute(cb), |
||||
Some(slice_transform_destructor_callback), |
||||
Some(transform_callback), |
||||
|
||||
// this is ugly, but I can't get the compiler
|
||||
// not to barf with "expected fn pointer, found fn item"
|
||||
// without this. sorry.
|
||||
if let Some(_) = in_domain_fn { |
||||
Some(in_domain_callback) |
||||
} else { |
||||
None |
||||
}, |
||||
|
||||
// this None points to the deprecated InRange callback
|
||||
None, |
||||
Some(slice_transform_name_callback), |
||||
) |
||||
}; |
||||
|
||||
SliceTransform { |
||||
inner: st |
||||
} |
||||
} |
||||
|
||||
pub fn create_fixed_prefix(len: size_t) -> SliceTransform { |
||||
SliceTransform { |
||||
inner: unsafe { |
||||
ffi::rocksdb_slicetransform_create_fixed_prefix(len) |
||||
}, |
||||
} |
||||
} |
||||
|
||||
pub fn create_noop() -> SliceTransform { |
||||
SliceTransform { |
||||
inner: unsafe { |
||||
ffi::rocksdb_slicetransform_create_noop() |
||||
}, |
||||
} |
||||
} |
||||
} |
||||
|
||||
pub type TransformFn = fn(&[u8]) -> Vec<u8>; |
||||
pub type InDomainFn = fn(&[u8]) -> bool; |
||||
|
||||
pub struct TransformCallback { |
||||
pub name: CString, |
||||
pub transform_fn: TransformFn, |
||||
pub in_domain_fn: Option<InDomainFn>, |
||||
} |
||||
|
||||
pub unsafe extern "C" fn slice_transform_destructor_callback( |
||||
raw_cb: *mut c_void |
||||
) { |
||||
let transform: Box<TransformCallback> = mem::transmute(raw_cb); |
||||
drop(transform); |
||||
} |
||||
|
||||
pub unsafe extern "C" fn slice_transform_name_callback( |
||||
raw_cb: *mut c_void |
||||
) -> *const c_char { |
||||
let cb = &mut *(raw_cb as *mut TransformCallback); |
||||
cb.name.as_ptr() |
||||
} |
||||
|
||||
pub unsafe extern "C" fn transform_callback( |
||||
raw_cb: *mut c_void, |
||||
raw_key: *const c_char, |
||||
key_len: size_t, |
||||
dst_length: *mut size_t, |
||||
) -> *mut c_char { |
||||
let cb = &mut *(raw_cb as *mut TransformCallback); |
||||
let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize); |
||||
let mut result = (cb.transform_fn)(key); |
||||
result.shrink_to_fit(); |
||||
|
||||
// copy the result into a C++ destroyable buffer
|
||||
let buf = libc::malloc(result.len() as size_t); |
||||
assert!(!buf.is_null()); |
||||
ptr::copy(result.as_ptr() as *mut c_void, &mut *buf, result.len()); |
||||
|
||||
*dst_length = result.len() as size_t; |
||||
buf as *mut c_char |
||||
} |
||||
|
||||
pub unsafe extern "C" fn in_domain_callback( |
||||
raw_cb: *mut c_void, |
||||
raw_key: *const c_char, |
||||
key_len: size_t, |
||||
) -> u8 { |
||||
let cb = &mut *(raw_cb as *mut TransformCallback); |
||||
let key = slice::from_raw_parts(raw_key as *const u8, key_len as usize); |
||||
|
||||
if (cb.in_domain_fn.unwrap())(key) { |
||||
1 |
||||
} else { |
||||
0 |
||||
} |
||||
} |
||||
@ -0,0 +1,48 @@ |
||||
extern crate rocksdb; |
||||
|
||||
use rocksdb::{DB, Options, SliceTransform}; |
||||
|
||||
#[test] |
||||
pub fn test_slice_transform() { |
||||
|
||||
let path = "_rust_rocksdb_slicetransform_test"; |
||||
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"); |
||||
|
||||
fn first_three(k: &[u8]) -> Vec<u8> { |
||||
k.iter().take(3).cloned().collect() |
||||
} |
||||
|
||||
let prefix_extractor = SliceTransform::create("first_three", first_three, None); |
||||
|
||||
let mut opts = Options::default(); |
||||
opts.create_if_missing(true); |
||||
opts.set_prefix_extractor(prefix_extractor); |
||||
|
||||
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()); |
||||
|
||||
fn cba(input: &Box<[u8]>) -> Box<[u8]> { |
||||
input.iter().cloned().collect::<Vec<_>>().into_boxed_slice() |
||||
} |
||||
|
||||
fn key(k: &[u8]) -> Box<[u8]> { k.to_vec().into_boxed_slice() } |
||||
|
||||
{ |
||||
let expected = vec![(cba(&a1), cba(&a1)), (cba(&a2), cba(&a2))]; |
||||
let a_iterator = db.prefix_iterator(b"aaa"); |
||||
assert_eq!(a_iterator.collect::<Vec<_>>(), expected) |
||||
} |
||||
|
||||
{ |
||||
let expected = vec![(cba(&b1), cba(&b1)), (cba(&b2), cba(&b2))]; |
||||
let b_iterator = db.prefix_iterator(b"bbb"); |
||||
assert_eq!(b_iterator.collect::<Vec<_>>(), expected) |
||||
} |
||||
} |
||||
Loading…
Reference in new issue