Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(node): avoid cloning in-memory data structures on apply_block #532

Merged
merged 8 commits into from
Oct 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
- Changed the `BlockWitness` to pass the inputs to the VM using only advice provider (#516).
- [BREAKING] Improved store API errors (return "not found" instead of "internal error" status if requested account(s) not found) (#518).
- [BREAKING] Migrated to v0.11 version of Miden VM (#528).
- Reduce cloning in the store's `apply_block` (#532).

## 0.5.1 (2024-09-12)

Expand Down
45 changes: 23 additions & 22 deletions crates/store/src/nullifier_tree.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
use miden_objects::{
crypto::{
hash::rpo::RpoDigest,
merkle::{Smt, SmtProof},
merkle::{MutationSet, Smt, SmtProof, SMT_DEPTH},
},
notes::Nullifier,
Felt, FieldElement, Word,
Expand All @@ -27,7 +27,7 @@ impl NullifierTree {
Ok(Self(inner))
}

/// Get SMT root.
/// Returns the root of the nullifier SMT.
pub fn root(&self) -> RpoDigest {
self.0.root()
}
Expand All @@ -37,26 +37,6 @@ impl NullifierTree {
self.0.open(&nullifier.inner())
}

/// Inserts block number in which nullifier was consumed.
pub fn insert(
&mut self,
nullifier: &Nullifier,
block_num: BlockNumber,
) -> Result<(), NullifierTreeError> {
let key = nullifier.inner();
let prev_value = self.0.get_value(&key);
if prev_value != Smt::EMPTY_VALUE {
return Err(NullifierTreeError::NullifierAlreadyExists {
nullifier: *nullifier,
block_num: Self::leaf_value_to_block_num(prev_value),
});
}

self.0.insert(key, Self::block_num_to_leaf_value(block_num));

Ok(())
}

/// Returns block number stored for the given nullifier or `None` if the nullifier wasn't
/// consumed.
pub fn get_block_num(&self, nullifier: &Nullifier) -> Option<BlockNumber> {
Expand All @@ -68,6 +48,27 @@ impl NullifierTree {
Some(Self::leaf_value_to_block_num(value))
}

/// Computes mutations for the nullifier SMT.
pub fn compute_mutations(
&self,
kv_pairs: impl IntoIterator<Item = (Nullifier, BlockNumber)>,
) -> MutationSet<SMT_DEPTH, RpoDigest, Word> {
self.0.compute_mutations(kv_pairs.into_iter().map(|(nullifier, block_num)| {
(nullifier.inner(), Self::block_num_to_leaf_value(block_num))
}))
}

/// Applies mutations to the nullifier SMT.
pub fn apply_mutations(
&mut self,
mutations: MutationSet<SMT_DEPTH, RpoDigest, Word>,
) -> Result<(), NullifierTreeError> {
self.0.apply_mutations(mutations).map_err(Into::into)
}

// HELPER FUNCTIONS
// --------------------------------------------------------------------------------------------

/// Returns the nullifier's leaf value in the SMT by its block number.
fn block_num_to_leaf_value(block: BlockNumber) -> Word {
[Felt::from(block), Felt::ZERO, Felt::ZERO, Felt::ZERO]
Expand Down
161 changes: 86 additions & 75 deletions crates/store/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -203,15 +203,21 @@ impl State {
// block store. Thus, such block should be considered as block candidates, but not
// finalized blocks. So we should check for the latest block when getting block from
// the store.
let store = self.block_store.clone();
let store = Arc::clone(&self.block_store);
let block_save_task =
tokio::spawn(async move { store.save_block(block_num, &block_data).await });

// scope to read in-memory data, validate the request, and compute intermediary values
let (account_tree, chain_mmr, nullifier_tree, notes) = {
// scope to read in-memory data, compute mutations required for updating account
// and nullifier trees, and validate the request
let (
nullifier_tree_old_root,
nullifier_tree_update,
account_tree_old_root,
account_tree_update,
) = {
let inner = self.inner.read().await;

let span = info_span!(target: COMPONENT, "update_in_memory_structs").entered();
let _span = info_span!(target: COMPONENT, "update_in_memory_structs").entered();

// nullifiers can be produced only once
let duplicate_nullifiers: Vec<_> = block
Expand All @@ -224,83 +230,72 @@ impl State {
return Err(ApplyBlockError::DuplicatedNullifiers(duplicate_nullifiers));
}

// update the in-memory data structures and compute the new block header. Important, the
// structures are not yet committed
// compute updates for the in-memory data structures

// update chain MMR
let chain_mmr = {
let mut chain_mmr = inner.chain_mmr.clone();

// new_block.chain_root must be equal to the chain MMR root prior to the update
let peaks = chain_mmr.peaks();
if peaks.hash_peaks() != header.chain_root() {
return Err(ApplyBlockError::NewBlockInvalidChainRoot);
}

chain_mmr.add(block_hash);
chain_mmr
};

// update nullifier tree
let nullifier_tree = {
let mut nullifier_tree = inner.nullifier_tree.clone();
for nullifier in block.nullifiers() {
nullifier_tree
.insert(nullifier, block_num)
.map_err(ApplyBlockError::FailedToUpdateNullifierTree)?;
}
// new_block.chain_root must be equal to the chain MMR root prior to the update
let peaks = inner.chain_mmr.peaks();
if peaks.hash_peaks() != header.chain_root() {
return Err(ApplyBlockError::NewBlockInvalidChainRoot);
}

if nullifier_tree.root() != header.nullifier_root() {
return Err(ApplyBlockError::NewBlockInvalidNullifierRoot);
}
nullifier_tree
};
// compute update for nullifier tree
let nullifier_tree_update = inner.nullifier_tree.compute_mutations(
block.nullifiers().iter().map(|nullifier| (*nullifier, block_num)),
);

// update account tree
let mut account_tree = inner.account_tree.clone();
for update in block.updated_accounts() {
account_tree.insert(
LeafIndex::new_max_depth(update.account_id().into()),
update.new_state_hash().into(),
);
if nullifier_tree_update.root() != header.nullifier_root() {
return Err(ApplyBlockError::NewBlockInvalidNullifierRoot);
}

if account_tree.root() != header.account_root() {
// compute update for account tree
let account_tree_update = inner.account_tree.compute_mutations(
block.updated_accounts().iter().map(|update| {
(
LeafIndex::new_max_depth(update.account_id().into()),
update.new_state_hash().into(),
)
}),
);

if account_tree_update.root() != header.account_root() {
return Err(ApplyBlockError::NewBlockInvalidAccountRoot);
}

// build note tree
let note_tree = block.build_note_tree();
if note_tree.root() != header.note_root() {
return Err(ApplyBlockError::NewBlockInvalidNoteRoot);
}
(
inner.nullifier_tree.root(),
nullifier_tree_update,
inner.account_tree.root(),
account_tree_update,
)
};

drop(span);

let notes = block
.notes()
.map(|(note_index, note)| {
let details = match note {
OutputNote::Full(note) => Some(note.to_bytes()),
OutputNote::Header(_) => None,
note => return Err(ApplyBlockError::InvalidOutputNoteType(note.clone())),
};

let merkle_path = note_tree.get_note_path(note_index);

Ok(NoteRecord {
block_num,
note_index,
note_id: note.id().into(),
metadata: *note.metadata(),
details,
merkle_path,
})
})
.collect::<Result<Vec<NoteRecord>, ApplyBlockError>>()?;
// build note tree
let note_tree = block.build_note_tree();
if note_tree.root() != header.note_root() {
return Err(ApplyBlockError::NewBlockInvalidNoteRoot);
}

(account_tree, chain_mmr, nullifier_tree, notes)
};
let notes = block
.notes()
.map(|(note_index, note)| {
let details = match note {
OutputNote::Full(note) => Some(note.to_bytes()),
OutputNote::Header(_) => None,
note => return Err(ApplyBlockError::InvalidOutputNoteType(note.clone())),
};

let merkle_path = note_tree.get_note_path(note_index);

Ok(NoteRecord {
block_num,
note_index,
note_id: note.id().into(),
metadata: *note.metadata(),
details,
merkle_path,
})
})
.collect::<Result<Vec<NoteRecord>, ApplyBlockError>>()?;

// Signals the transaction is ready to be committed, and the write lock can be acquired
let (allow_acquire, acquired_allowed) = oneshot::channel::<()>();
Expand All @@ -311,7 +306,7 @@ impl State {
// overlapping. Namely, the DB transaction only proceeds after this task acquires the
// in-memory write lock. This requires the DB update to run concurrently, so a new task is
// spawned.
let db = self.db.clone();
let db = Arc::clone(&self.db);
let db_update_task =
tokio::spawn(
async move { db.apply_block(allow_acquire, acquire_done, block, notes).await },
Expand All @@ -332,6 +327,16 @@ impl State {
// successfully.
let mut inner = self.inner.write().await;

// We need to check that neither the nullifier tree nor the account tree have changed
// while we were waiting for the DB preparation task to complete. If either of them
// did change, we do not proceed with in-memory and database updates, since it may
// lead to an inconsistent state.
if inner.nullifier_tree.root() != nullifier_tree_old_root
|| inner.account_tree.root() != account_tree_old_root
{
return Err(ApplyBlockError::ConcurrentWrite);
}

// Notify the DB update task that the write lock has been acquired, so it can commit
// the DB transaction
let _ = inform_acquire_done.send(());
Expand All @@ -343,9 +348,15 @@ impl State {
db_update_task.await??;

// Update the in-memory data structures after successful commit of the DB transaction
inner.chain_mmr = chain_mmr;
inner.nullifier_tree = nullifier_tree;
inner.account_tree = account_tree;
inner
.nullifier_tree
.apply_mutations(nullifier_tree_update)
.expect("Unreachable: old nullifier tree root must be checked before this step");
inner
.account_tree
.apply_mutations(account_tree_update)
.expect("Unreachable: old account tree root must be checked before this step");
inner.chain_mmr.add(block_hash);
}

info!(%block_hash, block_num, COMPONENT, "apply_block successful");
Expand Down
Loading