use crate::{
index::Unordered as UnorderedIndex,
journal::{
contiguous::{Contiguous, Mutable},
Error as JournalError,
},
merkle::{
self, hasher::Hasher as _, mem::Mem, storage::Storage as MerkleStorage, Graftable,
Location, Position,
},
metadata::{Config as MConfig, Metadata},
qmdb::{
self,
any::{
self,
operation::{update::Update, Operation},
},
current::{
batch::BitmapBatch,
grafting,
proof::{OperationProof, OpsRootWitness, RangeProof, RangeProofSpec},
},
operation::Operation as _,
Error,
},
Context,
};
use commonware_codec::{Codec, CodecShared, DecodeExt};
use commonware_cryptography::{Digest, DigestOf, Hasher};
use commonware_macros::boxed;
use commonware_parallel::Strategy;
use commonware_runtime::telemetry::metrics::{
histogram::{ScopedTimer, Timed},
Counter, Gauge, GaugeExt as _, MetricsExt as _,
};
use commonware_utils::{
bitmap::{self, Readable as _},
sequence::prefixed_u64::U64,
};
use core::{num::NonZeroU64, ops::Range};
use futures::future::try_join_all;
use std::{collections::BTreeMap, sync::Arc};
use tracing::{error, warn};
const NODE_PREFIX: u8 = 0;
const PRUNED_CHUNKS_PREFIX: u8 = 1;
pub(crate) struct Metrics<E: Context> {
pruned_chunks: Gauge,
sync_boundary: Gauge,
pub apply_batch_calls: Counter,
apply_batch_duration: Timed,
pub sync_calls: Counter,
sync_duration: Timed,
pub prune_calls: Counter,
prune_duration: Timed,
clock: Arc<E>,
}
impl<E: Context> Metrics<E> {
pub fn new(context: E) -> Self {
Self {
pruned_chunks: context.gauge("pruned_chunks", "Number of pruned bitmap chunks"),
sync_boundary: context
.gauge("sync_boundary", "Most recent safe sync boundary location"),
apply_batch_calls: context.counter("apply_batch_calls", "Number of apply-batch calls"),
apply_batch_duration: Timed::register(
&context,
"apply_batch_duration",
"Duration of apply-batch calls",
),
sync_calls: context.counter("sync_calls", "Number of sync calls"),
sync_duration: Timed::register(&context, "sync_duration", "Duration of sync calls"),
prune_calls: context.counter("prune_calls", "Number of prune calls"),
prune_duration: Timed::register(&context, "prune_duration", "Duration of prune calls"),
clock: Arc::new(context),
}
}
pub fn apply_batch_timer(&self) -> ScopedTimer<E> {
self.apply_batch_duration.scoped(&self.clock)
}
pub fn sync_timer(&self) -> ScopedTimer<E> {
self.sync_duration.scoped(&self.clock)
}
pub fn prune_timer(&self) -> ScopedTimer<E> {
self.prune_duration.scoped(&self.clock)
}
pub fn update(&self, pruned_chunks: u64, sync_boundary: u64) {
let _ = self.pruned_chunks.try_set(pruned_chunks);
let _ = self.sync_boundary.try_set(sync_boundary);
}
}
pub struct Db<
F: merkle::Graftable,
E: Context,
C: Contiguous<Item: CodecShared>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
U: Send + Sync,
const N: usize,
S: Strategy,
> {
pub(super) any: any::db::Db<F, E, C, I, H, U, N, S>,
pub(super) grafted_tree: Arc<Mem<F, H::Digest>>,
pub(super) metadata: Metadata<E, U64, Vec<u8>>,
pub(super) strategy: S,
pub(super) root: DigestOf<H>,
pub(super) metrics: Metrics<E>,
#[cfg(test)]
pub(super) halt_before_prune_log: bool,
}
impl<F, E, C, I, H, U, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
U: Update,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
S: Strategy,
Operation<F, U>: Codec,
{
#[cfg(any(test, feature = "test-traits"))]
pub(crate) const fn inactivity_floor_loc(&self) -> Location<F> {
self.any.inactivity_floor_loc()
}
pub const fn is_empty(&self) -> bool {
self.any.is_empty()
}
pub async fn get_metadata(&self) -> Result<Option<U::Value>, Error<F>> {
self.any.get_metadata().await
}
pub fn bounds(&self) -> std::ops::Range<Location<F>> {
self.any.bounds()
}
pub fn verify_range_proof(
proof: &RangeProof<F, H::Digest>,
start_loc: Location<F>,
ops: &[Operation<F, U>],
chunks: &[[u8; N]],
root: &H::Digest,
) -> bool {
proof.verify::<H, _, N>(start_loc, ops, chunks, root)
}
}
impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
U: Update,
C: Contiguous<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
S: Strategy,
Operation<F, U>: Codec,
{
fn grafted_storage(&self) -> impl MerkleStorage<F, Digest = H::Digest> + '_ {
grafting::Storage::<F, H, _, _>::new(
&self.grafted_tree,
grafting::height::<N>(),
&self.any.log.merkle,
)
}
pub const fn root(&self) -> H::Digest {
self.root
}
pub const fn strategy(&self) -> &S {
&self.strategy
}
pub const fn ops_root(&self) -> H::Digest {
self.any.root()
}
pub async fn ops_root_witness(&self) -> Result<OpsRootWitness<F, H::Digest>, Error<F>> {
let storage = self.grafted_storage();
let ops_size = storage.size();
let ops_leaves = Location::<F>::try_from(ops_size)?;
let grafted_root = compute_grafted_root::<F, H, _, _, N>(
self.any.bitmap.as_ref(),
&storage,
ops_leaves,
self.any.inactivity_floor_loc,
)
.await?;
let hasher = qmdb::hasher::<H>();
let partial_chunk = partial_chunk::<_, N>(self.any.bitmap.as_ref())
.map(|(chunk, next_bit)| (next_bit, hasher.digest(chunk.as_slice())));
let pending_chunk_digest: F::PendingChunk<H::Digest> = pending_chunk::<F, _, N>(
self.any.bitmap.as_ref(),
ops_leaves,
grafting::height::<N>(),
)?
.map(|chunk| hasher.digest(chunk.as_slice()))
.try_into()
.expect("pending_chunk must be consistent with family");
Ok(OpsRootWitness {
grafted_root,
pending_chunk_digest,
partial_chunk,
})
}
pub(super) fn grafted_snapshot(&self) -> Arc<merkle::batch::MerkleizedBatch<F, H::Digest, S>> {
merkle::batch::MerkleizedBatch::from_mem_with_strategy(
&self.grafted_tree,
self.strategy.clone(),
)
}
pub fn new_batch(&self) -> super::batch::UnmerkleizedBatch<F, H, U, N, S> {
super::batch::UnmerkleizedBatch::new(
self.any.new_batch(),
self.grafted_snapshot(),
BitmapBatch::Base(Arc::clone(&self.any.bitmap)),
)
}
pub(super) async fn operation_proof(
&self,
loc: Location<F>,
) -> Result<OperationProof<F, H::Digest, N>, Error<F>> {
let storage = self.grafted_storage();
let ops_root = self.any.root();
OperationProof::new::<H, _>(
self.any.bitmap.as_ref(),
&storage,
self.any.inactivity_floor_loc,
loc,
ops_root,
)
.await
}
#[allow(clippy::type_complexity)]
#[tracing::instrument(
name = "qmdb.current.db.range_proof",
level = "info",
skip_all,
fields(
start_loc = *start_loc,
max_ops = max_ops.get(),
),
)]
pub async fn range_proof(
&self,
start_loc: Location<F>,
max_ops: NonZeroU64,
) -> Result<(RangeProof<F, H::Digest>, Vec<Operation<F, U>>, Vec<[u8; N]>), Error<F>> {
let storage = self.grafted_storage();
let ops_root = self.any.root();
RangeProof::new_with_ops::<H, _, _, N>(
self.any.bitmap.as_ref(),
&storage,
&self.any.log,
RangeProofSpec {
start_loc,
max_ops,
inactivity_floor: self.any.inactivity_floor_loc,
ops_root,
},
)
.await
}
}
impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
U: Update,
C: Mutable<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
S: Strategy,
Operation<F, U>: Codec,
{
pub async fn ops_historical_proof(
&self,
historical_size: Location<F>,
start_loc: Location<F>,
max_ops: NonZeroU64,
) -> Result<(merkle::Proof<F, H::Digest>, Vec<Operation<F, U>>), Error<F>> {
self.any
.historical_proof(historical_size, start_loc, max_ops)
.await
}
pub async fn pinned_nodes_at(&self, loc: Location<F>) -> Result<Vec<H::Digest>, Error<F>> {
self.any.pinned_nodes_at(loc).await
}
pub fn sync_boundary(&self) -> Location<F> {
sync_boundary::<F, N>(
*self.any.inactivity_floor_loc / bitmap::Prunable::<N>::CHUNK_SIZE_BITS,
*self.any.last_commit_loc + 1,
)
}
pub(super) fn update_metrics(&self) {
self.metrics.update(
self.any.bitmap.pruned_chunks() as u64,
*self.sync_boundary(),
);
}
fn delayed_merge_rewind_floor(&self) -> Option<u64> {
pair_absorption_threshold::<F, N>(self.any.bitmap.pruned_chunks() as u64)
}
fn prune_grafted_tree_to_bitmap(&mut self) -> Result<(), Error<F>> {
let pruned_chunks = self.any.bitmap.pruned_chunks() as u64;
if pruned_chunks == 0 {
return Ok(());
}
let prune_loc = Location::<F>::new(pruned_chunks);
if prune_loc <= self.grafted_tree.bounds().start {
return Ok(());
}
let prune_pos = Position::try_from(prune_loc)
.map_err(|_| Error::<F>::DataCorrupted("prune location overflow"))?;
let size = self.grafted_tree.size();
let mut pinned = BTreeMap::new();
for pos in F::nodes_to_pin(prune_loc) {
let digest = self
.grafted_tree
.get_node(pos)
.ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
pinned.insert(pos, digest);
}
let mut retained = Vec::with_capacity((*size - *prune_pos) as usize);
for p in *prune_pos..*size {
let digest = self
.grafted_tree
.get_node(Position::new(p))
.ok_or(Error::<F>::DataCorrupted("missing retained grafted node"))?;
retained.push(digest);
}
self.grafted_tree = Arc::new(Mem::from_pruned_with_retained(prune_pos, pinned, retained));
Ok(())
}
#[tracing::instrument(name = "qmdb.current.db.prune", level = "info", skip_all)]
pub async fn prune(&mut self, prune_loc: Location<F>) -> Result<(), Error<F>> {
let _timer = self.metrics.prune_timer();
self.metrics.prune_calls.inc();
let sync_boundary = self.sync_boundary();
if prune_loc > sync_boundary {
return Err(Error::PruneBeyondMinRequired(prune_loc, sync_boundary));
}
self.any.log.commit().await?;
self.any.prune_bitmap(sync_boundary);
self.prune_grafted_tree_to_bitmap()?;
self.sync_metadata().await?;
#[cfg(test)]
if self.halt_before_prune_log {
std::future::pending::<()>().await;
}
self.any.prune_log(prune_loc).await?;
self.any.update_metrics();
self.update_metrics();
Ok(())
}
#[tracing::instrument(name = "qmdb.current.db.rewind", level = "info", skip_all)]
pub async fn rewind(&mut self, size: Location<F>) -> Result<(), Error<F>> {
let rewind_size = *size;
let current_size = *self.any.last_commit_loc + 1;
if rewind_size == current_size {
return Ok(());
}
if rewind_size == 0 || rewind_size > current_size {
return Err(Error::Journal(JournalError::InvalidRewind(rewind_size)));
}
let pruned_chunks = self.any.bitmap.pruned_chunks();
let pruned_bits = (pruned_chunks as u64)
.checked_mul(bitmap::Prunable::<N>::CHUNK_SIZE_BITS)
.ok_or_else(|| Error::DataCorrupted("pruned ops leaves overflow"))?;
if rewind_size < pruned_bits {
return Err(Error::Journal(JournalError::ItemPruned(rewind_size - 1)));
}
if let Some(rewind_floor) = self.delayed_merge_rewind_floor() {
if rewind_size < rewind_floor {
return Err(Error::Journal(JournalError::ItemPruned(rewind_size - 1)));
}
}
{
let rewind_last_loc = Location::<F>::new(rewind_size - 1);
let rewind_last_op = self.any.log.read(*rewind_last_loc).await?;
let Some(rewind_floor) = rewind_last_op.has_floor() else {
return Err(Error::<F>::UnexpectedData(rewind_last_loc));
};
if *rewind_floor < pruned_bits {
return Err(Error::<F>::Journal(JournalError::ItemPruned(*rewind_floor)));
}
}
let pinned_nodes = if pruned_chunks > 0 {
let grafted_leaves = Location::<F>::new(pruned_chunks as u64);
let mut pinned_nodes = Vec::new();
for pos in F::nodes_to_pin(grafted_leaves) {
let digest = self
.grafted_tree
.get_node(pos)
.ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
pinned_nodes.push(digest);
}
pinned_nodes
} else {
Vec::new()
};
self.any.rewind(size).await?;
let ops_size = self.any.log.merkle.size();
let ops_leaves = Location::<F>::try_from(ops_size)?;
let grafted_tree = build_grafted_tree::<F, H, S, N>(
self.any.bitmap.as_ref(),
&pinned_nodes,
&self.any.log.merkle,
ops_leaves,
&self.strategy,
)
.await?;
let storage = grafting::Storage::<F, H, _, _>::new(
&grafted_tree,
grafting::height::<N>(),
&self.any.log.merkle,
);
let partial_chunk = partial_chunk(self.any.bitmap.as_ref());
let ops_root = self.any.root();
let root = compute_db_root::<F, H, _, _, N>(
self.any.bitmap.as_ref(),
&storage,
ops_leaves,
partial_chunk,
self.any.inactivity_floor_loc,
&ops_root,
)
.await?;
self.grafted_tree = Arc::new(grafted_tree);
self.root = root;
self.update_metrics();
Ok(())
}
pub(crate) async fn sync_metadata(&mut self) -> Result<(), Error<F>> {
self.metadata.clear();
let pruned_chunks_u64 = self.any.bitmap.pruned_chunks() as u64;
let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
self.metadata
.put(key, pruned_chunks_u64.to_be_bytes().to_vec());
let pruned_chunks = Location::<F>::new(pruned_chunks_u64);
for (i, grafted_pos) in F::nodes_to_pin(pruned_chunks).enumerate() {
let digest = self
.grafted_tree
.get_node(grafted_pos)
.ok_or(Error::<F>::DataCorrupted("missing grafted pinned node"))?;
let key = U64::new(NODE_PREFIX, i as u64);
self.metadata.put(key, digest.to_vec());
}
self.metadata.sync().await?;
Ok(())
}
}
pub(crate) fn sync_boundary<F: Graftable, const N: usize>(
mut floor_chunks: u64,
ops_leaves: u64,
) -> Location<F> {
let chunk_bits = bitmap::Prunable::<N>::CHUNK_SIZE_BITS;
let grafting_height = grafting::height::<N>();
while floor_chunks > 0 {
let required_ops = pair_absorption_threshold::<F, N>(floor_chunks).unwrap_or_else(|| {
let youngest_start = (floor_chunks - 1) * chunk_bits;
let pos = F::subtree_root_position(Location::<F>::new(youngest_start), grafting_height);
F::peak_birth_size(pos, grafting_height)
});
if ops_leaves >= required_ops {
break;
}
floor_chunks -= 1;
}
Location::new(floor_chunks * chunk_bits)
}
fn pair_absorption_threshold<F: Graftable, const N: usize>(chunk_count: u64) -> Option<u64> {
if chunk_count == 0 {
return None;
}
let grafting_height = grafting::height::<N>();
let youngest = chunk_count - 1;
let youngest_start = youngest << grafting_height;
let youngest_end = (youngest + 1) << grafting_height;
let youngest_pos =
F::subtree_root_position(Location::<F>::new(youngest_start), grafting_height);
if F::peak_birth_size(youngest_pos, grafting_height) <= youngest_end {
return None;
}
let pair_chunk = youngest & !1;
let pair_start = pair_chunk << grafting_height;
let pair_pos = F::subtree_root_position(Location::<F>::new(pair_start), grafting_height + 1);
Some(F::peak_birth_size(pair_pos, grafting_height + 1))
}
impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
U: Update,
C: Mutable<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
S: Strategy,
Operation<F, U>: Codec,
{
#[tracing::instrument(name = "qmdb.current.db.commit", level = "info", skip_all)]
pub async fn commit(&mut self) -> Result<(), Error<F>> {
self.any.commit().await
}
#[tracing::instrument(name = "qmdb.current.db.sync", level = "info", skip_all)]
pub async fn sync(&mut self) -> Result<(), Error<F>> {
let _timer = self.metrics.sync_timer();
self.metrics.sync_calls.inc();
self.any.sync().await?;
self.sync_metadata().await?;
self.update_metrics();
Ok(())
}
#[boxed]
pub async fn destroy(self) -> Result<(), Error<F>> {
let Self { any, metadata, .. } = self;
metadata.destroy().await?;
any.destroy().await
}
}
impl<F, E, U, C, I, H, const N: usize, S> Db<F, E, C, I, H, U, N, S>
where
F: merkle::Graftable,
E: Context,
U: Update + 'static,
C: Mutable<Item = Operation<F, U>>,
I: UnorderedIndex<Value = Location<F>>,
H: Hasher,
S: Strategy,
Operation<F, U>: Codec,
{
#[tracing::instrument(name = "qmdb.current.db.apply_batch", level = "info", skip_all)]
pub async fn apply_batch(
&mut self,
batch: Arc<super::batch::MerkleizedBatch<F, H::Digest, U, N, S>>,
) -> Result<Range<Location<F>>, Error<F>> {
let _timer = self.metrics.apply_batch_timer();
self.metrics.apply_batch_calls.inc();
let range = self.any.apply_batch(Arc::clone(&batch.inner)).await?;
Arc::make_mut(&mut self.grafted_tree).apply_batch(&batch.grafted)?;
self.root = batch.canonical_root;
self.update_metrics();
Ok(range)
}
}
pub(super) fn partial_chunk<B: bitmap::Readable<N>, const N: usize>(
bitmap: &B,
) -> Option<([u8; N], u64)> {
let (last_chunk, next_bit) = bitmap.last_chunk();
if next_bit == bitmap::Prunable::<N>::CHUNK_SIZE_BITS {
None
} else {
Some((last_chunk, next_bit))
}
}
fn graftable_chunk_window<F: merkle::Graftable, B: bitmap::Readable<N>, const N: usize>(
bitmap: &B,
ops_leaves: Location<F>,
grafting_height: u32,
) -> Result<(u64, u64), Error<F>> {
let complete = bitmap.complete_chunks() as u64;
let graftable = grafting::graftable_chunks::<F>(*ops_leaves, grafting_height).min(complete);
let pending = complete - graftable;
if pending > 1 {
return Err(Error::DataCorrupted("multiple pending bitmap chunks"));
}
let pruned = bitmap.pruned_chunks() as u64;
if pruned > graftable {
return Err(Error::DataCorrupted(
"pruned chunks exceed graftable chunks",
));
}
Ok((complete, graftable))
}
pub(super) fn pending_chunk<F: merkle::Graftable, B: bitmap::Readable<N>, const N: usize>(
bitmap: &B,
ops_leaves: Location<F>,
grafting_height: u32,
) -> Result<Option<[u8; N]>, Error<F>> {
let (complete, graftable) =
graftable_chunk_window::<F, B, N>(bitmap, ops_leaves, grafting_height)?;
if complete - graftable != 1 {
return Ok(None);
}
Ok(Some(bitmap.get_chunk(graftable as usize)))
}
pub(super) fn combine_roots<H: Hasher>(
ops_root: &H::Digest,
grafted_root: &H::Digest,
pending: Option<&H::Digest>,
partial: Option<(u64, &H::Digest)>,
) -> H::Digest {
let hasher = qmdb::hasher::<H>();
match (pending, partial) {
(None, None) => hasher.hash([ops_root.as_ref(), grafted_root.as_ref()]),
(Some(pe), None) => hasher.hash([ops_root.as_ref(), grafted_root.as_ref(), pe.as_ref()]),
(None, Some((nb, p))) => {
let nb_bytes = nb.to_be_bytes();
hasher.hash([
ops_root.as_ref(),
grafted_root.as_ref(),
nb_bytes.as_slice(),
p.as_ref(),
])
}
(Some(pe), Some((nb, p))) => {
let nb_bytes = nb.to_be_bytes();
hasher.hash([
ops_root.as_ref(),
grafted_root.as_ref(),
pe.as_ref(),
nb_bytes.as_slice(),
p.as_ref(),
])
}
}
}
#[allow(clippy::too_many_arguments)]
pub(super) async fn compute_db_root<
F: merkle::Graftable,
H: Hasher,
B: bitmap::Readable<N>,
S: MerkleStorage<F, Digest = H::Digest>,
const N: usize,
>(
status: &B,
storage: &S,
ops_leaves: Location<F>,
partial_chunk: Option<([u8; N], u64)>,
inactivity_floor: Location<F>,
ops_root: &H::Digest,
) -> Result<H::Digest, Error<F>> {
let grafted_root =
compute_grafted_root::<F, H, B, S, N>(status, storage, ops_leaves, inactivity_floor)
.await?;
let hasher = qmdb::hasher::<H>();
let pending = pending_chunk::<F, B, N>(status, ops_leaves, grafting::height::<N>())?
.map(|chunk| hasher.digest(&chunk));
let partial = partial_chunk.map(|(chunk, next_bit)| {
let digest = hasher.digest(&chunk);
(next_bit, digest)
});
Ok(combine_roots::<H>(
ops_root,
&grafted_root,
pending.as_ref(),
partial.as_ref().map(|(nb, d)| (*nb, d)),
))
}
pub(super) async fn compute_grafted_root<
F: merkle::Graftable,
H: Hasher,
B: bitmap::Readable<N>,
S: MerkleStorage<F, Digest = H::Digest>,
const N: usize,
>(
status: &B,
storage: &S,
ops_leaves: Location<F>,
inactivity_floor: Location<F>,
) -> Result<H::Digest, Error<F>> {
let size = storage.size();
let leaves = Location::try_from(size)?;
let mut peaks: Vec<H::Digest> = Vec::new();
for (peak_pos, _) in F::peaks(size) {
let digest = storage
.get_node(peak_pos)
.await?
.ok_or_else(|| merkle::Error::<F>::MissingNode(peak_pos))?;
peaks.push(digest);
}
let grafting_height = grafting::height::<N>();
let (_complete_chunks, _graftable_chunks) =
graftable_chunk_window::<F, B, N>(status, ops_leaves, grafting_height)?;
let inactive_peaks =
grafting::chunk_aligned_inactive_peaks::<F>(leaves, inactivity_floor, grafting_height)?;
let hasher = qmdb::hasher::<H>();
Ok(hasher.root(leaves, inactive_peaks, peaks.iter())?)
}
pub(super) async fn read_graft_inputs<F: merkle::Graftable, D: Digest, const N: usize>(
ops_tree: &impl MerkleStorage<F, Digest = D>,
chunks: impl IntoIterator<Item = (usize, [u8; N])>,
) -> Result<Vec<(usize, D, [u8; N])>, Error<F>> {
let grafting_height = grafting::height::<N>();
try_join_all(chunks.into_iter().map(|(chunk_idx, chunk)| async move {
let leaf_start = Location::<F>::new((chunk_idx as u64) << grafting_height);
let pos = F::subtree_root_position(leaf_start, grafting_height);
let chunk_ops_digest = ops_tree
.get_node(pos)
.await?
.ok_or(merkle::Error::<F>::MissingGraftedLeaf(pos))?;
Ok::<_, Error<F>>((chunk_idx, chunk_ops_digest, chunk))
}))
.await
}
pub(super) async fn compute_grafted_leaves<
F: merkle::Graftable,
H: Hasher,
S: Strategy,
const N: usize,
>(
ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
chunks: impl IntoIterator<Item = (usize, [u8; N])>,
strategy: &S,
) -> Result<Vec<(usize, H::Digest)>, Error<F>> {
let inputs = read_graft_inputs::<F, _, N>(ops_tree, chunks).await?;
Ok(grafting::graft_chunk_digests::<H, _, N>(strategy, inputs))
}
pub(super) async fn build_grafted_tree<
F: merkle::Graftable,
H: Hasher,
S: Strategy,
const N: usize,
>(
bitmap: &impl bitmap::Readable<N>,
pinned_nodes: &[H::Digest],
ops_tree: &impl MerkleStorage<F, Digest = H::Digest>,
ops_leaves: Location<F>,
strategy: &S,
) -> Result<Mem<F, H::Digest>, Error<F>> {
let grafting_height = grafting::height::<N>();
let pruned_chunks = bitmap.pruned_chunks();
let complete_chunks = bitmap.complete_chunks();
let graftable_chunks = grafting::graftable_chunks::<F>(*ops_leaves, grafting_height)
.min(complete_chunks as u64) as usize;
assert!(
pruned_chunks <= graftable_chunks && graftable_chunks <= complete_chunks,
"invariant violated: pruned={pruned_chunks} graftable={graftable_chunks} complete={complete_chunks}"
);
let leaves = compute_grafted_leaves::<F, H, S, N>(
ops_tree,
(pruned_chunks..graftable_chunks).map(|chunk_idx| (chunk_idx, bitmap.get_chunk(chunk_idx))),
strategy,
)
.await?;
let mut grafted_tree = if pruned_chunks > 0 {
let grafted_pruning_boundary = Location::<F>::new(pruned_chunks as u64);
Mem::from_components(Vec::new(), grafted_pruning_boundary, pinned_nodes.to_vec())
.map_err(|_| Error::<F>::DataCorrupted("grafted tree rebuild failed"))?
} else {
Mem::new()
};
if !leaves.is_empty() {
let batch = {
let batch = grafted_tree.new_batch_with_strategy(strategy.clone());
let batch = batch.add_leaf_digests(leaves.iter().map(|&(_, digest)| digest));
let grafted_hasher = grafting::hasher::<F, H>(grafting_height);
batch.merkleize(&grafted_tree, &grafted_hasher)
};
grafted_tree.apply_batch(&batch)?;
}
Ok(grafted_tree)
}
pub(super) async fn init_metadata<F: merkle::Graftable, E: Context, D: Digest>(
context: E,
partition: &str,
) -> Result<(Metadata<E, U64, Vec<u8>>, usize, Vec<D>), Error<F>> {
let metadata_cfg = MConfig {
partition: partition.into(),
codec_config: ((0..).into(), ()),
};
let metadata =
Metadata::<_, U64, Vec<u8>>::init(context.child("metadata"), metadata_cfg).await?;
let key = U64::new(PRUNED_CHUNKS_PREFIX, 0);
let pruned_chunks = match metadata.get(&key) {
Some(bytes) => u64::from_be_bytes(bytes.as_slice().try_into().map_err(|_| {
error!("pruned chunks value not a valid u64");
Error::<F>::DataCorrupted("pruned chunks value not a valid u64")
})?),
None => {
warn!("bitmap metadata does not contain pruned chunks, initializing as empty");
0
}
} as usize;
let pinned_nodes = if pruned_chunks > 0 {
let pruned_loc = Location::<F>::new(pruned_chunks as u64);
if !pruned_loc.is_valid() {
return Err(Error::DataCorrupted("pruned chunks exceeds MAX_LEAVES"));
}
let mut pinned = Vec::new();
for (index, _pos) in F::nodes_to_pin(pruned_loc).enumerate() {
let metadata_key = U64::new(NODE_PREFIX, index as u64);
let Some(bytes) = metadata.get(&metadata_key) else {
return Err(Error::DataCorrupted(
"missing pinned node in grafted tree metadata",
));
};
let digest = D::decode(bytes.as_ref())
.map_err(|_| Error::<F>::DataCorrupted("invalid pinned node digest"))?;
pinned.push(digest);
}
pinned
} else {
Vec::new()
};
Ok((metadata, pruned_chunks, pinned_nodes))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{
merkle::{hasher::Standard as StandardHasher, mmb, mmr, Bagging::ForwardFold},
qmdb::{
any::traits::{DbAny, UnmerkleizedBatch as _},
current::{tests::fixed_config, unordered::fixed},
},
translator::OneCap,
};
use commonware_codec::FixedSize;
use commonware_cryptography::{sha256, Sha256};
use commonware_macros::test_traced;
use commonware_runtime::{deterministic, Runner as _, Supervisor as _};
use commonware_utils::bitmap::Prunable as PrunableBitMap;
const N: usize = sha256::Digest::SIZE;
#[test]
fn partial_chunk_single_bit() {
let mut bm = PrunableBitMap::<N>::new();
bm.push(true);
let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
assert!(result.is_some());
let (chunk, next_bit) = result.unwrap();
assert_eq!(next_bit, 1);
assert_eq!(chunk[0], 1); }
#[test]
fn partial_chunk_aligned() {
let mut bm = PrunableBitMap::<N>::new();
for _ in 0..PrunableBitMap::<N>::CHUNK_SIZE_BITS {
bm.push(true);
}
let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
assert!(result.is_none());
}
#[test]
fn partial_chunk_partial() {
let mut bm = PrunableBitMap::<N>::new();
for _ in 0..(PrunableBitMap::<N>::CHUNK_SIZE_BITS + 5) {
bm.push(true);
}
let result = partial_chunk::<PrunableBitMap<N>, N>(&bm);
assert!(result.is_some());
let (_chunk, next_bit) = result.unwrap();
assert_eq!(next_bit, 5);
}
#[test]
fn combine_roots_deterministic() {
let ops = Sha256::hash(b"ops");
let grafted = Sha256::hash(b"grafted");
let r1 = combine_roots::<Sha256>(&ops, &grafted, None, None);
let r2 = combine_roots::<Sha256>(&ops, &grafted, None, None);
assert_eq!(r1, r2);
}
#[test]
fn combine_roots_with_partial_differs() {
let ops = Sha256::hash(b"ops");
let grafted = Sha256::hash(b"grafted");
let partial_digest = Sha256::hash(b"partial");
let without = combine_roots::<Sha256>(&ops, &grafted, None, None);
let with = combine_roots::<Sha256>(&ops, &grafted, None, Some((5, &partial_digest)));
assert_ne!(without, with);
}
#[test]
fn combine_roots_with_pending_differs() {
let ops = Sha256::hash(b"ops");
let grafted = Sha256::hash(b"grafted");
let pending_digest = Sha256::hash(b"pending");
let without = combine_roots::<Sha256>(&ops, &grafted, None, None);
let with = combine_roots::<Sha256>(&ops, &grafted, Some(&pending_digest), None);
assert_ne!(without, with);
}
#[test]
fn combine_roots_pending_and_partial_independent() {
let ops = Sha256::hash(b"ops");
let grafted = Sha256::hash(b"grafted");
let pending_digest = Sha256::hash(b"pending");
let partial_digest = Sha256::hash(b"partial");
let only_pending = combine_roots::<Sha256>(&ops, &grafted, Some(&pending_digest), None);
let only_partial =
combine_roots::<Sha256>(&ops, &grafted, None, Some((5, &partial_digest)));
let both = combine_roots::<Sha256>(
&ops,
&grafted,
Some(&pending_digest),
Some((5, &partial_digest)),
);
assert_ne!(only_pending, only_partial);
assert_ne!(only_pending, both);
assert_ne!(only_partial, both);
}
#[test]
fn combine_roots_different_ops_root() {
let ops_a = Sha256::hash(b"ops_a");
let ops_b = Sha256::hash(b"ops_b");
let grafted = Sha256::hash(b"grafted");
let r1 = combine_roots::<Sha256>(&ops_a, &grafted, None, None);
let r2 = combine_roots::<Sha256>(&ops_b, &grafted, None, None);
assert_ne!(r1, r2);
}
#[test]
fn combine_roots_format_golden() {
let hasher = StandardHasher::<Sha256>::new(ForwardFold);
let ops = Sha256::hash(b"ops");
let grafted = Sha256::hash(b"grafted");
let pending = Sha256::hash(b"pending");
let partial = Sha256::hash(b"partial");
let next_bit: u64 = 0x1122_3344_5566_7788;
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, None, None),
hasher.hash([ops.as_ref(), grafted.as_ref()])
);
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, Some(&pending), None),
hasher.hash([ops.as_ref(), grafted.as_ref(), pending.as_ref()])
);
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, None, Some((next_bit, &partial))),
hasher.hash([
ops.as_ref(),
grafted.as_ref(),
next_bit.to_be_bytes().as_slice(),
partial.as_ref(),
])
);
assert_eq!(
combine_roots::<Sha256>(&ops, &grafted, Some(&pending), Some((next_bit, &partial))),
hasher.hash([
ops.as_ref(),
grafted.as_ref(),
pending.as_ref(),
next_bit.to_be_bytes().as_slice(),
partial.as_ref(),
])
);
}
type MmrDb = fixed::Db<
mmr::Family,
deterministic::Context,
sha256::Digest,
sha256::Digest,
Sha256,
OneCap,
32,
commonware_parallel::Sequential,
>;
type MmbDb = fixed::Db<
mmb::Family,
deterministic::Context,
sha256::Digest,
sha256::Digest,
Sha256,
OneCap,
32,
commonware_parallel::Sequential,
>;
async fn populate_fixed_db<F, DB>(db: &mut DB, start: u64, count: u64)
where
F: merkle::Graftable,
DB: DbAny<F, Key = sha256::Digest, Value = sha256::Digest>,
{
let mut batch = db.new_batch();
for idx in start..start + count {
let key = Sha256::hash(&idx.to_be_bytes());
let value = Sha256::hash(&(idx + count).to_be_bytes());
batch = batch.write(key, Some(value));
}
let merkleized = batch.merkleize(db, None).await.unwrap();
db.apply_batch(merkleized).await.unwrap();
db.commit().await.unwrap();
}
#[test_traced]
fn test_current_prune_dropped_before_log_prune() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let mut db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("prune-park", &ctx),
)
.await
.unwrap();
populate_fixed_db::<mmr::Family, _>(&mut db, 0, 512).await;
let durable_floor = db.inactivity_floor_loc();
{
let mut batch = db.new_batch();
for idx in 0..512u64 {
let key = Sha256::hash(&idx.to_be_bytes());
let value = Sha256::hash(&(idx + 1024).to_be_bytes());
batch = batch.write(key, Some(value));
}
let merkleized = batch.merkleize(&db, None).await.unwrap();
db.apply_batch(merkleized).await.unwrap();
}
assert!(db.sync_boundary() > durable_floor);
let bounds = db.bounds();
let floor = db.inactivity_floor_loc();
let root = db.root();
db.halt_before_prune_log = true;
{
let fut = db.prune(db.sync_boundary());
futures::pin_mut!(fut);
assert!(
futures::poll!(fut.as_mut()).is_pending(),
"prune must park before the log prune"
);
}
let pruned_bits = db.any.bitmap.pruned_bits();
assert!(pruned_bits > *durable_floor);
drop(db);
let db = MmrDb::init(
ctx.child("reopen"),
fixed_config::<OneCap>("prune-park", &ctx),
)
.await
.expect("prune crash must leave the db recoverable");
assert_eq!(db.bounds(), bounds);
assert_eq!(db.inactivity_floor_loc(), floor);
assert_eq!(db.root(), root);
assert_eq!(db.any.bitmap.pruned_bits(), pruned_bits);
db.destroy().await.unwrap();
});
}
#[test_traced]
fn test_ops_root_witness_verifies_without_partial_chunk() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let mut db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-full", &ctx),
)
.await
.unwrap();
let mut next_idx = 0;
populate_fixed_db::<mmr::Family, _>(&mut db, next_idx, 256).await;
next_idx += 256;
while partial_chunk::<_, 32>(db.any.bitmap.as_ref()).is_some() {
populate_fixed_db::<mmr::Family, _>(&mut db, next_idx, 1).await;
next_idx += 1;
}
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.partial_chunk.is_none());
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
let wrong_ops_root = Sha256::hash(b"wrong ops root");
assert!(!witness.verify::<Sha256>(&wrong_ops_root, &canonical_root));
let wrong_canonical_root = Sha256::hash(b"wrong canonical root");
assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
let mut tampered = witness;
tampered.grafted_root = Sha256::hash(b"wrong grafted root");
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
});
}
#[test_traced]
fn test_ops_root_witness_verifies_with_partial_chunk() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let mut db = MmbDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-partial", &ctx),
)
.await
.unwrap();
populate_fixed_db::<mmb::Family, _>(&mut db, 0, 260).await;
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.partial_chunk.is_some());
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
let wrong_ops_root = Sha256::hash(b"wrong ops root");
assert!(!witness.verify::<Sha256>(&wrong_ops_root, &canonical_root));
let wrong_canonical_root = Sha256::hash(b"wrong canonical root");
assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
let mut tampered = witness.clone();
tampered.grafted_root = Sha256::hash(b"wrong grafted root");
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
let mut tampered = witness.clone();
tampered.partial_chunk.as_mut().unwrap().0 += 1;
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
let mut tampered = witness;
tampered.partial_chunk.as_mut().unwrap().1 = Sha256::hash(b"wrong partial chunk");
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
});
}
#[test_traced]
fn test_ops_root_witness_verifies_with_pruned_db() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let mut db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-pruned", &ctx),
)
.await
.unwrap();
for _ in 0..5 {
populate_fixed_db::<mmr::Family, _>(&mut db, 0, 512).await;
}
db.prune(db.sync_boundary()).await.unwrap();
assert!(
db.any.bitmap.pruned_chunks() > 0,
"test requires at least one pruned chunk to exercise the zero-chunk path"
);
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
let wrong_canonical_root = Sha256::hash(b"wrong canonical root");
assert!(!witness.verify::<Sha256>(&ops_root, &wrong_canonical_root));
let mut tampered = witness;
tampered.grafted_root = Sha256::hash(b"wrong grafted root");
assert!(!tampered.verify::<Sha256>(&ops_root, &canonical_root));
});
}
#[test_traced]
fn test_ops_root_witness_verifies_on_fresh_db() {
let executor = deterministic::Runner::default();
executor.start(|ctx| async move {
let db = MmrDb::init(
ctx.child("storage"),
fixed_config::<OneCap>("ops-root-witness-fresh", &ctx),
)
.await
.unwrap();
let witness = db.ops_root_witness().await.unwrap();
let ops_root = db.ops_root();
let canonical_root = db.root();
assert!(witness.verify::<Sha256>(&ops_root, &canonical_root));
});
}
}