massa_event_cache/
rocksdb_operator.rs

1use rocksdb::MergeOperands;
2
3pub fn counter_merge(
4    _key: &[u8],
5    existing_val: Option<&[u8]>,
6    operands: &MergeOperands,
7) -> Option<Vec<u8>> {
8    let counter_current_value = if let Some(existing_val) = existing_val {
9        u64::from_be_bytes(existing_val.try_into().unwrap())
10    } else {
11        0
12    };
13
14    let counter_value = operands.iter().fold(counter_current_value, |mut acc, x| {
15        let incr_value = i64::from_be_bytes(x.try_into().unwrap());
16        acc = acc.saturating_add_signed(incr_value);
17        acc
18    });
19
20    Some(counter_value.to_be_bytes().to_vec())
21}
22
23#[cfg(test)]
24mod tests {
25    use super::*;
26    // std
27    // third-party
28    use rocksdb::{Options, DB};
29    use serial_test::serial;
30    use tempfile::TempDir;
31
32    #[test]
33    #[serial]
34    fn test_operator() {
35        let tmp_path = TempDir::new().unwrap().path().to_path_buf();
36        let options = {
37            let mut opts = Options::default();
38            opts.create_if_missing(true);
39            opts.set_merge_operator_associative("counter merge operator", counter_merge);
40            opts
41        };
42        let db = DB::open(&options, tmp_path).unwrap();
43        let key_1 = "foo1";
44        let key_2 = "baz42";
45        db.put(key_1, 0u64.to_be_bytes()).unwrap();
46        db.put(key_2, 0u64.to_be_bytes()).unwrap();
47
48        let value = db.get(key_1).unwrap().unwrap();
49        assert_eq!(u64::from_be_bytes(value.try_into().unwrap()), 0);
50        let value2 = db.get(key_2).unwrap().unwrap();
51        assert_eq!(u64::from_be_bytes(value2.try_into().unwrap()), 0);
52
53        // key_1 counter += 1
54        db.merge(key_1, 1i64.to_be_bytes()).unwrap();
55
56        let value = db.get(key_1).unwrap().unwrap();
57        assert_eq!(u64::from_be_bytes(value.try_into().unwrap()), 1);
58        let value2 = db.get(key_2).unwrap().unwrap();
59        assert_eq!(u64::from_be_bytes(value2.try_into().unwrap()), 0);
60
61        // key_2 counter += 9
62        db.merge(key_2, 9i64.to_be_bytes()).unwrap();
63        // key_2 counter += 1
64        db.merge(key_2, 1i64.to_be_bytes()).unwrap();
65        // key_2 counter += 32
66        db.merge(key_2, 32i64.to_be_bytes()).unwrap();
67
68        let value = db.get(key_1).unwrap().unwrap();
69        assert_eq!(u64::from_be_bytes(value.try_into().unwrap()), 1);
70        let value2 = db.get(key_2).unwrap().unwrap();
71        assert_eq!(u64::from_be_bytes(value2.try_into().unwrap()), 42);
72    }
73
74    #[test]
75    #[serial]
76    fn test_operator_2() {
77        let tmp_path = TempDir::new().unwrap().path().to_path_buf();
78        let options = {
79            let mut opts = Options::default();
80            opts.create_if_missing(true);
81            opts.set_merge_operator_associative("counter merge operator", counter_merge);
82            opts
83        };
84        let db = DB::open(&options, tmp_path).unwrap();
85        let key_1 = "foo1";
86        let key_2 = "baz42";
87        db.put(key_1, 0u64.to_be_bytes()).unwrap();
88        db.put(key_2, 0u64.to_be_bytes()).unwrap();
89
90        let value = db.get(key_1).unwrap().unwrap();
91        assert_eq!(u64::from_be_bytes(value.try_into().unwrap()), 0);
92        let value2 = db.get(key_2).unwrap().unwrap();
93        assert_eq!(u64::from_be_bytes(value2.try_into().unwrap()), 0);
94
95        // key_1 counter += 1
96        db.merge(key_1, 1i64.to_be_bytes()).unwrap();
97
98        let value = db.get(key_1).unwrap().unwrap();
99        assert_eq!(u64::from_be_bytes(value.try_into().unwrap()), 1);
100        let value2 = db.get(key_2).unwrap().unwrap();
101        assert_eq!(u64::from_be_bytes(value2.try_into().unwrap()), 0);
102
103        db.merge(key_1, (-3i64).to_be_bytes()).unwrap();
104
105        let value = db.get(key_1).unwrap().unwrap();
106        assert_eq!(u64::from_be_bytes(value.try_into().unwrap()), 0);
107        let value2 = db.get(key_2).unwrap().unwrap();
108        assert_eq!(u64::from_be_bytes(value2.try_into().unwrap()), 0);
109    }
110}
massa_event_cache/rocksdb_operator.rs

massa_event_cache/
rocksdb_operator.rs
