cnidarium/snapshot/
rocks_wrapper.rs

1use std::fmt::{Debug, Formatter};
2use std::ops::Deref;
3use std::sync::Arc;
4
5/// A wrapper type that acts as a `rocksdb::Snapshot` of an `Arc`'d database
6/// handle.
7///
8/// This works around a limitation of the `rocksdb` API: the `rocksdb::Snapshot`
9/// can only take a borrowed database handle, not an `Arc`'d one, so the
10/// lifetime of the `rocksdb::Snapshot` is bound to the lifetime of the borrowed
11/// handle.  Instead, this wrapper type bundles an `Arc`'d handle together with
12/// the `rocksdb::Snapshot`, so that the database is guaranteed to live at least
13/// as long as any snapshot of it.
14pub struct RocksDbSnapshot {
15    /// The snapshot itself.  It's not really `'static`, so it's on us to ensure
16    /// that the database stays live as long as the snapshot does.
17    inner: rocksdb::Snapshot<'static>,
18    /// The raw pointer form of the Arc<DB> we use to guarantee the database
19    /// lives at least as long as the snapshot.  We create this from the Arc<DB>
20    /// in the constructor, pass it to the snapshot on creation, and then
21    /// convert it back into an Arc in the drop impl to decrement the refcount.
22    ///
23    /// Arc::into_raw consumes the Arc instance but does not decrement the
24    /// refcount.  This means that we cannot accidentally drop the Arc while
25    /// using the raw pointer.  Instead, we must explicitly convert the raw
26    /// pointer back into an Arc when we're finished using it, and only then
27    /// drop it.
28    raw_db: *const rocksdb::DB,
29}
30
31impl Debug for RocksDbSnapshot {
32    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
33        f.debug_struct("RocksDbSnapshot").finish()
34    }
35}
36
37// Safety requires that the inner snapshot instance must never live longer than
38// the wrapper.  We're assured that this is the case, because we only return a
39// borrow of the inner snapshot, and because `rocksdb::Snapshot` is neither
40// `Copy` nor `Clone`.
41//
42// We're also reasonably certain that the upstream crate will not add such an
43// implementation in the future, because its drop impl is used to make the FFI
44// call that discards the in-memory snapshot, so it would not be safe to add
45// such an implementation.
46impl Deref for RocksDbSnapshot {
47    type Target = rocksdb::Snapshot<'static>;
48
49    fn deref(&self) -> &Self::Target {
50        &self.inner
51    }
52}
53
54impl RocksDbSnapshot {
55    /// Creates a new snapshot of the given `db`.
56    pub fn new(db: Arc<rocksdb::DB>) -> Self {
57        // First, convert the Arc<DB> into a raw pointer.
58        let raw_db = Arc::into_raw(db);
59        // Next, use the raw pointer to construct a &DB instance with a fake
60        // 'static lifetime, and use that instance to construct the inner
61        // Snapshot.
62        let static_db: &'static rocksdb::DB = unsafe { &*raw_db };
63        let inner = rocksdb::Snapshot::new(static_db);
64
65        Self { inner, raw_db }
66    }
67}
68
69impl Drop for RocksDbSnapshot {
70    fn drop(&mut self) {
71        // Now that we know we're finished with the `Snapshot`, we can
72        // reconstruct the `Arc` and drop it, to decrement the DB refcount.
73        unsafe {
74            let db = Arc::from_raw(self.raw_db);
75            std::mem::drop(db);
76        }
77    }
78}
79
80/// The `Send` implementation is safe because the `rocksdb::Snapshot` is `Send`.
81unsafe impl Send for RocksDbSnapshot {}
82/// The `Sync` implementation is safe because the `rocksdb::Snapshot` is `Sync`.
83unsafe impl Sync for RocksDbSnapshot {}