zebra_state/service/finalized_state/
column_family.rs

1//! Type-safe column family access.
2
3// When these types aren't exported, they become dead code.
4#![cfg_attr(not(any(test, feature = "proptest-impl")), allow(dead_code))]
5
6use std::{
7    any::type_name,
8    collections::{BTreeMap, HashMap},
9    fmt::Debug,
10    hash::Hash,
11    marker::PhantomData,
12    ops::RangeBounds,
13};
14
15use crate::service::finalized_state::{DiskWriteBatch, FromDisk, IntoDisk, ReadDisk, WriteDisk};
16
17use super::DiskDb;
18
19/// A type-safe read-only column family reference.
20///
21/// Use this struct instead of raw [`ReadDisk`] access, because it is type-safe.
22/// So you only have to define the types once, and you can't accidentally use different types for
23/// reading and writing. (Which is a source of subtle database bugs.)
24#[derive(Clone)]
25pub struct TypedColumnFamily<'cf, Key, Value>
26where
27    Key: IntoDisk + FromDisk + Debug,
28    Value: IntoDisk + FromDisk,
29{
30    /// The database.
31    db: DiskDb,
32
33    /// The column family reference in the database.
34    cf: rocksdb::ColumnFamilyRef<'cf>,
35
36    /// The column family name, only used for debugging and equality checking.
37    _cf_name: String,
38
39    /// A marker type used to bind the key and value types to the struct.
40    _marker: PhantomData<(Key, Value)>,
41}
42
43/// A type-safe and drop-safe batch write to a column family.
44///
45/// Use this struct instead of raw [`WriteDisk`] access, because it is type-safe.
46/// So you only have to define the types once, and you can't accidentally use different types for
47/// reading and writing. (Which is a source of subtle database bugs.)
48///
49/// This type is also drop-safe: unwritten batches have to be specifically ignored.
50#[must_use = "batches must be written to the database"]
51#[derive(Debug, Eq, PartialEq)]
52pub struct WriteTypedBatch<'cf, Key, Value, Batch>
53where
54    Key: IntoDisk + FromDisk + Debug,
55    Value: IntoDisk + FromDisk,
56    Batch: WriteDisk,
57{
58    inner: TypedColumnFamily<'cf, Key, Value>,
59
60    batch: Batch,
61}
62
63impl<Key, Value> Debug for TypedColumnFamily<'_, Key, Value>
64where
65    Key: IntoDisk + FromDisk + Debug,
66    Value: IntoDisk + FromDisk,
67{
68    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69        f.debug_struct(&format!(
70            "TypedColumnFamily<{}, {}>",
71            type_name::<Key>(),
72            type_name::<Value>()
73        ))
74        .field("db", &self.db)
75        .field("cf", &self._cf_name)
76        .finish()
77    }
78}
79
80impl<Key, Value> PartialEq for TypedColumnFamily<'_, Key, Value>
81where
82    Key: IntoDisk + FromDisk + Debug,
83    Value: IntoDisk + FromDisk,
84{
85    fn eq(&self, other: &Self) -> bool {
86        self.db == other.db && self._cf_name == other._cf_name
87    }
88}
89
90impl<Key, Value> Eq for TypedColumnFamily<'_, Key, Value>
91where
92    Key: IntoDisk + FromDisk + Debug,
93    Value: IntoDisk + FromDisk,
94{
95}
96
97impl<'cf, Key, Value> TypedColumnFamily<'cf, Key, Value>
98where
99    Key: IntoDisk + FromDisk + Debug,
100    Value: IntoDisk + FromDisk,
101{
102    // Creation
103
104    /// Returns a new typed column family, if it exists in the database.
105    pub fn new(db: &'cf DiskDb, cf_name: &str) -> Option<Self> {
106        let cf = db.cf_handle(cf_name)?;
107
108        Some(Self {
109            db: db.clone(),
110            cf,
111            _cf_name: cf_name.to_string(),
112            _marker: PhantomData,
113        })
114    }
115
116    // Writing
117
118    /// Returns a typed writer for this column family for a new batch.
119    ///
120    /// These methods are the only way to get a `WriteTypedBatch`, which ensures
121    /// that the read and write types are consistent.
122    pub fn new_batch_for_writing(self) -> WriteTypedBatch<'cf, Key, Value, DiskWriteBatch> {
123        WriteTypedBatch {
124            inner: self,
125            batch: DiskWriteBatch::new(),
126        }
127    }
128
129    /// Wraps an existing write batch, and returns a typed writer for this column family.
130    ///
131    /// These methods are the only way to get a `WriteTypedBatch`, which ensures
132    /// that the read and write types are consistent.
133    pub fn take_batch_for_writing(
134        self,
135        batch: DiskWriteBatch,
136    ) -> WriteTypedBatch<'cf, Key, Value, DiskWriteBatch> {
137        WriteTypedBatch { inner: self, batch }
138    }
139
140    /// Wraps an existing write batch reference, and returns a typed writer for this column family.
141    ///
142    /// These methods are the only way to get a `WriteTypedBatch`, which ensures
143    /// that the read and write types are consistent.
144    pub fn with_batch_for_writing(
145        self,
146        batch: &mut DiskWriteBatch,
147    ) -> WriteTypedBatch<'cf, Key, Value, &mut DiskWriteBatch> {
148        WriteTypedBatch { inner: self, batch }
149    }
150
151    // Reading
152
153    /// Returns true if this rocksdb column family does not contain any entries.
154    pub fn zs_is_empty(&self) -> bool {
155        self.db.zs_is_empty(&self.cf)
156    }
157
158    /// Returns the value for `key` in this rocksdb column family, if present.
159    pub fn zs_get(&self, key: &Key) -> Option<Value> {
160        self.db.zs_get(&self.cf, key)
161    }
162
163    /// Check if this rocksdb column family contains the serialized form of `key`.
164    pub fn zs_contains(&self, key: &Key) -> bool {
165        self.db.zs_contains(&self.cf, key)
166    }
167
168    /// Returns the lowest key in this column family, and the corresponding value.
169    ///
170    /// Returns `None` if this column family is empty.
171    pub fn zs_first_key_value(&self) -> Option<(Key, Value)> {
172        self.db.zs_first_key_value(&self.cf)
173    }
174
175    /// Returns the highest key in this column family, and the corresponding value.
176    ///
177    /// Returns `None` if this column family is empty.
178    pub fn zs_last_key_value(&self) -> Option<(Key, Value)> {
179        self.db.zs_last_key_value(&self.cf)
180    }
181
182    /// Returns the first key greater than or equal to `lower_bound` in this column family,
183    /// and the corresponding value.
184    ///
185    /// Returns `None` if there are no keys greater than or equal to `lower_bound`.
186    pub fn zs_next_key_value_from(&self, lower_bound: &Key) -> Option<(Key, Value)> {
187        self.db.zs_next_key_value_from(&self.cf, lower_bound)
188    }
189
190    /// Returns the first key strictly greater than `lower_bound` in this column family,
191    /// and the corresponding value.
192    ///
193    /// Returns `None` if there are no keys greater than `lower_bound`.
194    pub fn zs_next_key_value_strictly_after(&self, lower_bound: &Key) -> Option<(Key, Value)> {
195        self.db
196            .zs_next_key_value_strictly_after(&self.cf, lower_bound)
197    }
198
199    /// Returns the first key less than or equal to `upper_bound` in this column family,
200    /// and the corresponding value.
201    ///
202    /// Returns `None` if there are no keys less than or equal to `upper_bound`.
203    pub fn zs_prev_key_value_back_from(&self, upper_bound: &Key) -> Option<(Key, Value)> {
204        self.db.zs_prev_key_value_back_from(&self.cf, upper_bound)
205    }
206
207    /// Returns the first key strictly less than `upper_bound` in this column family,
208    /// and the corresponding value.
209    ///
210    /// Returns `None` if there are no keys less than `upper_bound`.
211    pub fn zs_prev_key_value_strictly_before(&self, upper_bound: &Key) -> Option<(Key, Value)> {
212        self.db
213            .zs_prev_key_value_strictly_before(&self.cf, upper_bound)
214    }
215
216    /// Returns a forward iterator over the items in this column family in `range`.
217    ///
218    /// Holding this iterator open might delay block commit transactions.
219    pub fn zs_forward_range_iter<Range>(
220        &self,
221        range: Range,
222    ) -> impl Iterator<Item = (Key, Value)> + '_
223    where
224        Range: RangeBounds<Key>,
225    {
226        self.db.zs_forward_range_iter(&self.cf, range)
227    }
228
229    /// Returns a reverse iterator over the items in this column family in `range`.
230    ///
231    /// Holding this iterator open might delay block commit transactions.
232    pub fn zs_reverse_range_iter<Range>(
233        &self,
234        range: Range,
235    ) -> impl Iterator<Item = (Key, Value)> + '_
236    where
237        Range: RangeBounds<Key>,
238    {
239        self.db.zs_reverse_range_iter(&self.cf, range)
240    }
241}
242
243impl<Key, Value> TypedColumnFamily<'_, Key, Value>
244where
245    Key: IntoDisk + FromDisk + Debug + Ord,
246    Value: IntoDisk + FromDisk,
247{
248    /// Returns the keys and values in this column family in `range`, in an ordered `BTreeMap`.
249    ///
250    /// Holding this iterator open might delay block commit transactions.
251    pub fn zs_items_in_range_ordered<Range>(&self, range: Range) -> BTreeMap<Key, Value>
252    where
253        Range: RangeBounds<Key>,
254    {
255        self.db.zs_items_in_range_ordered(&self.cf, range)
256    }
257}
258
259impl<Key, Value> TypedColumnFamily<'_, Key, Value>
260where
261    Key: IntoDisk + FromDisk + Debug + Hash + Eq,
262    Value: IntoDisk + FromDisk,
263{
264    /// Returns the keys and values in this column family in `range`, in an unordered `HashMap`.
265    ///
266    /// Holding this iterator open might delay block commit transactions.
267    pub fn zs_items_in_range_unordered<Range>(&self, range: Range) -> HashMap<Key, Value>
268    where
269        Range: RangeBounds<Key>,
270    {
271        self.db.zs_items_in_range_unordered(&self.cf, range)
272    }
273}
274
275impl<Key, Value, Batch> WriteTypedBatch<'_, Key, Value, Batch>
276where
277    Key: IntoDisk + FromDisk + Debug,
278    Value: IntoDisk + FromDisk,
279    Batch: WriteDisk,
280{
281    // Batching before writing
282
283    /// Serialize and insert the given key and value into this column family,
284    /// overwriting any existing `value` for `key`.
285    pub fn zs_insert(mut self, key: &Key, value: &Value) -> Self {
286        self.batch.zs_insert(&self.inner.cf, key, value);
287
288        self
289    }
290
291    /// Remove the given key from this column family, if it exists.
292    pub fn zs_delete(mut self, key: &Key) -> Self {
293        self.batch.zs_delete(&self.inner.cf, key);
294
295        self
296    }
297
298    /// Delete the given key range from this rocksdb column family, if it exists, including `from`
299    /// and excluding `until_strictly_before`.
300    //.
301    // TODO: convert zs_delete_range() to take std::ops::RangeBounds
302    //       see zs_range_iter() for an example of the edge cases
303    pub fn zs_delete_range(mut self, from: &Key, until_strictly_before: &Key) -> Self {
304        self.batch
305            .zs_delete_range(&self.inner.cf, from, until_strictly_before);
306
307        self
308    }
309}
310
311// Writing a batch to the database requires an owned batch.
312impl<Key, Value> WriteTypedBatch<'_, Key, Value, DiskWriteBatch>
313where
314    Key: IntoDisk + FromDisk + Debug,
315    Value: IntoDisk + FromDisk,
316{
317    // Writing batches
318
319    /// Writes this batch to this column family in the database,
320    /// taking ownership and consuming it.
321    pub fn write_batch(self) -> Result<(), rocksdb::Error> {
322        self.inner.db.write(self.batch)
323    }
324}