commonware_storage/kv/
batch.rs

1//! Support for batching changes to an underlying key-value store.
2
3use super::{Deletable, Gettable, Updatable};
4use crate::qmdb::{operation::Key, Error};
5use commonware_codec::CodecShared;
6use std::{collections::BTreeMap, future::Future};
7
8/// A batch of changes which may be written to an underlying store with [Batchable::write_batch].
9/// Writes and deletes to a batch are not applied to the store until the batch is written but
10/// will be reflected in reads from the batch.
11pub struct Batch<'a, K, V, D>
12where
13    K: Key,
14    V: CodecShared + Clone,
15    D: Gettable<Key = K, Value = V, Error = Error> + Sync,
16{
17    /// The underlying k/v store.
18    db: &'a D,
19    /// The diff of changes to the store.
20    ///
21    /// If the value is Some, the key is being created or updated.
22    /// If the value is None, the key is being deleted.
23    ///
24    /// We use a BTreeMap instead of HashMap to allow for a deterministic iteration order.
25    diff: BTreeMap<K, Option<V>>,
26}
27
28impl<'a, K, V, D> Batch<'a, K, V, D>
29where
30    K: Key,
31    V: CodecShared + Clone,
32    D: Gettable<Key = K, Value = V, Error = Error> + Sync,
33{
34    /// Returns a new batch of changes that may be written to the store.
35    pub const fn new(db: &'a D) -> Self {
36        Self {
37            db,
38            diff: BTreeMap::new(),
39        }
40    }
41
42    /// Deletes `key` from the batch without checking if it is present in the batch or store.
43    pub fn delete_unchecked(&mut self, key: K) -> Result<(), Error> {
44        self.diff.insert(key, None);
45
46        Ok(())
47    }
48}
49
50impl<'a, K, V, D> Gettable for Batch<'a, K, V, D>
51where
52    K: Key,
53    V: CodecShared + Clone,
54    D: Gettable<Key = K, Value = V, Error = Error> + Sync,
55{
56    type Key = K;
57    type Value = V;
58    type Error = Error;
59
60    /// Returns the value of `key` in the batch, or the value in the store if it is not present
61    /// in the batch.
62    async fn get(&self, key: &K) -> Result<Option<V>, Error> {
63        if let Some(value) = self.diff.get(key) {
64            return Ok(value.clone());
65        }
66
67        self.db.get(key).await
68    }
69}
70
71impl<'a, K, V, D> Updatable for Batch<'a, K, V, D>
72where
73    K: Key,
74    V: CodecShared + Clone,
75    D: Gettable<Key = K, Value = V, Error = Error> + Sync,
76{
77    /// Updates the value of `key` to `value` in the batch.
78    async fn update(&mut self, key: K, value: V) -> Result<(), Error> {
79        self.diff.insert(key, Some(value));
80
81        Ok(())
82    }
83}
84
85impl<'a, K, V, D> Deletable for Batch<'a, K, V, D>
86where
87    K: Key,
88    V: CodecShared + Clone,
89    D: Gettable<Key = K, Value = V, Error = Error> + Sync,
90{
91    /// Deletes `key` from the batch.
92    /// Returns true if the key was in the batch or store, false otherwise.
93    async fn delete(&mut self, key: K) -> Result<bool, Error> {
94        if let Some(entry) = self.diff.get_mut(&key) {
95            match entry {
96                Some(_) => {
97                    *entry = None;
98                    return Ok(true);
99                }
100                None => return Ok(false),
101            }
102        }
103
104        if self.db.get(&key).await?.is_some() {
105            self.diff.insert(key, None);
106            return Ok(true);
107        }
108
109        Ok(false)
110    }
111}
112
113impl<'a, K, V, D> IntoIterator for Batch<'a, K, V, D>
114where
115    K: Key,
116    V: CodecShared + Clone,
117    D: Gettable<Key = K, Value = V, Error = Error> + Sync,
118{
119    type Item = (K, Option<V>);
120    type IntoIter = std::collections::btree_map::IntoIter<K, Option<V>>;
121
122    fn into_iter(self) -> Self::IntoIter {
123        self.diff.into_iter()
124    }
125}
126
127/// A k/v store that supports making batched changes.
128pub trait Batchable: Gettable<Key: Key, Value: CodecShared + Clone, Error = Error> {
129    /// Returns a new empty batch of changes.
130    fn new_batch(&self) -> Batch<'_, Self::Key, Self::Value, Self>
131    where
132        Self: Sized + Sync,
133        Self::Value: Send + Sync,
134    {
135        Batch {
136            db: self,
137            diff: BTreeMap::new(),
138        }
139    }
140
141    /// Writes a batch of changes to the store.
142    fn write_batch<'a, Iter>(
143        &'a mut self,
144        iter: Iter,
145    ) -> impl Future<Output = Result<(), Error>> + Send + use<'a, Self, Iter>
146    where
147        Self: Send,
148        Iter: IntoIterator<Item = (Self::Key, Option<Self::Value>)> + Send + 'a,
149        Iter::IntoIter: Send;
150}
151
152#[cfg(test)]
153mod tests {
154    use super::*;
155    use crate::{
156        kv::tests::{assert_deletable, assert_gettable, assert_updatable},
157        qmdb::store::db::Db,
158        translator::TwoCap,
159    };
160    use commonware_cryptography::sha256::Digest;
161    use commonware_runtime::deterministic::Context;
162
163    type TestStore = Db<Context, Digest, Vec<u8>, TwoCap>;
164    type TestBatch<'a> = Batch<'a, Digest, Vec<u8>, TestStore>;
165
166    #[allow(dead_code)]
167    fn assert_batch_futures_are_send(batch: &mut TestBatch<'_>, key: Digest) {
168        assert_gettable(batch, &key);
169        assert_updatable(batch, key, vec![]);
170        assert_deletable(batch, key);
171    }
172}
commonware_storage/kv/batch.rs

commonware_storage/kv/
batch.rs
