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
|
#![allow(dead_code)]
use std::{
cmp::Ordering,
convert::TryInto,
path::{Path, PathBuf},
};
use rocksdb::{Error, Options, DB};
/// Temporary database path which calls DB::Destroy when DBPath is dropped.
pub struct DBPath {
dir: tempfile::TempDir, // kept for cleaning up during drop
path: PathBuf,
}
impl DBPath {
/// Produces a fresh (non-existent) temporary path which will be DB::destroy'ed automatically.
pub fn new(prefix: &str) -> DBPath {
let dir = tempfile::Builder::new()
.prefix(prefix)
.tempdir()
.expect("Failed to create temporary path for db.");
let path = dir.path().join("db");
DBPath { dir, path }
}
}
impl Drop for DBPath {
fn drop(&mut self) {
let opts = Options::default();
DB::destroy(&opts, &self.path).expect("Failed to destroy temporary DB");
}
}
/// Convert a DBPath ref to a Path ref.
/// We don't implement this for DBPath values because we want them to
/// exist until the end of their scope, not get passed into functions and
/// dropped early.
impl AsRef<Path> for &DBPath {
fn as_ref(&self) -> &Path {
&self.path
}
}
type Pair = (Box<[u8]>, Box<[u8]>);
pub fn pair(left: &[u8], right: &[u8]) -> Pair {
(Box::from(left), Box::from(right))
}
#[track_caller]
pub fn assert_iter(iter: impl Iterator<Item = Result<Pair, Error>>, want: &[Pair]) {
let got = iter.collect::<Result<Vec<_>, _>>().unwrap();
assert_eq!(got.as_slice(), want);
}
#[track_caller]
pub fn assert_iter_reversed(iter: impl Iterator<Item = Result<Pair, Error>>, want: &[Pair]) {
let mut got = iter.collect::<Result<Vec<_>, _>>().unwrap();
got.reverse();
assert_eq!(got.as_slice(), want);
}
/// A timestamp type we use in testing [user-defined timestamp](https://github.com/facebook/rocksdb/wiki/User-defined-Timestamp).
/// This is a `u64` in little endian encoding.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct U64Timestamp([u8; Self::SIZE]);
impl U64Timestamp {
pub const SIZE: usize = 8;
pub fn new(ts: u64) -> Self {
Self(ts.to_le_bytes())
}
}
impl From<&[u8]> for U64Timestamp {
fn from(slice: &[u8]) -> Self {
assert_eq!(
slice.len(),
Self::SIZE,
"incorrect timestamp length: {}, should be {}",
slice.len(),
Self::SIZE
);
Self(slice.try_into().unwrap())
}
}
impl From<U64Timestamp> for Vec<u8> {
fn from(ts: U64Timestamp) -> Self {
ts.0.to_vec()
}
}
impl AsRef<[u8]> for U64Timestamp {
fn as_ref(&self) -> &[u8] {
&self.0
}
}
impl PartialOrd for U64Timestamp {
fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
Some(self.cmp(other))
}
}
impl Ord for U64Timestamp {
fn cmp(&self, other: &Self) -> Ordering {
let lhs = u64::from_le_bytes(self.0);
let rhs = u64::from_le_bytes(other.0);
lhs.cmp(&rhs)
}
}
/// A comparator for use in column families with [user-defined timestamp](https://github.com/facebook/rocksdb/wiki/User-defined-Timestamp)
/// enabled. This comparator assumes `u64` timestamp in little endian encoding.
/// This is the same behavior as RocksDB's built-in comparator.
///
/// Adapted from C++ and Golang implementations from:
/// - [rocksdb](https://github.com/facebook/rocksdb/blob/v9.4.0/test_util/testutil.cc#L112)
/// - [gorocksdb](https://github.com/linxGnu/grocksdb/blob/v1.9.2/db_ts_test.go#L167)
/// - [SeiDB](https://github.com/sei-protocol/sei-db/blob/v0.0.41/ss/rocksdb/comparator.go)
pub struct U64Comparator;
impl U64Comparator {
pub const NAME: &'static str = "rust-rocksdb.U64Comparator";
pub fn compare(a: &[u8], b: &[u8]) -> Ordering {
// First, compare the keys without timestamps. If the keys are different,
// then we don't have to consider the timestamps at all.
let ord = Self::compare_without_ts(a, true, b, true);
if ord != Ordering::Equal {
return ord;
}
// The keys are the same, so now we compare the timestamps.
// The larger (i.e. newer) key should come first, hence the `reverse`.
Self::compare_ts(
extract_timestamp_from_user_key(a),
extract_timestamp_from_user_key(b),
)
.reverse()
}
pub fn compare_ts(bz1: &[u8], bz2: &[u8]) -> Ordering {
let ts1 = U64Timestamp::from(bz1);
let ts2 = U64Timestamp::from(bz2);
ts1.cmp(&ts2)
}
pub fn compare_without_ts(
mut a: &[u8],
a_has_ts: bool,
mut b: &[u8],
b_has_ts: bool,
) -> Ordering {
if a_has_ts {
a = strip_timestamp_from_user_key(a);
}
if b_has_ts {
b = strip_timestamp_from_user_key(b);
}
a.cmp(b)
}
}
fn extract_timestamp_from_user_key(key: &[u8]) -> &[u8] {
&key[(key.len() - U64Timestamp::SIZE)..]
}
fn strip_timestamp_from_user_key(key: &[u8]) -> &[u8] {
&key[..(key.len() - U64Timestamp::SIZE)]
}
|
