Skip to main content

commonware_storage/qmdb/immutable/
batch.rs

1//! Batch mutation API for Immutable QMDBs.
2
3use super::Immutable;
4use crate::{
5    Context,
6    journal::{authenticated, contiguous::Mutable},
7    merkle::{Family, Location, Proof},
8    qmdb::{
9        Error,
10        any::{ValueEncoding, batch::lookup_sorted},
11        chain::{self, Bounds, Commitment, Compatible},
12        immutable::operation::Operation,
13        operation::Key,
14    },
15    translator::Translator,
16};
17use commonware_codec::EncodeShared;
18use commonware_cryptography::{Digest, Hasher};
19use commonware_parallel::Strategy;
20use commonware_utils::iter::zip_eq;
21use std::{
22    collections::BTreeMap,
23    iter,
24    sync::{Arc, Weak},
25};
26
27type DiffVec<K, F, V> = Vec<(K, DiffEntry<F, V>)>;
28
29/// What happened to a key in this batch.
30#[derive(Clone)]
31pub(crate) struct DiffEntry<F: Family, V> {
32    pub(crate) value: V,
33    pub(crate) loc: Location<F>,
34}
35
36/// Look up `key` in `diff`, ignoring an entry below the inactivity `floor`.
37///
38/// A reopened database indexes only operations at or above its floor, so batch reads skip
39/// entries below the floor to answer alike on every database at the same commitment.
40fn lookup_active<'a, K: Ord, F: Family, V>(
41    diff: &'a [(K, DiffEntry<F, V>)],
42    key: &K,
43    floor: Location<F>,
44) -> Option<&'a DiffEntry<F, V>> {
45    lookup_sorted(diff, key).filter(|entry| entry.loc >= floor)
46}
47
48/// Result of merkleizing a batch.
49type MerkleizeResult<F, D, K, V, S> = Result<Arc<MerkleizedBatch<F, D, K, V, S>>, Error<F>>;
50
51/// A speculative batch of operations whose root digest has not yet been computed, in contrast
52/// to [`MerkleizedBatch`].
53///
54/// Consuming [`UnmerkleizedBatch::merkleize`] produces a merkleized batch.
55/// Methods that need the DB (e.g. [`get`](Self::get)) accept it as a parameter.
56#[allow(clippy::type_complexity)]
57pub struct UnmerkleizedBatch<F, H, K, V, S: Strategy>
58where
59    F: Family,
60    K: Key,
61    V: ValueEncoding,
62    H: Hasher,
63{
64    /// Authenticated journal batch for computing the speculative Merkle root.
65    journal_batch: authenticated::UnmerkleizedBatch<F, H, Operation<F, K, V>, S>,
66
67    /// Pending mutations.
68    mutations: BTreeMap<K, V::Value>,
69
70    /// Parent batch in the chain. `None` for batches created directly from the DB.
71    parent: Option<Arc<MerkleizedBatch<F, H::Digest, K, V, S>>>,
72
73    /// The state immediately before this batch's operations.
74    /// This batch's i-th operation lands at location `base.size + i`.
75    base: Commitment<F, H::Digest>,
76}
77
78/// Merkleized authenticated-journal batch wrapping an [`Operation`] payload.
79type JournalBatch<F, D, K, V, S> = Arc<authenticated::MerkleizedBatch<F, D, Operation<F, K, V>, S>>;
80
81/// A speculative batch of operations whose root digest has been computed, in contrast to
82/// [`UnmerkleizedBatch`]. Reads through it return [`crate::qmdb::Error::StaleRead`] once the
83/// database is at no state of its chain, for example after a sibling or one of its descendants
84/// is applied (see [`crate::qmdb::chain`]).
85#[derive(Clone)]
86pub struct MerkleizedBatch<F: Family, D: Digest, K: Key, V: ValueEncoding, S: Strategy> {
87    /// Authenticated journal batch (Merkle state + local items).
88    pub(super) journal_batch: JournalBatch<F, D, K, V, S>,
89
90    /// This batch's local key-level changes only (not accumulated from ancestors).
91    /// Sorted by key with no duplicates; queried via `lookup_sorted` (binary search).
92    pub(super) diff: Arc<DiffVec<K, F, V::Value>>,
93
94    /// The parent batch in the chain, if any.
95    pub(super) parent: Option<Weak<Self>>,
96
97    /// Arc refs to the diffs of the reachable ancestors, collected during `merkleize()`. Reads
98    /// consult them (see `read_diffs`), and `apply_batch` applies the unapplied ones. 1:1 with
99    /// `bounds.ancestors` (same length, same ordering).
100    pub(super) ancestor_diffs: Vec<Arc<DiffVec<K, F, V::Value>>>,
101
102    /// Position and floor bounds for this batch chain.
103    pub(super) bounds: chain::Bounds<F, D>,
104}
105
106impl<F, H, K, V, S: Strategy> UnmerkleizedBatch<F, H, K, V, S>
107where
108    F: Family,
109    K: Key,
110    V: ValueEncoding,
111    H: Hasher,
112    Operation<F, K, V>: EncodeShared,
113{
114    /// Create a batch from the DB (no parent chain).
115    pub(super) fn new<E, C, T>(
116        immutable: &Immutable<F, E, K, V, C, H, T, S>,
117        base: Commitment<F, H::Digest>,
118    ) -> Self
119    where
120        E: Context,
121        C: Mutable<Item = Operation<F, K, V>>,
122        C::Item: EncodeShared,
123        T: Translator,
124    {
125        Self {
126            journal_batch: immutable.journal.new_batch(),
127            mutations: BTreeMap::new(),
128            parent: None,
129            base,
130        }
131    }
132
133    /// The database boundary for this batch chain.
134    ///
135    /// A batch created from the database uses its base. A child inherits its parent's `db`.
136    fn db(&self) -> Commitment<F, H::Digest> {
137        self.parent
138            .as_ref()
139            .map_or(self.base, |parent| parent.bounds.db)
140    }
141
142    /// The inactivity floor this batch builds on: its parent's, or `db`'s for a batch with no
143    /// parent.
144    fn floor<E, C, T>(&self, db: &Immutable<F, E, K, V, C, H, T, S>) -> Location<F>
145    where
146        E: Context,
147        C: Mutable<Item = Operation<F, K, V>>,
148        C::Item: EncodeShared,
149        T: Translator,
150    {
151        self.parent.as_ref().map_or_else(
152            || db.inactivity_floor_loc(),
153            |parent| parent.bounds.inactivity_floor,
154        )
155    }
156
157    /// Check that `db` is at a state of this batch's chain and return the [`Compatible`]
158    /// witness that database reads require.
159    #[allow(clippy::type_complexity)]
160    fn compatible<'a, E, C, T>(
161        &self,
162        db: &'a Immutable<F, E, K, V, C, H, T, S>,
163    ) -> Result<Compatible<'a, Immutable<F, E, K, V, C, H, T, S>>, Error<F>>
164    where
165        E: Context,
166        C: Mutable<Item = Operation<F, K, V>>,
167        C::Item: EncodeShared,
168        T: Translator,
169    {
170        self.parent.as_ref().map_or_else(
171            || self.base.compatible(db, db.commitment()),
172            |parent| parent.bounds.compatible(db, db.commitment()),
173        )
174    }
175
176    /// Set a key to a value.
177    ///
178    /// The key must not already be set in the database history or in any ancestor batch.
179    pub fn set(mut self, key: K, value: V::Value) -> Self {
180        self.mutations.insert(key, value);
181        self
182    }
183
184    /// Read through: mutations -> ancestor diffs -> database.
185    ///
186    /// Only operations at or above the inactivity floor this batch builds on are read, so every
187    /// database at the same commitment answers alike, however far it has pruned.
188    ///
189    /// # Errors
190    ///
191    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
192    pub async fn get<E, C, T>(
193        &self,
194        key: &K,
195        db: &Immutable<F, E, K, V, C, H, T, S>,
196    ) -> Result<Option<V::Value>, Error<F>>
197    where
198        E: Context,
199        C: Mutable<Item = Operation<F, K, V>>,
200        C::Item: EncodeShared,
201        T: Translator,
202    {
203        let db = self.compatible(db)?;
204        // Check this batch's pending mutations.
205        if let Some(value) = self.mutations.get(key) {
206            return Ok(Some(value.clone()));
207        }
208        let floor = self.floor(&*db);
209        if let Some(parent) = self.parent.as_ref() {
210            for diff in parent.read_diffs(db.bounds().end) {
211                if let Some(entry) = lookup_active(diff, key, floor) {
212                    return Ok(Some(entry.value.clone()));
213                }
214            }
215        }
216        // Fall through to base DB.
217        db.get_from(key, floor).await
218    }
219
220    /// Batch read multiple keys.
221    ///
222    /// Returns results in the same order as the input keys, matching [`Self::get`] for each key.
223    ///
224    /// # Errors
225    ///
226    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
227    pub async fn get_many<E, C, T>(
228        &self,
229        keys: &[&K],
230        db: &Immutable<F, E, K, V, C, H, T, S>,
231    ) -> Result<Vec<Option<V::Value>>, Error<F>>
232    where
233        E: Context,
234        C: Mutable<Item = Operation<F, K, V>>,
235        C::Item: EncodeShared,
236        T: Translator,
237    {
238        let db = self.compatible(db)?;
239        if keys.is_empty() {
240            return Ok(Vec::new());
241        }
242        let floor = self.floor(&*db);
243
244        let mut results: Vec<Option<V::Value>> = Vec::with_capacity(keys.len());
245        let mut db_indices = Vec::new();
246        let mut db_keys = Vec::new();
247
248        for (i, key) in keys.iter().enumerate() {
249            // Check local mutations.
250            if let Some(value) = self.mutations.get(*key) {
251                results.push(Some(value.clone()));
252                continue;
253            }
254
255            // Check the parent's retained diff chain.
256            if let Some(entry) = self.parent.as_ref().and_then(|parent| {
257                parent
258                    .read_diffs(db.bounds().end)
259                    .find_map(|diff| lookup_active(diff, *key, floor))
260            }) {
261                results.push(Some(entry.value.clone()));
262                continue;
263            }
264
265            // Need DB fallthrough.
266            db_indices.push(i);
267            db_keys.push(*key);
268            results.push(None);
269        }
270
271        if !db_keys.is_empty() {
272            let db_results = db.get_many_from(&db_keys, floor).await?;
273            for (slot, value) in zip_eq(db_indices, db_results) {
274                results[slot] = value;
275            }
276        }
277
278        Ok(results)
279    }
280
281    /// Resolve mutations into operations and return a merkleized batch.
282    ///
283    /// `inactivity_floor` declares that all operations before this location are inactive. It must
284    /// be at least the floor this batch builds on (its parent's, or the database's) and at most
285    /// this batch's own commit location.
286    ///
287    /// # Errors
288    ///
289    /// - Returns [`Error::StaleBatch`] if `db` is not at a state of this batch's chain or an
290    ///   unapplied ancestor was dropped.
291    /// - Returns [`Error::FloorRegressed`] if `inactivity_floor` is below the floor this batch
292    ///   builds on, and [`Error::FloorBeyondSize`] if it is past the commit location.
293    #[tracing::instrument(name = "qmdb.immutable.batch.merkleize", level = "info", skip_all)]
294    pub async fn merkleize<E, C, T>(
295        self,
296        db: &Immutable<F, E, K, V, C, H, T, S>,
297        metadata: Option<V::Value>,
298        inactivity_floor: Location<F>,
299    ) -> MerkleizeResult<F, H::Digest, K, V, S>
300    where
301        E: Context,
302        C: Mutable<Item = Operation<F, K, V>>,
303        C::Item: EncodeShared,
304        T: Translator,
305    {
306        let base = self.base.size;
307
308        let live_ancestors: Vec<_> =
309            chain::parent_and_ancestors(self.parent.as_ref(), |parent| parent.ancestors())
310                .collect();
311        let boundary = chain::effective_boundary(
312            self.db(),
313            live_ancestors.last().map(|oldest| oldest.bounds.base),
314        );
315
316        // Compute the batch chain bounds.
317        let mut ancestor_diffs = Vec::new();
318        let mut ancestors = Vec::new();
319        for batch in &live_ancestors {
320            ancestor_diffs.push(Arc::clone(&batch.diff));
321            ancestors.push(batch.commitment());
322        }
323        let db = chain::merkleizable(db, db.commitment(), boundary, ancestors.iter().copied())?;
324        let start_floor = self.floor(&*db);
325
326        // Build operations: one Set per key, then Commit. `self.mutations` is a BTreeMap, so
327        // iteration yields keys in sorted order, which `diff` relies on for binary search.
328        let mut ops: Vec<Operation<F, K, V>> = Vec::with_capacity(self.mutations.len() + 1);
329        let mut diff: DiffVec<K, F, V::Value> = Vec::with_capacity(self.mutations.len());
330
331        for (key, value) in self.mutations {
332            let loc = base + ops.len() as u64;
333            ops.push(Operation::Set(key.clone(), value.clone()));
334            diff.push((key, DiffEntry { value, loc }));
335        }
336
337        ops.push(Operation::Commit(metadata, inactivity_floor));
338
339        let total_size = base + ops.len() as u64;
340        chain::validate_merkleize_floor::<F>(start_floor, inactivity_floor, total_size - 1)?;
341        let inactive_peaks = F::inactive_peaks(total_size, inactivity_floor);
342
343        // Leaf and node hashing dominate merkleization, so run them as one job through the
344        // strategy (see `Journal::merkleize`).
345        let (journal, root) = db
346            .journal
347            .merkleize(self.journal_batch, ops, inactive_peaks)
348            .await?;
349
350        // Hold ancestor batches until the journal has captured their operations and nodes.
351        drop(live_ancestors);
352
353        Ok(Arc::new(MerkleizedBatch {
354            journal_batch: journal,
355            diff: Arc::new(diff),
356            parent: self.parent.as_ref().map(Arc::downgrade),
357            ancestor_diffs,
358            bounds: chain::Bounds {
359                base: self.base,
360                db: boundary,
361                tip: Commitment::new(total_size, root),
362                ancestors,
363                inactivity_floor,
364            },
365        }))
366    }
367}
368
369impl<F: Family, D: Digest, K: Key, V: ValueEncoding, S: Strategy> MerkleizedBatch<F, D, K, V, S>
370where
371    Operation<F, K, V>: EncodeShared,
372{
373    /// Return the speculative root.
374    pub const fn root(&self) -> D {
375        self.bounds.tip.root
376    }
377
378    /// Return the [`Bounds`] of the batch.
379    pub const fn bounds(&self) -> &Bounds<F, D> {
380        &self.bounds
381    }
382
383    /// Return the operations this batch appends to the log and the location of the first.
384    #[allow(clippy::type_complexity)]
385    pub fn operations(&self) -> (Location<F>, Arc<Vec<Operation<F, K, V>>>) {
386        (
387            self.bounds.base.size,
388            Arc::clone(self.journal_batch.items()),
389        )
390    }
391
392    /// Inclusion proof for the operations returned by [`Self::operations`], anchored at
393    /// this batch's tip. The pair verifies against [`Self::root`] via
394    /// [`crate::qmdb::verify_proof`]. Together with [`Self::pinned_nodes`] they verify via
395    /// [`crate::qmdb::verify_proof_and_pinned_nodes`].
396    ///
397    /// Nodes of unapplied ancestors are read through the chain, so those ancestors must still be
398    /// held. Nodes below the chain are read from `db`'s [Merkle store][crate::merkle::mem::Mem],
399    /// which retains them at least until this batch's changes are flushed (by a commit or sync
400    /// after apply).
401    ///
402    /// # Errors
403    ///
404    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain,
405    /// [`crate::merkle::Error::ElementPruned`] if a required node has been pruned or belongs to a
406    /// dropped unapplied ancestor, and [`crate::merkle::Error::Empty`] if the batch has no
407    /// operations (such as a batch from [`Immutable::to_batch`]).
408    pub fn proof<E, C, H, T>(
409        &self,
410        db: &Immutable<F, E, K, V, C, H, T, S>,
411    ) -> Result<Proof<F, D>, Error<F>>
412    where
413        E: Context,
414        C: Mutable<Item = Operation<F, K, V>>,
415        H: Hasher<Digest = D>,
416        T: Translator,
417    {
418        let db = self.bounds.compatible(db, db.commitment())?;
419        let inactive_peaks = F::inactive_peaks(self.bounds.tip.size, self.bounds.inactivity_floor);
420        db.journal
421            .speculative_proof(&self.journal_batch, inactive_peaks)
422            .map_err(Into::into)
423    }
424
425    /// The Merkle frontier at the first operation returned by [`Self::operations`]
426    /// ([`Family::nodes_to_pin`]), which lets a consumer holding only this batch's base rebuild
427    /// compact state and replay the operations. The operations, [`Self::proof`], and pinned
428    /// nodes verify against [`Self::root`] via [`crate::qmdb::verify_proof_and_pinned_nodes`].
429    ///
430    /// Nodes of unapplied ancestors are read through the chain, so those ancestors must still be
431    /// held. Nodes below the chain are read from `db`'s [Merkle store][crate::merkle::mem::Mem],
432    /// which retains them at least until this batch's changes are flushed (by a commit or sync
433    /// after apply).
434    ///
435    /// # Errors
436    ///
437    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain, and
438    /// [`crate::merkle::Error::ElementPruned`] if a required node has been pruned or belongs to a
439    /// dropped unapplied ancestor.
440    pub fn pinned_nodes<E, C, H, T>(
441        &self,
442        db: &Immutable<F, E, K, V, C, H, T, S>,
443    ) -> Result<Vec<D>, Error<F>>
444    where
445        E: Context,
446        C: Mutable<Item = Operation<F, K, V>>,
447        H: Hasher<Digest = D>,
448        T: Translator,
449    {
450        let db = self.bounds.compatible(db, db.commitment())?;
451        db.journal
452            .speculative_pinned_nodes(&self.journal_batch)
453            .map_err(Into::into)
454    }
455
456    /// Yield this batch's diff, then the ancestor diffs a read consults before a compatible
457    /// database at `db_size` operations: each reachable or unapplied ancestor's, newest first.
458    ///
459    /// The walk ends at the first ancestor that is both unreachable and applied. The database
460    /// holds that ancestor's writes and every older layer's. A reachable applied ancestor stays in
461    /// the walk so its writes resolve from memory. The batch keeps every retained diff, because a
462    /// bounded initialization can make those layers unapplied again.
463    fn read_diffs(
464        &self,
465        db_size: Location<F>,
466    ) -> impl Iterator<Item = &[(K, DiffEntry<F, V::Value>)]> {
467        let mut reachable = iter::once_with(|| self.ancestors()).flatten();
468        iter::once(self.diff.as_slice()).chain(
469            self.ancestor_diffs
470                .iter()
471                .zip(&self.bounds.ancestors)
472                .take_while(move |(_, state)| reachable.next().is_some() || state.size > db_size)
473                .map(|(diff, _)| diff.as_slice()),
474        )
475    }
476
477    /// Iterate over ancestor batches (parent first, then grandparent, etc.).
478    pub(super) fn ancestors(&self) -> impl Iterator<Item = Arc<Self>> + use<F, D, K, V, S> {
479        chain::ancestors(self.parent.clone(), |batch| batch.parent.as_ref())
480    }
481
482    /// The [`Commitment`] this batch commits to.
483    pub(super) const fn commitment(&self) -> Commitment<F, D> {
484        self.bounds.tip
485    }
486
487    /// Read through: local diff -> ancestor diffs -> database.
488    ///
489    /// Only operations at or above this batch's inactivity floor are read, so every database at
490    /// the same commitment answers alike, however far it has pruned.
491    ///
492    /// # Errors
493    ///
494    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
495    pub async fn get<E, C, H, T>(
496        &self,
497        key: &K,
498        db: &Immutable<F, E, K, V, C, H, T, S>,
499    ) -> Result<Option<V::Value>, Error<F>>
500    where
501        E: Context,
502        C: Mutable<Item = Operation<F, K, V>>,
503        C::Item: EncodeShared,
504        H: Hasher<Digest = D>,
505        T: Translator,
506    {
507        let db = self.bounds.compatible(db, db.commitment())?;
508        let floor = self.bounds.inactivity_floor;
509        for diff in self.read_diffs(db.bounds().end) {
510            if let Some(entry) = lookup_active(diff, key, floor) {
511                return Ok(Some(entry.value.clone()));
512            }
513        }
514        db.get_from(key, floor).await
515    }
516
517    /// Batch read multiple keys.
518    ///
519    /// Returns results in the same order as the input keys, matching [`Self::get`] for each key.
520    ///
521    /// # Errors
522    ///
523    /// Returns [`Error::StaleRead`] if `db` is not at a state of this batch's chain.
524    pub async fn get_many<E, C, H, T>(
525        &self,
526        keys: &[&K],
527        db: &Immutable<F, E, K, V, C, H, T, S>,
528    ) -> Result<Vec<Option<V::Value>>, Error<F>>
529    where
530        E: Context,
531        C: Mutable<Item = Operation<F, K, V>>,
532        C::Item: EncodeShared,
533        H: Hasher<Digest = D>,
534        T: Translator,
535    {
536        let db = self.bounds.compatible(db, db.commitment())?;
537        if keys.is_empty() {
538            return Ok(Vec::new());
539        }
540        let floor = self.bounds.inactivity_floor;
541
542        let mut results: Vec<Option<V::Value>> = Vec::with_capacity(keys.len());
543        let mut db_indices = Vec::new();
544        let mut db_keys = Vec::new();
545
546        for (i, key) in keys.iter().enumerate() {
547            if let Some(entry) = self
548                .read_diffs(db.bounds().end)
549                .find_map(|diff| lookup_active(diff, *key, floor))
550            {
551                results.push(Some(entry.value.clone()));
552                continue;
553            }
554
555            // Need DB fallthrough.
556            db_indices.push(i);
557            db_keys.push(*key);
558            results.push(None);
559        }
560
561        if !db_keys.is_empty() {
562            let db_results = db.get_many_from(&db_keys, floor).await?;
563            for (slot, value) in zip_eq(db_indices, db_results) {
564                results[slot] = value;
565            }
566        }
567
568        Ok(results)
569    }
570
571    /// Create a new speculative batch of operations with this batch as its parent.
572    ///
573    /// Every unapplied ancestor in the chain must be held until the child (or any descendant) is
574    /// merkleized. Otherwise, `merkleize` returns [`Error::StaleBatch`].
575    ///
576    /// A child can be created from a stale parent. While the database is at none of the chain's
577    /// states, the child's reads return [`Error::StaleRead`], and merkleizing it returns
578    /// [`Error::StaleBatch`].
579    pub fn new_batch<H>(self: &Arc<Self>) -> UnmerkleizedBatch<F, H, K, V, S>
580    where
581        H: Hasher<Digest = D>,
582    {
583        UnmerkleizedBatch {
584            journal_batch: self.journal_batch.new_batch::<H>(),
585            mutations: BTreeMap::new(),
586            parent: Some(Arc::clone(self)),
587            base: self.commitment(),
588        }
589    }
590}
591
592impl<F, E, K, V, C, H, T, S> Immutable<F, E, K, V, C, H, T, S>
593where
594    F: Family,
595    E: Context,
596    K: Key,
597    V: ValueEncoding,
598    C: Mutable<Item = Operation<F, K, V>>,
599    C::Item: EncodeShared,
600    H: Hasher,
601    T: Translator,
602    S: Strategy,
603{
604    /// Create an initial [`MerkleizedBatch`] from the DB's current state.
605    pub fn to_batch(&self) -> Arc<MerkleizedBatch<F, H::Digest, K, V, S>> {
606        Arc::new(MerkleizedBatch {
607            journal_batch: self.journal.to_merkleized_batch(),
608            diff: Arc::new(Vec::new()),
609            parent: None,
610            ancestor_diffs: Vec::new(),
611            bounds: chain::Bounds::from_db(self.commitment(), self.inactivity_floor_loc),
612        })
613    }
614}