keyvaluedb/
lib.rs

1//! Key-Value store abstraction.
2
3#![deny(clippy::all)]
4
5mod io_stats;
6
7pub use io_stats::{IoStats, Kind as IoStatsKind};
8use std::future::Future;
9use std::io;
10use std::pin::Pin;
11
12/// Required length of prefixes.
13pub const PREFIX_LEN: usize = 12;
14
15/// Database value.
16pub type DBValue = Vec<u8>;
17pub type DBKey = Vec<u8>;
18pub type DBKeyValue = (DBKey, DBValue);
19pub type DBKeyRef<'a> = &'a DBKey;
20pub type DBKeyValueRef<'a> = (&'a DBKey, &'a DBValue);
21
22/// Write transaction. Batches a sequence of put/delete operations for efficiency.
23#[derive(Default, Debug, Clone, Eq, PartialEq)]
24pub struct DBTransaction {
25    /// Database operations.
26    pub ops: Vec<DBOp>,
27}
28
29/// Database operation.
30#[derive(Clone, Debug, PartialEq, Eq)]
31pub enum DBOp {
32    Insert {
33        col: u32,
34        key: DBKey,
35        value: DBValue,
36    },
37    Delete {
38        col: u32,
39        key: DBKey,
40    },
41    DeletePrefix {
42        col: u32,
43        prefix: DBKey,
44    },
45}
46
47impl DBOp {
48    /// Returns the key associated with this operation.
49    pub fn key(&self) -> &DBKey {
50        match *self {
51            DBOp::Insert { ref key, .. } => key,
52            DBOp::Delete { ref key, .. } => key,
53            DBOp::DeletePrefix { ref prefix, .. } => prefix,
54        }
55    }
56    /// Returns the value associated with this operation.
57    pub fn value(&self) -> Option<&DBValue> {
58        match *self {
59            DBOp::Insert { ref value, .. } => Some(value),
60            DBOp::Delete { .. } => None,
61            DBOp::DeletePrefix { .. } => None,
62        }
63    }
64    /// Returns the column associated with this operation.
65    pub fn col(&self) -> u32 {
66        match *self {
67            DBOp::Insert { col, .. } => col,
68            DBOp::Delete { col, .. } => col,
69            DBOp::DeletePrefix { col, .. } => col,
70        }
71    }
72}
73
74impl DBTransaction {
75    /// Create new transaction.
76    pub fn new() -> DBTransaction {
77        DBTransaction::with_capacity(256)
78    }
79
80    /// Create new transaction with capacity.
81    pub fn with_capacity(cap: usize) -> DBTransaction {
82        DBTransaction {
83            ops: Vec::with_capacity(cap),
84        }
85    }
86
87    /// Insert a key-value pair in the transaction. Any existing value will be overwritten upon write.
88    pub fn put<K, V>(&mut self, col: u32, key: K, value: V)
89    where
90        K: AsRef<[u8]>,
91        V: AsRef<[u8]>,
92    {
93        self.ops.push(DBOp::Insert {
94            col,
95            key: DBKey::from(key.as_ref()),
96            value: DBValue::from(value.as_ref()),
97        })
98    }
99    pub fn put_owned(&mut self, col: u32, key: Vec<u8>, value: Vec<u8>) {
100        self.ops.push(DBOp::Insert { col, key, value })
101    }
102
103    /// Delete value by key.
104    pub fn delete<K>(&mut self, col: u32, key: K)
105    where
106        K: AsRef<[u8]>,
107    {
108        self.ops.push(DBOp::Delete {
109            col,
110            key: DBKey::from(key.as_ref()),
111        });
112    }
113    pub fn delete_owned(&mut self, col: u32, key: Vec<u8>) {
114        self.ops.push(DBOp::Delete { col, key });
115    }
116
117    /// Delete all values with the given key prefix.
118    /// Using an empty prefix here will remove all keys
119    /// (all keys start with the empty prefix).
120    pub fn delete_prefix<K>(&mut self, col: u32, prefix: K)
121    where
122        K: AsRef<[u8]>,
123    {
124        self.ops.push(DBOp::DeletePrefix {
125            col,
126            prefix: DBKey::from(prefix.as_ref()),
127        });
128    }
129    pub fn delete_prefix_owned(&mut self, col: u32, prefix: Vec<u8>) {
130        self.ops.push(DBOp::DeletePrefix { col, prefix });
131    }
132}
133
134/// Transaction Result, returns the transaction unchanged upon error
135#[derive(Debug)]
136pub struct DBTransactionError {
137    pub error: io::Error,
138    pub transaction: DBTransaction,
139}
140
141impl std::fmt::Display for DBTransactionError {
142    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
143        f.debug_struct("TransactionError")
144            .field("error", &self.error)
145            .finish()
146    }
147}
148impl std::error::Error for DBTransactionError {}
149
150pub type KeyValueDBPinBoxFuture<'a, T> = Pin<Box<dyn Future<Output = T> + Send + 'a>>;
151
152/// Generic key-value database.
153///
154/// The `KeyValueDB` deals with "column families", which can be thought of as distinct
155/// stores within a database. Keys written in one column family will not be accessible from
156/// any other. The number of column families must be specified at initialization, with a
157/// differing interface for each database.
158///
159/// The API laid out here, along with the `Sync` bound implies interior synchronization for
160/// implementation.
161/// Clone is here so we can pass an owned self to async functions, requiring interior locked mutability.
162pub trait KeyValueDB: Sync + Send + Clone + 'static {
163    /// Helper to create a new transaction.
164    fn transaction(&self) -> DBTransaction {
165        DBTransaction::new()
166    }
167
168    /// Get a value by key.
169    fn get<'a>(
170        &'a self,
171        col: u32,
172        key: &'a [u8],
173    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBValue>>>;
174
175    /// Remove a value by key, returning the old value
176    fn delete<'a>(
177        &'a self,
178        col: u32,
179        key: &'a [u8],
180    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBValue>>>;
181
182    /// Write a transaction of changes to the backing store.
183    fn write(
184        &self,
185        transaction: DBTransaction,
186    ) -> KeyValueDBPinBoxFuture<'_, Result<(), DBTransactionError>>;
187
188    /// Iterate over the data for a given column.
189    /// Return all key/value pairs, optionally where the key starts with the given prefix.
190    /// Iterator closure returns true for more items, false to stop iteration.
191    fn iter<
192        'a,
193        T: Send + 'static,
194        C: Send + 'static,
195        F: FnMut(&mut C, DBKeyValueRef) -> io::Result<Option<T>> + Send + Sync + 'static,
196    >(
197        &'a self,
198        col: u32,
199        prefix: Option<&'a [u8]>,
200        context: C,
201        f: F,
202    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>>;
203
204    /// Iterate over the data for a given column.
205    /// Return all keys, optionally where the key starts with the given prefix.
206    /// Iterator closure returns true for more items, false to stop iteration.
207    fn iter_keys<
208        'a,
209        T: Send + 'static,
210        C: Send + 'static,
211        F: FnMut(&mut C, DBKeyRef) -> io::Result<Option<T>> + Send + Sync + 'static,
212    >(
213        &'a self,
214        col: u32,
215        prefix: Option<&'a [u8]>,
216        context: C,
217        f: F,
218    ) -> KeyValueDBPinBoxFuture<'a, io::Result<(C, Option<T>)>>;
219
220    /// Query statistics.
221    ///
222    /// Not all keyvaluedb implementations are able or expected to implement this, so by
223    /// default, empty statistics is returned. Also, not all keyvaluedb implementation
224    /// can return every statistic or configured to do so (some statistics gathering
225    /// may impede the performance and might be off by default).
226    fn io_stats(&self, _kind: IoStatsKind) -> IoStats {
227        IoStats::empty()
228    }
229
230    /// The number of column families in the db.
231    fn num_columns(&self) -> io::Result<u32>;
232
233    /// The number of keys in a column (estimated).
234    fn num_keys(&self, col: u32) -> KeyValueDBPinBoxFuture<'_, io::Result<u64>>;
235
236    /// Check for the existence of a value by key.
237    fn has_key<'a>(
238        &'a self,
239        col: u32,
240        key: &'a [u8],
241    ) -> KeyValueDBPinBoxFuture<'a, io::Result<bool>> {
242        let this = self.clone();
243        Box::pin(async move { Ok(this.get(col, key).await?.is_some()) })
244    }
245
246    /// Check for the existence of a value by prefix.
247    fn has_prefix<'a>(
248        &'a self,
249        col: u32,
250        prefix: &'a [u8],
251    ) -> KeyValueDBPinBoxFuture<'a, io::Result<bool>> {
252        let this = self.clone();
253        Box::pin(async move {
254            let (_, out) = this
255                .iter_keys(col, Some(prefix), (), |_, _| Ok(Some(())))
256                .await?;
257            Ok(out.is_some())
258        })
259    }
260
261    /// Get the first value matching the given prefix.
262    fn first_with_prefix<'a>(
263        &'a self,
264        col: u32,
265        prefix: &'a [u8],
266    ) -> KeyValueDBPinBoxFuture<'a, io::Result<Option<DBKeyValue>>> {
267        let this = self.clone();
268        Box::pin(async move {
269            let (_, out) = this
270                .iter(col, Some(prefix), (), |_, (k, v)| {
271                    Ok(Some((k.to_vec(), v.to_vec())))
272                })
273                .await?;
274            Ok(out)
275        })
276    }
277
278    /// Cleanup/Vacuum database
279    fn cleanup(&self) -> KeyValueDBPinBoxFuture<'_, io::Result<()>> {
280        Box::pin(async { Ok(()) })
281    }
282}
283
284/// For a given start prefix (inclusive), returns the correct end prefix (non-inclusive).
285/// This assumes the key bytes are ordered in lexicographical order.
286/// Since key length is not limited, for some case we return `None` because there is
287/// no bounded limit (every keys in the serie `[]`, `[255]`, `[255, 255]` ...).
288pub fn end_prefix(prefix: &[u8]) -> Option<Vec<u8>> {
289    let mut end_range = prefix.to_vec();
290    while let Some(0xff) = end_range.last() {
291        end_range.pop();
292    }
293    if let Some(byte) = end_range.last_mut() {
294        *byte += 1;
295        Some(end_range)
296    } else {
297        None
298    }
299}
300
301#[cfg(test)]
302mod test {
303    use super::end_prefix;
304
305    #[test]
306    fn end_prefix_test() {
307        assert_eq!(end_prefix(&[5, 6, 7]), Some(vec![5, 6, 8]));
308        assert_eq!(end_prefix(&[5, 6, 255]), Some(vec![5, 7]));
309        // This is not equal as the result is before start.
310        assert_ne!(end_prefix(&[5, 255, 255]), Some(vec![5, 255]));
311        // This is equal ([5, 255] will not be deleted because
312        // it is before start).
313        assert_eq!(end_prefix(&[5, 255, 255]), Some(vec![6]));
314        assert_eq!(end_prefix(&[255, 255, 255]), None);
315
316        assert_eq!(end_prefix(&[0x00, 0xff]), Some(vec![0x01]));
317        assert_eq!(end_prefix(&[0xff]), None);
318        assert_eq!(end_prefix(&[]), None);
319        assert_eq!(end_prefix(b"0"), Some(b"1".to_vec()));
320    }
321}
keyvaluedb/lib.rs

keyvaluedb/
lib.rs
