rust_rocksdb/
comparator.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
16use crate::{ffi, ffi_util::CStrLike};
17use libc::{c_char, c_int, c_uchar, c_void, size_t};
18use std::cmp::Ordering;
19use std::ffi::CString;
20use std::ptr::NonNull;
21use std::slice;
22
23// RocksDB calls these from whatever thread reaches the comparison, including
24// background flush and compaction threads, and from several at once. They have
25// to be shareable across threads for `Options` and `Comparator` to be `Sync`.
26pub type CompareFn = dyn Fn(&[u8], &[u8]) -> Ordering + Send + Sync;
27
28pub type CompareTsFn = dyn Fn(&[u8], &[u8]) -> Ordering + Send + Sync;
29
30pub type CompareWithoutTsFn = dyn Fn(&[u8], bool, &[u8], bool) -> Ordering + Send + Sync;
31
32pub struct ComparatorCallback {
33    pub name: CString,
34    pub compare_fn: Box<CompareFn>,
35}
36
37impl ComparatorCallback {
38    pub unsafe extern "C" fn destructor_callback(raw_cb: *mut c_void) {
39        unsafe {
40            drop(Box::from_raw(raw_cb as *mut Self));
41        }
42    }
43
44    pub unsafe extern "C" fn name_callback(raw_cb: *mut c_void) -> *const c_char {
45        unsafe {
46            let cb: &Self = &*(raw_cb as *const Self);
47            let ptr = cb.name.as_ptr();
48            ptr as *const c_char
49        }
50    }
51
52    pub unsafe extern "C" fn compare_callback(
53        raw_cb: *mut c_void,
54        a_raw: *const c_char,
55        a_len: size_t,
56        b_raw: *const c_char,
57        b_len: size_t,
58    ) -> c_int {
59        unsafe {
60            let cb: &Self = &*(raw_cb as *const Self);
61            let a: &[u8] = slice::from_raw_parts(a_raw.cast::<u8>(), a_len);
62            let b: &[u8] = slice::from_raw_parts(b_raw.cast::<u8>(), b_len);
63            (cb.compare_fn)(a, b) as c_int
64        }
65    }
66}
67
68pub struct ComparatorWithTsCallback {
69    pub name: CString,
70    pub compare_fn: Box<CompareFn>,
71    pub compare_ts_fn: Box<CompareTsFn>,
72    pub compare_without_ts_fn: Box<CompareWithoutTsFn>,
73}
74
75impl ComparatorWithTsCallback {
76    pub unsafe extern "C" fn destructor_callback(raw_cb: *mut c_void) {
77        unsafe {
78            drop(Box::from_raw(raw_cb as *mut Self));
79        }
80    }
81
82    pub unsafe extern "C" fn name_callback(raw_cb: *mut c_void) -> *const c_char {
83        unsafe {
84            let cb: &Self = &*(raw_cb as *const Self);
85            let ptr = cb.name.as_ptr();
86            ptr as *const c_char
87        }
88    }
89
90    pub unsafe extern "C" fn compare_callback(
91        raw_cb: *mut c_void,
92        a_raw: *const c_char,
93        a_len: size_t,
94        b_raw: *const c_char,
95        b_len: size_t,
96    ) -> c_int {
97        unsafe {
98            let cb: &Self = &*(raw_cb as *const Self);
99            let a: &[u8] = slice::from_raw_parts(a_raw.cast::<u8>(), a_len);
100            let b: &[u8] = slice::from_raw_parts(b_raw.cast::<u8>(), b_len);
101            (cb.compare_fn)(a, b) as c_int
102        }
103    }
104
105    pub unsafe extern "C" fn compare_ts_callback(
106        raw_cb: *mut c_void,
107        a_ts_raw: *const c_char,
108        a_ts_len: size_t,
109        b_ts_raw: *const c_char,
110        b_ts_len: size_t,
111    ) -> c_int {
112        unsafe {
113            let cb: &Self = &*(raw_cb as *const Self);
114            let a_ts: &[u8] = slice::from_raw_parts(a_ts_raw.cast::<u8>(), a_ts_len);
115            let b_ts: &[u8] = slice::from_raw_parts(b_ts_raw.cast::<u8>(), b_ts_len);
116            (cb.compare_ts_fn)(a_ts, b_ts) as c_int
117        }
118    }
119
120    pub unsafe extern "C" fn compare_without_ts_callback(
121        raw_cb: *mut c_void,
122        a_raw: *const c_char,
123        a_len: size_t,
124        a_has_ts_raw: c_uchar,
125        b_raw: *const c_char,
126        b_len: size_t,
127        b_has_ts_raw: c_uchar,
128    ) -> c_int {
129        unsafe {
130            let cb: &Self = &*(raw_cb as *const Self);
131            let a: &[u8] = slice::from_raw_parts(a_raw.cast::<u8>(), a_len);
132            let a_has_ts = a_has_ts_raw != 0;
133            let b: &[u8] = slice::from_raw_parts(b_raw.cast::<u8>(), b_len);
134            let b_has_ts = b_has_ts_raw != 0;
135            (cb.compare_without_ts_fn)(a, a_has_ts, b, b_has_ts) as c_int
136        }
137    }
138}
139
140/// A key ordering, owned by Rust and shareable with anything that needs one.
141///
142/// RocksDB never takes ownership of a comparator: `Options`, `WriteBatchWithIndex`
143/// and the rest all keep a borrowed pointer to it. Wrap it in an `Arc` and hand
144/// out clones so it outlives every user, which is what
145/// [`Options::set_comparator`](crate::Options::set_comparator) does internally.
146pub struct Comparator {
147    pub(crate) inner: NonNull<ffi::rocksdb_comparator_t>,
148}
149
150impl Comparator {
151    /// Builds a comparator from a total order over keys.
152    ///
153    /// `name` is persisted in the SST files. Changing the ordering without
154    /// changing the name gives RocksDB no way to notice, and it will read the
155    /// old files with the new ordering, so pick a new name whenever the order
156    /// changes.
157    ///
158    /// # Panics
159    ///
160    /// Panics if `name` contains an interior nul byte.
161    pub fn new(name: impl CStrLike, compare_fn: Box<CompareFn>) -> Self {
162        let cb = Box::new(ComparatorCallback {
163            name: name.into_c_string().unwrap(),
164            compare_fn,
165        });
166        let inner = unsafe {
167            ffi::rocksdb_comparator_create(
168                Box::into_raw(cb).cast::<c_void>(),
169                Some(ComparatorCallback::destructor_callback),
170                Some(ComparatorCallback::compare_callback),
171                Some(ComparatorCallback::name_callback),
172            )
173        };
174        Self {
175            inner: NonNull::new(inner).expect("rocksdb_comparator_create returned null"),
176        }
177    }
178
179    /// Builds a comparator for a column family that uses user-defined timestamps.
180    ///
181    /// `timestamp_size` is the fixed width of the timestamp suffix on every key.
182    /// `compare_fn` orders whole keys including the timestamp, `compare_ts_fn`
183    /// orders two bare timestamps, and `compare_without_ts_fn` orders keys
184    /// ignoring their timestamps. See the RocksDB [user-defined timestamp] docs.
185    ///
186    /// # Panics
187    ///
188    /// Panics if `name` contains an interior nul byte.
189    ///
190    /// [user-defined timestamp]: https://github.com/facebook/rocksdb/wiki/User-defined-Timestamp
191    pub fn with_ts(
192        name: impl CStrLike,
193        timestamp_size: usize,
194        compare_fn: Box<CompareFn>,
195        compare_ts_fn: Box<CompareTsFn>,
196        compare_without_ts_fn: Box<CompareWithoutTsFn>,
197    ) -> Self {
198        let cb = Box::new(ComparatorWithTsCallback {
199            name: name.into_c_string().unwrap(),
200            compare_fn,
201            compare_ts_fn,
202            compare_without_ts_fn,
203        });
204        let inner = unsafe {
205            ffi::rocksdb_comparator_with_ts_create(
206                Box::into_raw(cb).cast::<c_void>(),
207                Some(ComparatorWithTsCallback::destructor_callback),
208                Some(ComparatorWithTsCallback::compare_callback),
209                Some(ComparatorWithTsCallback::compare_ts_callback),
210                Some(ComparatorWithTsCallback::compare_without_ts_callback),
211                Some(ComparatorWithTsCallback::name_callback),
212                timestamp_size,
213            )
214        };
215        Self {
216            inner: NonNull::new(inner).expect("rocksdb_comparator_with_ts_create returned null"),
217        }
218    }
219
220    pub(crate) fn from_raw(inner: NonNull<ffi::rocksdb_comparator_t>) -> Self {
221        Self { inner }
222    }
223}
224
225impl Drop for Comparator {
226    fn drop(&mut self) {
227        unsafe {
228            ffi::rocksdb_comparator_destroy(self.inner.as_ptr());
229        }
230    }
231}
232
233// Safe because the callbacks are `Fn + Send + Sync` and the only other state is
234// a `CString` that is read but never mutated. Destroying the comparator needs
235// `&mut self`, so it cannot race with a comparison.
236unsafe impl Send for Comparator {}
237unsafe impl Sync for Comparator {}
rust_rocksdb/comparator.rs

rust_rocksdb/
comparator.rs
