1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
// Copyright 2020 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 libc::{c_char, c_int, c_uchar, c_void, size_t};
use std::ffi::{CStr, CString};
use std::slice;

/// Decision about how to handle compacting an object
///
/// This is returned by a compaction filter callback. Depending
/// on the value, the object may be kept, removed, or changed
/// in the database during a compaction.
pub enum Decision {
    /// Keep the old value
    Keep,
    /// Remove the object from the database
    Remove,
    /// Change the value for the key
    Change(&'static [u8]),
}

/// CompactionFilter allows an application to modify/delete a key-value at
/// the time of compaction.
pub trait CompactionFilter {
    /// The compaction process invokes this
    /// method for kv that is being compacted. The application can inspect
    /// the existing value of the key and make decision based on it.
    ///
    /// Key-Values that are results of merge operation during compaction are not
    /// passed into this function. Currently, when you have a mix of Put()s and
    /// Merge()s on a same key, we only guarantee to process the merge operands
    /// through the compaction filters. Put()s might be processed, or might not.
    ///
    /// When the value is to be preserved, the application has the option
    /// to modify the existing_value and pass it back through new_value.
    /// value_changed needs to be set to true in this case.
    ///
    /// Note that RocksDB snapshots (i.e. call GetSnapshot() API on a
    /// DB* object) will not guarantee to preserve the state of the DB with
    /// CompactionFilter. Data seen from a snapshot might disappear after a
    /// compaction finishes. If you use snapshots, think twice about whether you
    /// want to use compaction filter and whether you are using it in a safe way.
    ///
    /// If the CompactionFilter was created by a factory, then it will only ever
    /// be used by a single thread that is doing the compaction run, and this
    /// call does not need to be thread-safe.  However, multiple filters may be
    /// in existence and operating concurrently.
    fn filter(&mut self, level: u32, key: &[u8], value: &[u8]) -> Decision;

    /// Returns a name that identifies this compaction filter.
    /// The name will be printed to LOG file on start up for diagnosis.
    fn name(&self) -> &CStr;
}

/// Function to filter compaction with.
///
/// This function takes the level of compaction, the key, and the existing value
/// and returns the decision about how to handle the Key-Value pair.
///
///  See [Options::set_compaction_filter][set_compaction_filter] for more details
///
///  [set_compaction_filter]: ../struct.Options.html#method.set_compaction_filter
pub trait CompactionFilterFn: FnMut(u32, &[u8], &[u8]) -> Decision {}
impl<F> CompactionFilterFn for F where F: FnMut(u32, &[u8], &[u8]) -> Decision + Send + 'static {}

pub struct CompactionFilterCallback<F>
where
    F: CompactionFilterFn,
{
    pub name: CString,
    pub filter_fn: F,
}

impl<F> CompactionFilter for CompactionFilterCallback<F>
where
    F: CompactionFilterFn,
{
    fn name(&self) -> &CStr {
        self.name.as_c_str()
    }

    fn filter(&mut self, level: u32, key: &[u8], value: &[u8]) -> Decision {
        (self.filter_fn)(level, key, value)
    }
}

pub unsafe extern "C" fn destructor_callback<F>(raw_cb: *mut c_void)
where
    F: CompactionFilter,
{
    drop(Box::from_raw(raw_cb as *mut F));
}

pub unsafe extern "C" fn name_callback<F>(raw_cb: *mut c_void) -> *const c_char
where
    F: CompactionFilter,
{
    let cb = &*(raw_cb as *mut F);
    cb.name().as_ptr()
}

pub unsafe extern "C" fn filter_callback<F>(
    raw_cb: *mut c_void,
    level: c_int,
    raw_key: *const c_char,
    key_length: size_t,
    existing_value: *const c_char,
    value_length: size_t,
    new_value: *mut *mut c_char,
    new_value_length: *mut size_t,
    value_changed: *mut c_uchar,
) -> c_uchar
where
    F: CompactionFilter,
{
    use self::Decision::{Change, Keep, Remove};

    let cb = &mut *(raw_cb as *mut F);
    let key = slice::from_raw_parts(raw_key as *const u8, key_length);
    let oldval = slice::from_raw_parts(existing_value as *const u8, value_length);
    let result = cb.filter(level as u32, key, oldval);
    match result {
        Keep => 0,
        Remove => 1,
        Change(newval) => {
            *new_value = newval.as_ptr() as *mut c_char;
            *new_value_length = newval.len() as size_t;
            *value_changed = 1_u8;
            0
        }
    }
}

#[cfg(test)]
#[allow(unused_variables)]
fn test_filter(level: u32, key: &[u8], value: &[u8]) -> Decision {
    use self::Decision::{Change, Keep, Remove};
    match key.first() {
        Some(&b'_') => Remove,
        Some(&b'%') => Change(b"secret"),
        _ => Keep,
    }
}

#[test]
fn compaction_filter_test() {
    use crate::{Options, DB};

    let path = "_rust_rocksdb_filter_test";
    let mut opts = Options::default();
    opts.create_if_missing(true);
    opts.set_compaction_filter("test", test_filter);
    {
        let db = DB::open(&opts, path).unwrap();
        let _r = db.put(b"k1", b"a");
        let _r = db.put(b"_k", b"b");
        let _r = db.put(b"%k", b"c");
        db.compact_range(None::<&[u8]>, None::<&[u8]>);
        assert_eq!(&*db.get(b"k1").unwrap().unwrap(), b"a");
        assert!(db.get(b"_k").unwrap().is_none());
        assert_eq!(&*db.get(b"%k").unwrap().unwrap(), b"secret");
    }
    let result = DB::destroy(&opts, path);
    assert!(result.is_ok());
}