persistent_queue/
store.rs

1//! The [`Store`] trait and its built-in backends.
2
3use std::error::Error as StdError;
4
5/// A single write in a [`Store::commit`] batch.
6pub enum Op<'a> {
7    /// Insert or overwrite `key` with `value`.
8    Put(&'a [u8], &'a [u8]),
9    /// Remove `key`.
10    Delete(&'a [u8]),
11}
12
13/// A key/value pair returned by a store seek.
14pub type KeyValue = (Vec<u8>, Vec<u8>);
15
16/// An ordered key/value byte store: the durable substrate under the queue.
17///
18/// Keys sort lexicographically. The queue needs forward and backward seeks and
19/// one atomic, optionally durable, batch write; it never asks the store to know
20/// anything about queues.
21pub trait Store: Send + Sync {
22    /// The backend's error type.
23    type Error: StdError + Send + Sync + 'static;
24
25    /// Value for an exact key.
26    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error>;
27
28    /// Smallest entry whose key is `>= from`.
29    fn seek(&self, from: &[u8]) -> Result<Option<KeyValue>, Self::Error>;
30
31    /// Greatest entry whose key is `<= upto`.
32    fn seek_back(&self, upto: &[u8]) -> Result<Option<KeyValue>, Self::Error>;
33
34    /// Apply `ops`. When `durable`, do not return until they survive a crash.
35    fn commit(&self, ops: &[Op<'_>], durable: bool) -> Result<(), Self::Error>;
36}
37
38// ---- mem ----
39
40/// In-memory [`Store`] backed by a `BTreeMap`. Not persistent; `durable` is a
41/// no-op. Useful as the default, for tests, and as a benchmark baseline.
42pub struct MemStore {
43    map: crate::sync::Mutex<std::collections::BTreeMap<Vec<u8>, Vec<u8>>>,
44}
45
46impl MemStore {
47    /// Create an empty store.
48    pub fn new() -> Self {
49        Self {
50            map: crate::sync::Mutex::new(std::collections::BTreeMap::new()),
51        }
52    }
53}
54
55impl Default for MemStore {
56    fn default() -> Self {
57        Self::new()
58    }
59}
60
61impl Store for MemStore {
62    type Error = std::convert::Infallible;
63
64    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
65        Ok(self.map.lock().unwrap().get(key).cloned())
66    }
67
68    fn seek(&self, from: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
69        Ok(self
70            .map
71            .lock()
72            .unwrap()
73            .range(from.to_vec()..)
74            .next()
75            .map(|(k, v)| (k.clone(), v.clone())))
76    }
77
78    fn seek_back(&self, upto: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
79        Ok(self
80            .map
81            .lock()
82            .unwrap()
83            .range(..=upto.to_vec())
84            .next_back()
85            .map(|(k, v)| (k.clone(), v.clone())))
86    }
87
88    fn commit(&self, ops: &[Op<'_>], _durable: bool) -> Result<(), Self::Error> {
89        let mut map = self.map.lock().unwrap();
90        for op in ops {
91            match op {
92                Op::Put(k, v) => {
93                    map.insert(k.to_vec(), v.to_vec());
94                }
95                Op::Delete(k) => {
96                    map.remove(*k);
97                }
98            }
99        }
100        Ok(())
101    }
102}
103
104// ---- sled ----
105
106/// [`Store`] backed by a [sled](https://docs.rs/sled) database. Requires the
107/// `sled` feature.
108///
109/// Only one process may open a given database directory at a time; sled takes an
110/// exclusive lock. A backend error on open, including a corrupt store, is surfaced
111/// by [`Builder::open`](crate::Builder::open) as
112/// [`OpenError::Store`](crate::OpenError::Store); the queue does not auto-repair.
113#[cfg(feature = "sled")]
114pub struct SledStore {
115    db: sled::Db,
116}
117
118#[cfg(feature = "sled")]
119impl SledStore {
120    /// Open (creating if needed) a sled database at `path`.
121    pub fn open(path: impl AsRef<std::path::Path>) -> sled::Result<Self> {
122        Ok(Self {
123            db: sled::open(path)?,
124        })
125    }
126
127    /// Wrap an already-open sled database.
128    pub fn from_db(db: sled::Db) -> Self {
129        Self { db }
130    }
131}
132
133#[cfg(feature = "sled")]
134impl Store for SledStore {
135    type Error = sled::Error;
136
137    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
138        Ok(self.db.get(key)?.map(|v| v.to_vec()))
139    }
140
141    fn seek(&self, from: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
142        match self.db.range(from.to_vec()..).next() {
143            Some(r) => {
144                let (k, v) = r?;
145                Ok(Some((k.to_vec(), v.to_vec())))
146            }
147            None => Ok(None),
148        }
149    }
150
151    fn seek_back(&self, upto: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
152        match self.db.range(..=upto.to_vec()).next_back() {
153            Some(r) => {
154                let (k, v) = r?;
155                Ok(Some((k.to_vec(), v.to_vec())))
156            }
157            None => Ok(None),
158        }
159    }
160
161    fn commit(&self, ops: &[Op<'_>], durable: bool) -> Result<(), Self::Error> {
162        let mut batch = sled::Batch::default();
163        for op in ops {
164            match op {
165                Op::Put(k, v) => batch.insert(*k, *v),
166                Op::Delete(k) => batch.remove(*k),
167            }
168        }
169        self.db.apply_batch(batch)?;
170        if durable {
171            self.db.flush()?;
172        }
173        Ok(())
174    }
175}
176
177// ---- redb ----
178
179#[cfg(feature = "redb")]
180const REDB_TABLE: redb::TableDefinition<'static, &[u8], &[u8]> =
181    redb::TableDefinition::new("entries");
182
183/// [`Store`] backed by a [redb](https://docs.rs/redb) database. Requires the
184/// `redb` feature.
185///
186/// Only one process may open a given database file at a time; redb takes an
187/// exclusive lock. A backend error on open, including a corrupt store, is surfaced
188/// by [`Builder::open`](crate::Builder::open) as
189/// [`OpenError::Store`](crate::OpenError::Store); the queue does not auto-repair.
190#[cfg(feature = "redb")]
191pub struct RedbStore {
192    db: redb::Database,
193}
194
195#[cfg(feature = "redb")]
196impl RedbStore {
197    /// Open (creating if needed) a redb database at `path`.
198    // The Store trait surfaces redb::Error unboxed, so keep open consistent.
199    #[allow(clippy::result_large_err)]
200    pub fn open(path: impl AsRef<std::path::Path>) -> Result<Self, redb::Error> {
201        let db = redb::Database::create(path)?;
202        let wtx = db.begin_write()?;
203        wtx.open_table(REDB_TABLE)?;
204        wtx.commit()?;
205        Ok(Self { db })
206    }
207
208    /// Wrap an already-open redb database.
209    pub fn from_db(db: redb::Database) -> Self {
210        Self { db }
211    }
212}
213
214#[cfg(feature = "redb")]
215impl Store for RedbStore {
216    type Error = redb::Error;
217
218    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
219        let rtx = self.db.begin_read()?;
220        let table = rtx.open_table(REDB_TABLE)?;
221        Ok(table.get(key)?.map(|g| g.value().to_vec()))
222    }
223
224    fn seek(&self, from: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
225        let rtx = self.db.begin_read()?;
226        let table = rtx.open_table(REDB_TABLE)?;
227        match table.range::<&[u8]>(from..)?.next() {
228            Some(r) => {
229                let (k, v) = r?;
230                Ok(Some((k.value().to_vec(), v.value().to_vec())))
231            }
232            None => Ok(None),
233        }
234    }
235
236    fn seek_back(&self, upto: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
237        let rtx = self.db.begin_read()?;
238        let table = rtx.open_table(REDB_TABLE)?;
239        match table.range::<&[u8]>(..=upto)?.next_back() {
240            Some(r) => {
241                let (k, v) = r?;
242                Ok(Some((k.value().to_vec(), v.value().to_vec())))
243            }
244            None => Ok(None),
245        }
246    }
247
248    fn commit(&self, ops: &[Op<'_>], durable: bool) -> Result<(), Self::Error> {
249        let mut wtx = self.db.begin_write()?;
250        if !durable {
251            wtx.set_durability(redb::Durability::None);
252        }
253        {
254            let mut table = wtx.open_table(REDB_TABLE)?;
255            for op in ops {
256                match op {
257                    Op::Put(k, v) => {
258                        table.insert(*k, *v)?;
259                    }
260                    Op::Delete(k) => {
261                        table.remove(*k)?;
262                    }
263                }
264            }
265        }
266        wtx.commit()?;
267        Ok(())
268    }
269}
270
271// ---- rocksdb ----
272
273/// [`Store`] backed by [RocksDB](https://docs.rs/rocksdb). Requires the `rocksdb`
274/// feature, which builds a bundled C++ RocksDB (needs a C++ toolchain).
275///
276/// Only one process may open a given database directory at a time; RocksDB takes an
277/// exclusive lock. A backend error on open, including a corrupt store, is surfaced
278/// by [`Builder::open`](crate::Builder::open) as
279/// [`OpenError::Store`](crate::OpenError::Store); the queue does not auto-repair.
280///
281/// The bundled RocksDB is compiled from source and statically linked into your binary
282/// (nothing extra to ship), but it needs a C++ toolchain to build, adds compile time
283/// and binary size, and links the C++ runtime - so a fully static (musl) binary is
284/// hard. For a small, pure-Rust, fully static binary, use `sled` or `redb`.
285#[cfg(feature = "rocksdb")]
286pub struct RocksStore {
287    db: rocksdb::DB,
288}
289
290#[cfg(feature = "rocksdb")]
291impl RocksStore {
292    /// Open (creating if needed) a RocksDB database at `path`.
293    pub fn open(path: impl AsRef<std::path::Path>) -> Result<Self, rocksdb::Error> {
294        Ok(Self {
295            db: rocksdb::DB::open_default(path)?,
296        })
297    }
298
299    /// Wrap an already-open RocksDB database.
300    pub fn from_db(db: rocksdb::DB) -> Self {
301        Self { db }
302    }
303}
304
305#[cfg(feature = "rocksdb")]
306impl Store for RocksStore {
307    type Error = rocksdb::Error;
308
309    fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>, Self::Error> {
310        self.db.get(key)
311    }
312
313    fn seek(&self, from: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
314        let mut iter = self.db.iterator(rocksdb::IteratorMode::From(
315            from,
316            rocksdb::Direction::Forward,
317        ));
318        match iter.next() {
319            Some(Ok((k, v))) => Ok(Some((k.to_vec(), v.to_vec()))),
320            Some(Err(e)) => Err(e),
321            None => Ok(None),
322        }
323    }
324
325    fn seek_back(&self, upto: &[u8]) -> Result<Option<KeyValue>, Self::Error> {
326        let mut iter = self.db.iterator(rocksdb::IteratorMode::From(
327            upto,
328            rocksdb::Direction::Reverse,
329        ));
330        match iter.next() {
331            Some(Ok((k, v))) => Ok(Some((k.to_vec(), v.to_vec()))),
332            Some(Err(e)) => Err(e),
333            None => Ok(None),
334        }
335    }
336
337    fn commit(&self, ops: &[Op<'_>], durable: bool) -> Result<(), Self::Error> {
338        let mut batch = rocksdb::WriteBatch::default();
339        for op in ops {
340            match op {
341                Op::Put(k, v) => batch.put(*k, *v),
342                Op::Delete(k) => batch.delete(*k),
343            }
344        }
345        let mut opts = rocksdb::WriteOptions::default();
346        opts.set_sync(durable);
347        self.db.write_opt(batch, &opts)
348    }
349}
350
351#[cfg(test)]
352mod tests {
353    use super::*;
354
355    #[test]
356    fn mem_store_contract() {
357        contract(MemStore::new());
358    }
359
360    #[cfg(feature = "sled")]
361    #[test]
362    fn sled_store_contract() {
363        let dir = tempfile::tempdir().unwrap();
364        contract(SledStore::open(dir.path().join("db")).unwrap());
365    }
366
367    #[cfg(feature = "redb")]
368    #[test]
369    fn redb_store_contract() {
370        let dir = tempfile::tempdir().unwrap();
371        contract(RedbStore::open(dir.path().join("db.redb")).unwrap());
372    }
373
374    #[cfg(feature = "rocksdb")]
375    #[test]
376    fn rocksdb_store_contract() {
377        let dir = tempfile::tempdir().unwrap();
378        contract(RocksStore::open(dir.path().join("db")).unwrap());
379    }
380
381    // A store reads keys back exactly, keeps them in byte-lexicographic order, and
382    // applies deletes; seeks find the nearest key in each direction.
383    fn contract<S: Store>(store: S) {
384        assert!(store.get(b"missing").unwrap().is_none());
385        assert!(store.seek(b"a").unwrap().is_none());
386
387        store
388            .commit(
389                &[
390                    Op::Put(b"b", b"2"),
391                    Op::Put(b"a", b"1"),
392                    Op::Put(b"c", b"3"),
393                ],
394                true,
395            )
396            .unwrap();
397
398        assert_eq!(store.get(b"a").unwrap().as_deref(), Some(&b"1"[..]));
399        assert_eq!(store.get(b"z").unwrap(), None);
400
401        let (k, v) = store.seek(b"a").unwrap().unwrap();
402        assert_eq!((k.as_slice(), v.as_slice()), (&b"a"[..], &b"1"[..]));
403        assert_eq!(store.seek(b"aa").unwrap().unwrap().0.as_slice(), b"b");
404        assert_eq!(store.seek_back(b"bz").unwrap().unwrap().0.as_slice(), b"b");
405        assert_eq!(
406            store.seek_back(b"\xff").unwrap().unwrap().0.as_slice(),
407            b"c"
408        );
409
410        store.commit(&[Op::Delete(b"b")], true).unwrap();
411        assert_eq!(store.get(b"b").unwrap(), None);
412        assert_eq!(store.seek(b"b").unwrap().unwrap().0.as_slice(), b"c");
413    }
414}
persistent_queue/store.rs

persistent_queue/
store.rs
