kaspa_database/
writer.rs

1use kaspa_utils::refs::Refs;
2use rocksdb::WriteBatch;
3
4use crate::prelude::DB;
5
6/// Abstraction over direct/batched DB writing
7pub trait DbWriter {
8    fn put<K, V>(&mut self, key: K, value: V) -> Result<(), rocksdb::Error>
9    where
10        K: AsRef<[u8]>,
11        V: AsRef<[u8]>;
12    fn delete<K: AsRef<[u8]>>(&mut self, key: K) -> Result<(), rocksdb::Error>;
13    fn delete_range<K>(&mut self, from: K, to: K) -> Result<(), rocksdb::Error>
14    where
15        K: AsRef<[u8]>;
16}
17
18/// A trait which is intentionally not implemented for the batch writer.
19/// Aimed for compile-time safety of operations which do not support batch writing semantics
20pub trait DirectWriter: DbWriter {}
21
22pub struct DirectDbWriter<'a> {
23    db: Refs<'a, DB>,
24}
25
26impl<'a> DirectDbWriter<'a> {
27    pub fn new(db: &'a DB) -> Self {
28        Self { db: db.into() }
29    }
30
31    pub fn from_arc(db: std::sync::Arc<DB>) -> Self {
32        Self { db: db.into() }
33    }
34}
35
36impl DbWriter for DirectDbWriter<'_> {
37    fn put<K, V>(&mut self, key: K, value: V) -> Result<(), rocksdb::Error>
38    where
39        K: AsRef<[u8]>,
40        V: AsRef<[u8]>,
41    {
42        self.db.put(key, value)
43    }
44
45    fn delete<K: AsRef<[u8]>>(&mut self, key: K) -> Result<(), rocksdb::Error> {
46        self.db.delete(key)
47    }
48
49    fn delete_range<K>(&mut self, from: K, to: K) -> Result<(), rocksdb::Error>
50    where
51        K: AsRef<[u8]>,
52    {
53        let mut batch = WriteBatch::default();
54        batch.delete_range(from, to);
55        self.db.write(batch)
56    }
57}
58
59impl DirectWriter for DirectDbWriter<'_> {}
60
61pub struct BatchDbWriter<'a> {
62    batch: &'a mut WriteBatch,
63}
64
65impl<'a> BatchDbWriter<'a> {
66    pub fn new(batch: &'a mut WriteBatch) -> Self {
67        Self { batch }
68    }
69}
70
71impl DbWriter for BatchDbWriter<'_> {
72    fn put<K, V>(&mut self, key: K, value: V) -> Result<(), rocksdb::Error>
73    where
74        K: AsRef<[u8]>,
75        V: AsRef<[u8]>,
76    {
77        self.batch.put(key, value);
78        Ok(())
79    }
80
81    fn delete<K: AsRef<[u8]>>(&mut self, key: K) -> Result<(), rocksdb::Error> {
82        self.batch.delete(key);
83        Ok(())
84    }
85
86    fn delete_range<K>(&mut self, from: K, to: K) -> Result<(), rocksdb::Error>
87    where
88        K: AsRef<[u8]>,
89    {
90        self.batch.delete_range(from, to);
91        Ok(())
92    }
93}
94
95impl<T: DbWriter> DbWriter for &mut T {
96    #[inline]
97    fn put<K, V>(&mut self, key: K, value: V) -> Result<(), rocksdb::Error>
98    where
99        K: AsRef<[u8]>,
100        V: AsRef<[u8]>,
101    {
102        (*self).put(key, value)
103    }
104
105    #[inline]
106    fn delete<K: AsRef<[u8]>>(&mut self, key: K) -> Result<(), rocksdb::Error> {
107        (*self).delete(key)
108    }
109
110    #[inline]
111    fn delete_range<K>(&mut self, from: K, to: K) -> Result<(), rocksdb::Error>
112    where
113        K: AsRef<[u8]>,
114    {
115        (*self).delete_range(from, to)
116    }
117}
118
119impl<T: DirectWriter> DirectWriter for &mut T {}
120
121/// A writer for memory stores which writes nothing to the DB
122#[derive(Default)]
123pub struct MemoryWriter;
124
125impl DbWriter for MemoryWriter {
126    fn put<K, V>(&mut self, _key: K, _value: V) -> Result<(), rocksdb::Error>
127    where
128        K: AsRef<[u8]>,
129        V: AsRef<[u8]>,
130    {
131        Ok(())
132    }
133
134    fn delete<K: AsRef<[u8]>>(&mut self, _key: K) -> Result<(), rocksdb::Error> {
135        Ok(())
136    }
137
138    fn delete_range<K>(&mut self, _from: K, _to: K) -> Result<(), rocksdb::Error>
139    where
140        K: AsRef<[u8]>,
141    {
142        Ok(())
143    }
144}
145
146impl DirectWriter for MemoryWriter {}
kaspa_database/writer.rs

kaspa_database/
writer.rs
