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