rocksdb/
merge_operator.rs

1// Copyright 2020 Tyler Neely
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7// http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14//
15
16//! rustic merge operator
17//!
18//! ```
19//! use rocksdb::{Options, DB, MergeOperands};
20//!
21//! fn concat_merge(new_key: &[u8],
22//!                 existing_val: Option<&[u8]>,
23//!                 operands: &MergeOperands)
24//!                 -> Option<Vec<u8>> {
25//!
26//!    let mut result: Vec<u8> = Vec::with_capacity(operands.len());
27//!    existing_val.map(|v| {
28//!        for e in v {
29//!            result.push(*e)
30//!        }
31//!    });
32//!    for op in operands {
33//!        for e in op {
34//!            result.push(*e)
35//!        }
36//!    }
37//!    Some(result)
38//! }
39//!
40//!let tempdir = tempfile::Builder::new()
41//!    .prefix("_rust_path_to_rocksdb")
42//!    .tempdir()
43//!    .expect("Failed to create temporary path for the _rust_path_to_rocksdb");
44//!let path = tempdir.path();
45//!let mut opts = Options::default();
46//!
47//!opts.create_if_missing(true);
48//!opts.set_merge_operator_associative("test operator", concat_merge);
49//!{
50//!    let db = DB::open(&opts, path).unwrap();
51//!    let p = db.put(b"k1", b"a");
52//!    db.merge(b"k1", b"b");
53//!    db.merge(b"k1", b"c");
54//!    db.merge(b"k1", b"d");
55//!    db.merge(b"k1", b"efg");
56//!    let r = db.get(b"k1");
57//!    assert_eq!(r.unwrap().unwrap(), b"abcdefg");
58//!}
59//!let _ = DB::destroy(&opts, path);
60//! ```
61
62use libc::{self, c_char, c_int, c_void, size_t};
63use std::ffi::CString;
64use std::mem;
65use std::ptr;
66use std::slice;
67
68pub trait MergeFn:
69    Fn(&[u8], Option<&[u8]>, &MergeOperands) -> Option<Vec<u8>> + Send + Sync + 'static
70{
71}
72impl<F> MergeFn for F where
73    F: Fn(&[u8], Option<&[u8]>, &MergeOperands) -> Option<Vec<u8>> + Send + Sync + 'static
74{
75}
76
77pub struct MergeOperatorCallback<F: MergeFn, PF: MergeFn> {
78    pub name: CString,
79    pub full_merge_fn: F,
80    pub partial_merge_fn: PF,
81}
82
83pub unsafe extern "C" fn destructor_callback<F: MergeFn, PF: MergeFn>(raw_cb: *mut c_void) {
84    drop(unsafe { Box::from_raw(raw_cb as *mut MergeOperatorCallback<F, PF>) });
85}
86
87pub unsafe extern "C" fn delete_callback(
88    _raw_cb: *mut c_void,
89    value: *const c_char,
90    value_length: size_t,
91) {
92    if !value.is_null() {
93        drop(unsafe {
94            Box::from_raw(slice::from_raw_parts_mut(
95                value.cast_mut().cast::<u8>(),
96                value_length,
97            ))
98        });
99    }
100}
101
102pub unsafe extern "C" fn name_callback<F: MergeFn, PF: MergeFn>(
103    raw_cb: *mut c_void,
104) -> *const c_char {
105    let cb = unsafe { &mut *(raw_cb as *mut MergeOperatorCallback<F, PF>) };
106    cb.name.as_ptr()
107}
108
109pub unsafe extern "C" fn full_merge_callback<F: MergeFn, PF: MergeFn>(
110    raw_cb: *mut c_void,
111    raw_key: *const c_char,
112    key_len: size_t,
113    existing_value: *const c_char,
114    existing_value_len: size_t,
115    operands_list: *const *const c_char,
116    operands_list_len: *const size_t,
117    num_operands: c_int,
118    success: *mut u8,
119    new_value_length: *mut size_t,
120) -> *mut c_char {
121    let cb = unsafe { &mut *(raw_cb as *mut MergeOperatorCallback<F, PF>) };
122    let operands = &MergeOperands::new(operands_list, operands_list_len, num_operands);
123    let key = unsafe { slice::from_raw_parts(raw_key.cast::<u8>(), key_len) };
124    let oldval = if existing_value.is_null() {
125        None
126    } else {
127        Some(unsafe { slice::from_raw_parts(existing_value.cast::<u8>(), existing_value_len) })
128    };
129    (cb.full_merge_fn)(key, oldval, operands).map_or_else(
130        || {
131            unsafe { *new_value_length = 0 };
132            unsafe { *success = 0_u8 };
133            ptr::null_mut() as *mut c_char
134        },
135        |result| {
136            unsafe { *new_value_length = result.len() as size_t };
137            unsafe { *success = 1_u8 };
138            Box::into_raw(result.into_boxed_slice()) as *mut c_char
139        },
140    )
141}
142
143pub unsafe extern "C" fn partial_merge_callback<F: MergeFn, PF: MergeFn>(
144    raw_cb: *mut c_void,
145    raw_key: *const c_char,
146    key_len: size_t,
147    operands_list: *const *const c_char,
148    operands_list_len: *const size_t,
149    num_operands: c_int,
150    success: *mut u8,
151    new_value_length: *mut size_t,
152) -> *mut c_char {
153    let cb = unsafe { &mut *(raw_cb as *mut MergeOperatorCallback<F, PF>) };
154    let operands = &MergeOperands::new(operands_list, operands_list_len, num_operands);
155    let key = unsafe { slice::from_raw_parts(raw_key.cast::<u8>(), key_len) };
156    (cb.partial_merge_fn)(key, None, operands).map_or_else(
157        || {
158            unsafe { *new_value_length = 0 };
159            unsafe { *success = 0_u8 };
160            ptr::null_mut::<c_char>()
161        },
162        |result| {
163            unsafe { *new_value_length = result.len() as size_t };
164            unsafe { *success = 1_u8 };
165            Box::into_raw(result.into_boxed_slice()) as *mut c_char
166        },
167    )
168}
169
170pub struct MergeOperands {
171    operands_list: *const *const c_char,
172    operands_list_len: *const size_t,
173    num_operands: usize,
174}
175
176impl MergeOperands {
177    fn new(
178        operands_list: *const *const c_char,
179        operands_list_len: *const size_t,
180        num_operands: c_int,
181    ) -> MergeOperands {
182        assert!(num_operands >= 0);
183        MergeOperands {
184            operands_list,
185            operands_list_len,
186            num_operands: num_operands as usize,
187        }
188    }
189
190    pub fn len(&self) -> usize {
191        self.num_operands
192    }
193
194    pub fn is_empty(&self) -> bool {
195        self.num_operands == 0
196    }
197
198    pub fn iter(&self) -> MergeOperandsIter {
199        MergeOperandsIter {
200            operands: self,
201            cursor: 0,
202        }
203    }
204
205    fn get_operand(&self, index: usize) -> Option<&[u8]> {
206        if index >= self.num_operands {
207            None
208        } else {
209            unsafe {
210                let base = self.operands_list as usize;
211                let base_len = self.operands_list_len as usize;
212                let spacing = mem::size_of::<*const *const u8>();
213                let spacing_len = mem::size_of::<*const size_t>();
214                let len_ptr = (base_len + (spacing_len * index)) as *const size_t;
215                let len = *len_ptr;
216                let ptr = base + (spacing * index);
217                Some(slice::from_raw_parts(*(ptr as *const *const u8), len))
218            }
219        }
220    }
221}
222
223pub struct MergeOperandsIter<'a> {
224    operands: &'a MergeOperands,
225    cursor: usize,
226}
227
228impl<'a> Iterator for MergeOperandsIter<'a> {
229    type Item = &'a [u8];
230
231    fn next(&mut self) -> Option<Self::Item> {
232        let operand = self.operands.get_operand(self.cursor)?;
233        self.cursor += 1;
234        Some(operand)
235    }
236
237    fn size_hint(&self) -> (usize, Option<usize>) {
238        let remaining = self.operands.num_operands - self.cursor;
239        (remaining, Some(remaining))
240    }
241}
242
243impl<'a> IntoIterator for &'a MergeOperands {
244    type Item = &'a [u8];
245    type IntoIter = MergeOperandsIter<'a>;
246
247    fn into_iter(self) -> Self::IntoIter {
248        Self::IntoIter {
249            operands: self,
250            cursor: 0,
251        }
252    }
253}
rocksdb/merge_operator.rs

rocksdb/
merge_operator.rs
