//! [`ManagedDb`] implementation for QMDB [`any`](commonware_storage::qmdb::any) databases. //! //! The QMDB batch API passes `&db` to `get()` and `merkleize()` for //! read-through to applied state. This module provides wrapper types //! that capture a [`Shared`] database handle alongside the raw batch so the //! [`Unmerkleized`](super::Unmerkleized) and [`Merkleized`](super::Merkleized) //! traits can be implemented without a DB parameter. use crate::stateful::db::{ BatchContext, ManagedDb, Merkleized as MerkleizedTrait, Shared, StateSyncDb, SyncEngineConfig, Unmerkleized as UnmerkleizedTrait, sync_standard_db, }; use commonware_codec::{Codec, Read as CodecRead}; use commonware_cryptography::Hasher; use commonware_parallel::Strategy; use commonware_runtime::{Handle, Spawner}; use commonware_storage::{ Context, index::{ Ordered as OrderedIndex, Unordered as UnorderedIndex, unordered::Index as UnorderedIdx, }, journal::contiguous::{ Contiguous, Mutable, fixed::Journal as FixedJournal, variable::Journal as VariableJournal, }, merkle::{Family, Location}, qmdb::{ Error, any::{ FixedConfig, VariableConfig, batch::{MerkleizedBatch, Staged, UnmerkleizedBatch}, db::Db, initial_root, operation::{Operation, Update}, ordered, unordered, value::{self, FixedEncoding, ValueEncoding, VariableEncoding}, }, operation::Key, sync::{self, Target as AnySyncTarget}, }, translator::Translator, }; use commonware_utils::{Array, channel::mpsc, non_empty_range}; use std::{ ops::{Deref, Range}, sync::Arc, }; // Matches commonware_storage::qmdb::any::BITMAP_CHUNK_BYTES, which is crate-private. const ANY_BITMAP_CHUNK_BYTES: usize = 64; /// Wraps a QMDB [`UnmerkleizedBatch`] with a reference to the parent /// database, implementing the [`Unmerkleized`](super::Unmerkleized) trait. pub struct AnyUnmerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex>, H: Hasher, S: Strategy, Operation: Codec, { batch: UnmerkleizedBatch, db: Shared>, metadata: Option, } /// Staged batch returned by [`AnyUnmerkleized::stage`], wrapping a QMDB [`Staged`] with a /// reference to the parent database. /// /// Like any speculative batch, this handle is a branch-scoped view of the shared database: it /// stays valid only while every batch finalized on the database is an ancestor of this batch /// (see [`MerkleizedBatch`]'s branch-validity contract). pub struct AnyStaged where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex>, H: Hasher, S: Strategy, Operation: Codec, { staged: Staged, db: Shared>, metadata: Option, } /// Key-value operations shared by both `any` update kinds. impl AnyUnmerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex> + 'static, H: Hasher, S: Strategy, Operation: Codec, { /// Set commit metadata included in the next /// [`merkleize`](UnmerkleizedTrait::merkleize) call. pub fn with_metadata(mut self, metadata: U::Value) -> Self { self.metadata = Some(metadata); self } /// Read a value by key, falling back to applied state. pub async fn get(&self, key: &U::Key) -> Result, Error> { let db = self.db.read().await; self.batch.get(key, &db).await } /// Read multiple values by key, falling back to applied state. /// /// Returns results in the same order as the input keys. pub async fn get_many(&self, keys: &[&U::Key]) -> Result>, Error> { let db = self.db.read().await; self.batch.get_many(keys, &db).await } /// Read multiple values and return a staged batch for the same keys. /// /// Returns results in the same order as the input keys. pub async fn stage( self, keys: &[&U::Key], ) -> Result<(Vec>, AnyStaged), Error> { let Self { batch, db, metadata, } = self; let (values, staged) = { let guard = db.read().await; batch.stage(keys, &guard).await? }; Ok(( values, AnyStaged { staged, db, metadata, }, )) } /// Record a mutation. `Some(value)` for upsert, `None` for delete. pub fn write(mut self, key: U::Key, value: Option) -> Self { self.batch = self.batch.write(key, value); self } } /// Wraps a QMDB [`MerkleizedBatch`] with a reference to the parent /// database, implementing the [`Merkleized`](super::Merkleized) trait. pub struct AnyMerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex>, H: Hasher, S: Strategy, Operation: Codec, { inner: Arc>, db: Shared>, } impl Clone for AnyMerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex>, H: Hasher, S: Strategy, Operation: Codec, { fn clone(&self) -> Self { Self { inner: Arc::clone(&self.inner), db: self.db.clone(), } } } impl Deref for AnyUnmerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex>, H: Hasher, S: Strategy, Operation: Codec, { type Target = UnmerkleizedBatch; fn deref(&self) -> &Self::Target { &self.batch } } impl Deref for AnyMerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex>, H: Hasher, S: Strategy, Operation: Codec, { type Target = MerkleizedBatch; fn deref(&self) -> &Self::Target { &self.inner } } /// Read-expansion operations for the `any` staged batch. impl AnyStaged where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex> + 'static, H: Hasher, S: Strategy, Operation: Codec, { /// Set commit metadata included in the [`merkleize`](Self::merkleize) call, replacing any /// metadata set before staging. pub fn with_metadata(mut self, metadata: U::Value) -> Self { self.metadata = Some(metadata); self } /// Expand this staged batch with more reads. /// /// Existing read indices remain stable. Newly read keys are appended to the staged read set and /// assigned the returned range. Expansion does not deduplicate against previously staged keys /// and does not observe values computed for earlier staged slots but not yet passed to /// `merkleize`. pub async fn expand( self, keys: &[&U::Key], ) -> Result<(Range, Vec>, Self), Error> { let Self { staged, db, metadata, } = self; let (range, values, staged) = { let guard = db.read().await; staged.expand(keys, &guard).await? }; Ok(( range, values, Self { staged, db, metadata, }, )) } } /// Staged merkleize for the `any` unordered update kind. impl AnyStaged, S> where F: Family, E: Context, K: Key, V: ValueEncoding + 'static, C: Mutable>>, I: UnorderedIndex> + 'static, H: Hasher, S: Strategy, Operation>: Codec, { /// Record updates for staged reads and upserts for unread keys, then merkleize. /// /// Consumes the staged handle and write vectors. Call [`expand`](AnyStaged::expand) before /// this method if more keys must be read into the staged index space. /// /// A `Some` value is an upsert. `None` is a delete. Update indices refer to the staged read /// set: the initial `stage` input followed by any [`expand`](AnyStaged::expand) ranges. Metadata /// set via [`with_metadata`](AnyStaged::with_metadata) (or before staging) is committed with the /// returned batch. /// /// # Panics /// /// Panics if any update's `read_index` is out of the staged read range. pub async fn merkleize( self, updates: Vec<(usize, Option)>, upserts: Vec<(K, Option)>, ) -> Result, S>, Error> { let Self { staged, db, metadata, } = self; let inner = { let guard = db.read().await; staged.merkleize(updates, upserts, metadata, &guard).await? }; Ok(AnyMerkleized { inner, db }) } } /// Staged merkleize for the `any` ordered update kind. impl AnyStaged, S> where F: Family, E: Context, K: Key, V: ValueEncoding + 'static, C: Mutable>>, I: OrderedIndex> + 'static, H: Hasher, S: Strategy, Operation>: Codec, { /// Record updates for staged reads and upserts for unread keys, then merkleize. /// /// Consumes the staged handle and write vectors. Call [`expand`](AnyStaged::expand) before /// this method if more keys must be read into the staged index space. /// /// A `Some` value is an upsert. `None` is a delete. Update indices refer to the staged read /// set: the initial `stage` input followed by any [`expand`](AnyStaged::expand) ranges. Metadata /// set via [`with_metadata`](AnyStaged::with_metadata) (or before staging) is committed with the /// returned batch. /// /// # Panics /// /// Panics if any update's `read_index` is out of the staged read range. pub async fn merkleize( self, updates: Vec<(usize, Option)>, upserts: Vec<(K, Option)>, ) -> Result, S>, Error> { let Self { staged, db, metadata, } = self; let inner = { let guard = db.read().await; staged.merkleize(updates, upserts, metadata, &guard).await? }; Ok(AnyMerkleized { inner, db }) } } /// Read-through operations for the `any` merkleized batch. impl AnyMerkleized where F: Family, E: Context, U: Update, C: Contiguous>, I: UnorderedIndex> + 'static, H: Hasher, S: Strategy, Operation: Codec, { /// Read a value by key, falling back to applied state. pub async fn get(&self, key: &U::Key) -> Result, Error> { let db = self.db.read().await; self.inner.get(key, &db).await } /// Read multiple values by key, falling back to applied state. /// /// Returns results in the same order as the input keys. pub async fn get_many(&self, keys: &[&U::Key]) -> Result>, Error> { let db = self.db.read().await; self.inner.get_many(keys, &db).await } } /// Implement [`Unmerkleized`](UnmerkleizedTrait) for the `any` unordered update kind. impl UnmerkleizedTrait for AnyUnmerkleized, S> where F: Family, E: Context, K: Key, V: ValueEncoding + 'static, C: Mutable>>, I: UnorderedIndex> + 'static, H: Hasher, S: Strategy, Operation>: Codec, { type Merkleized = AnyMerkleized, S>; type Error = Error; async fn merkleize(self) -> Result> { let db = self.db.read().await; let merkleized = self.batch.merkleize(&db, self.metadata).await?; Ok(AnyMerkleized { inner: merkleized, db: self.db.clone(), }) } } /// Implement [`Unmerkleized`](UnmerkleizedTrait) for the `any` ordered update kind. impl UnmerkleizedTrait for AnyUnmerkleized, S> where F: Family, E: Context, K: Key, V: ValueEncoding + 'static, C: Mutable>>, I: OrderedIndex> + 'static, H: Hasher, S: Strategy, Operation>: Codec, { type Merkleized = AnyMerkleized, S>; type Error = Error; async fn merkleize(self) -> Result> { let db = self.db.read().await; let merkleized = self.batch.merkleize(&db, self.metadata).await?; Ok(AnyMerkleized { inner: merkleized, db: self.db.clone(), }) } } /// Implement [`Merkleized`](MerkleizedTrait) for all supported `any` update kinds. impl MerkleizedTrait for AnyMerkleized where F: Family, E: Context, U: Update, C: Mutable>, I: UnorderedIndex> + 'static, H: Hasher, S: Strategy, Operation: Codec, AnyUnmerkleized: UnmerkleizedTrait, { type Digest = H::Digest; type Unmerkleized = AnyUnmerkleized; fn root(&self) -> H::Digest { self.inner.root() } fn new_batch(&self) -> Self::Unmerkleized { AnyUnmerkleized { batch: self.inner.new_batch::(), db: self.db.clone(), metadata: None, } } } /// Implement [`ManagedDb`] for unordered QMDB databases with fixed-size values. /// /// `new_batch` captures the [`Shared`] database handle in the returned /// wrapper so that `get()` and `merkleize()` can read through to /// applied state. impl ManagedDb for Db< F, E, FixedJournal>>>, UnorderedIdx>, H, unordered::Update>, ANY_BITMAP_CHUNK_BYTES, S, > where F: Family, E: Context + Spawner, K: Array, V: value::FixedValue + 'static, H: Hasher + 'static, T: Translator, S: Strategy, { type Unmerkleized = AnyUnmerkleized< F, E, FixedJournal>>>, UnorderedIdx>, H, unordered::Update>, S, >; type Merkleized = AnyMerkleized< F, E, FixedJournal>>>, UnorderedIdx>, H, unordered::Update>, S, >; type Error = Error; type Config = FixedConfig; type SyncTarget = AnySyncTarget; async fn init(context: E, config: Self::Config) -> Result> { ::init(context, config).await } fn initial_sync_target() -> Self::SyncTarget { AnySyncTarget::new( initial_root::>, H>(), non_empty_range!(Location::new(0), Location::new(1)), ) } fn new_batch(database: BatchContext<'_, Self>) -> Self::Unmerkleized { let (database, shared) = database.into_parts(); AnyUnmerkleized { batch: database.new_batch(), db: shared, metadata: None, } } fn matches_sync_target(batch: &Self::Merkleized, target: &Self::SyncTarget) -> bool { batch.root() == target.root && *target.range.start() == batch.bounds().inactivity_floor && *target.range.end() == batch.bounds().tip.size } async fn apply(self, batch: Self::Merkleized) -> Result> { let (db, _) = self.apply_batch(batch.inner).await?; Ok(db) } async fn finalize(self) -> Result<(Self, Handle<()>), Error> { self.start_sync().await } async fn prune(self, target: &Self::SyncTarget) -> Result> { self.prune((*target.range.start()).into()).await } fn sync_target(&self) -> Self::SyncTarget { let bounds = self.bounds(); AnySyncTarget::new( self.root(), non_empty_range!(self.sync_boundary(), bounds.end), ) } async fn rewind_to_target(self, target: Self::SyncTarget) -> Result> { let db = self.rewind(target.range.end()).await?; let db = db.sync().await?; let rewound_target = db.sync_target(); assert_eq!( rewound_target, target, "rewound database target mismatch after rewind", ); Ok(db) } } /// Implement [`ManagedDb`] for unordered QMDB databases with variable-size values. impl ManagedDb for Db< F, E, VariableJournal>>>, UnorderedIdx>, H, unordered::Update>, ANY_BITMAP_CHUNK_BYTES, S, > where F: Family, E: Context + Spawner, K: Key, V: value::VariableValue + 'static, H: Hasher, T: Translator, S: Strategy, Operation>>: Codec, { type Unmerkleized = AnyUnmerkleized< F, E, VariableJournal>>>, UnorderedIdx>, H, unordered::Update>, S, >; type Merkleized = AnyMerkleized< F, E, VariableJournal>>>, UnorderedIdx>, H, unordered::Update>, S, >; type Error = Error; type Config = VariableConfig< T, >> as CodecRead>::Cfg, S, >; type SyncTarget = AnySyncTarget; async fn init(context: E, config: Self::Config) -> Result> { ::init(context, config).await } fn initial_sync_target() -> Self::SyncTarget { AnySyncTarget::new( initial_root::>, H>(), non_empty_range!(Location::new(0), Location::new(1)), ) } fn new_batch(database: BatchContext<'_, Self>) -> Self::Unmerkleized { let (database, shared) = database.into_parts(); AnyUnmerkleized { batch: database.new_batch(), db: shared, metadata: None, } } fn matches_sync_target(batch: &Self::Merkleized, target: &Self::SyncTarget) -> bool { batch.root() == target.root && *target.range.start() == batch.bounds().inactivity_floor && *target.range.end() == batch.bounds().tip.size } async fn apply(self, batch: Self::Merkleized) -> Result> { let (db, _) = self.apply_batch(batch.inner).await?; Ok(db) } async fn finalize(self) -> Result<(Self, Handle<()>), Error> { self.start_sync().await } async fn prune(self, target: &Self::SyncTarget) -> Result> { self.prune((*target.range.start()).into()).await } fn sync_target(&self) -> Self::SyncTarget { let bounds = self.bounds(); AnySyncTarget::new( self.root(), non_empty_range!(self.sync_boundary(), bounds.end), ) } async fn rewind_to_target(self, target: Self::SyncTarget) -> Result> { let db = self.rewind(target.range.end()).await?; let db = db.sync().await?; let rewound_target = db.sync_target(); assert_eq!( rewound_target, target, "rewound database target mismatch after rewind", ); Ok(db) } } impl StateSyncDb for Db< F, E, FixedJournal>>>, UnorderedIdx>, H, unordered::Update>, ANY_BITMAP_CHUNK_BYTES, S, > where F: Family, E: Context + Spawner, K: Array, V: value::FixedValue + 'static, H: Hasher, T: Translator, S: Strategy, R: sync::SourceFor, { type SyncError = sync::Error; async fn sync_db( context: E, config: Self::Config, source: R, target: Self::SyncTarget, tip_updates: mpsc::Receiver, finish: Option>, reached_target: Option>, sync_config: SyncEngineConfig, ) -> Result { sync_standard_db( context, config, source, target, tip_updates, finish, reached_target, sync_config, ) .await } } impl StateSyncDb for Db< F, E, VariableJournal>>>, UnorderedIdx>, H, unordered::Update>, ANY_BITMAP_CHUNK_BYTES, S, > where F: Family, E: Context + Spawner, K: Key, V: value::VariableValue + 'static, H: Hasher, T: Translator, S: Strategy, Operation>>: Codec, R: sync::SourceFor, { type SyncError = sync::Error; async fn sync_db( context: E, config: Self::Config, source: R, target: Self::SyncTarget, tip_updates: mpsc::Receiver, finish: Option>, reached_target: Option>, sync_config: SyncEngineConfig, ) -> Result { sync_standard_db( context, config, source, target, tip_updates, finish, reached_target, sync_config, ) .await } } #[cfg(test)] mod tests { use super::*; use commonware_cryptography::{Sha256, sha256::Digest}; use commonware_parallel::Sequential; use commonware_runtime::{ BufferPooler, Runner as _, Supervisor as _, buffer::paged::CacheRef, deterministic, mocks::{DelayedSyncContext, PendingSyncs, drive_pending_syncs, release_pending_syncs}, }; use commonware_storage::{ journal::contiguous::fixed::Config as FixedJournalConfig, merkle::{full::Config as MerkleConfig, mmr}, qmdb::any::unordered::fixed, translator::TwoCap, }; use commonware_utils::{NZU16, NZU64, NZUsize}; use std::num::{NonZeroU16, NonZeroUsize}; type UnorderedFixedDb = fixed::Db; const PAGE_SIZE: NonZeroU16 = NZU16!(101); const PAGE_CACHE_SIZE: NonZeroUsize = NZUsize!(11); fn fixed_config(suffix: &str, pooler: &impl BufferPooler) -> FixedConfig { let page_cache = CacheRef::from_pooler(pooler, PAGE_SIZE, PAGE_CACHE_SIZE); FixedConfig { merkle_config: MerkleConfig { journal_partition: format!("stateful-any-journal-{suffix}"), metadata_partition: format!("stateful-any-metadata-{suffix}"), items_per_blob: NZU64!(11), write_buffer: NZUsize!(1024), replay_buffer: NZUsize!(1024), strategy: Sequential, page_cache: page_cache.clone(), }, journal_config: FixedJournalConfig { partition: format!("stateful-any-log-{suffix}"), items_per_blob: NZU64!(7), page_cache, write_buffer: NZUsize!(1024), replay_buffer: NZUsize!(1024), }, translator: TwoCap, init_cache_size: Some(NZUsize!(1024)), init_buffer: NZUsize!(1 << 21), init_concurrency: (), } } #[test] fn unmerkleized_batch_falls_through_to_applied_state() { deterministic::Runner::default().start(|context| async move { let config = fixed_config("unordered-fixed-live-fallback", &context); let db = >::init(context.child("db"), config) .await .unwrap(); let db = Shared::new("test", db); let key = Sha256::hash(&[b"key"]); let value = Sha256::hash(&[b"winner"]); let pre_finalization = db.new_batch_for_test::<_>().await; let winner = db.new_batch_for_test::<_>().await.write(key, Some(value)); let winner = crate::stateful::db::Unmerkleized::merkleize(winner) .await .unwrap(); let (slot, database) = db.write().await; let database = >::apply(database, winner) .await .unwrap(); let (database, sync) = >::finalize(database) .await .unwrap(); slot.put(database); sync.await.expect("database sync failed"); // The batch commitment does not provide a historical database view. assert_eq!(pre_finalization.get(&key).await.unwrap(), Some(value)); }); } /// The glue staged wrapper (`AnyUnmerkleized::stage` -> `AnyStaged::expand` -> /// `AnyStaged::merkleize`) must return the same values and root as an explicit `get_many` + /// `write` + `merkleize`, including a staged delete, an upsert, and metadata flow (both set /// on the staged handle via `with_metadata` and carried from before staging). This guards /// metadata flow and db-handle pairing through the wrapper. #[test] fn unordered_fixed_staged_merkleize_matches_explicit_writes() { deterministic::Runner::default().start(|context| async move { let config = fixed_config("unordered-fixed-glue-staged", &context); let db = >::init(context.child("db"), config) .await .unwrap(); let db = Shared::new("test", db); let key = |i: u64| Sha256::hash(&[&i.to_be_bytes()]); let val = |i: u64| Sha256::hash(&[&(i + 10_000).to_be_bytes()]); let metadata = Sha256::hash(&[b"metadata"]); // Seed keys 0..50 and persist them. let mut seed = db.new_batch_for_test::<_>().await; for i in 0..50u64 { seed = seed.write(key(i), Some(val(i))); } let merkleized = crate::stateful::db::Unmerkleized::merkleize(seed) .await .unwrap(); { let (slot, database) = db.write().await; let database = >::apply(database, merkleized) .await .unwrap(); let (database, sync) = >::finalize(database) .await .unwrap(); slot.put(database); sync.await.expect("database sync failed"); } // Read set: key(1) updated, key(2) deleted, key(999) missing -> created. let read_keys = [key(1), key(2), key(999)]; let keys: Vec<&Digest> = read_keys.iter().collect(); let indexed_updates = vec![(0, Some(val(1_000))), (1, None), (2, Some(val(1_001)))]; let upserts = vec![(key(3), Some(val(1_002)))]; // Explicit path. let mut explicit = db.new_batch_for_test::<_>().await; let explicit_values = explicit.get_many(&keys).await.unwrap(); for (slot, value) in &indexed_updates { explicit = explicit.write(read_keys[*slot], *value); } for (k, v) in &upserts { explicit = explicit.write(*k, *v); } let explicit_root = crate::stateful::db::Unmerkleized::merkleize(explicit.with_metadata(metadata)) .await .unwrap() .root(); // Staged path, with metadata set on the staged handle. let staged_batch = db.new_batch_for_test::<_>().await; let split = 2; let (mut staged_values, staged) = staged_batch.stage(&keys[..split]).await.unwrap(); let (range, suffix_values, staged) = staged.expand(&keys[split..]).await.unwrap(); assert_eq!(range, split..keys.len()); staged_values.extend(suffix_values); let staged_root = staged .with_metadata(metadata) .merkleize(indexed_updates.clone(), upserts.clone()) .await .unwrap() .root(); assert_eq!(explicit_values, staged_values); assert_eq!(explicit_root, staged_root); // Metadata set before staging must be carried through to staged merkleize. let carried_batch = db.new_batch_for_test::<_>().await.with_metadata(metadata); let (carried_values, staged) = carried_batch.stage(&keys).await.unwrap(); let carried_root = staged .merkleize(indexed_updates.clone(), upserts.clone()) .await .unwrap() .root(); assert_eq!(explicit_values, carried_values); assert_eq!(explicit_root, carried_root); }); } type DelayedFixedDb = fixed::Db< mmr::Family, DelayedSyncContext, Digest, Digest, Sha256, TwoCap, Sequential, >; #[test] fn apply_does_not_finalize() { deterministic::Runner::default().start(|context| async move { let pending = PendingSyncs::default(); let delayed = DelayedSyncContext { inner: context, pending: pending.clone(), }; let config = fixed_config("unordered-fixed-deferred", &delayed); let db = drive_pending_syncs( &pending, >::init(delayed.child("db"), config), ) .await .unwrap(); let db = Shared::new("test", db); let key = Sha256::hash(&[b"key"]); let value = Sha256::hash(&[b"value"]); let batch = db.new_batch_for_test::<_>().await.write(key, Some(value)); let merkleized = crate::stateful::db::Unmerkleized::merkleize(batch) .await .unwrap(); let starts = pending.starts(); let (slot, database) = db.write().await; let database = >::apply(database, merkleized) .await .unwrap(); slot.put(database); assert_eq!(pending.starts(), starts, "apply must not start durability"); { let guard = db.read().await; assert_eq!(guard.get(&key).await.unwrap(), Some(value)); } let (slot, database) = db.write().await; let (database, sync) = >::finalize(database) .await .unwrap(); slot.put(database); assert!( pending.starts() > pending.completions(), "finalize must leave its flush parked", ); release_pending_syncs(&pending); drive_pending_syncs(&pending, sync) .await .expect("flush must succeed once released"); }); } }