fix(store/state): perform potentially blocking RocksDB access in a blocking context

CHANGELOG.md

@@ -24,6 +24,7 @@
 - [BREAKING] Renamed `--url` CLI flags and `*_URL` env vars to `--listen` / `*_LISTEN` across all components.
 - [BREAKING] Removed `miden-node validator` subcommand and created a separate `miden-validator` binary ([#2053](https://github.com/0xMiden/node/pull/2053)).
 - [BREAKING] Removed `miden-node ntx-builder` subcommand and created a separate `miden-ntx-builder` binary ([#2067](https://github.com/0xMiden/node/pull/2067)).
+- Replaced blocking-in-async LargeSmt and account state forest operations in the store with wrappers using Tokio's `block_in_place()` ([#2076](https://github.com/0xMiden/node/pull/2076)).
 - [BREAKING] Reworked note proto types for multi-attachment support: `NoteMetadata` now carries `attachment_schemes` (repeated) and `attachments_commitment` instead of a single `attachment`. `Note` and `NetworkNote` gained an `attachments` field. `NoteSyncRecord` now embeds full `NoteMetadata` instead of `NoteMetadataHeader`. Removed `NoteAttachmentKind` enum and `NoteMetadataHeader` message ([#2078](https://github.com/0xMiden/node/pull/2078)).
 - [BREAKING] Changed `SyncChainMmr` endpoint: the upper end of the block range we're syncing is now the chain tip with the requested finality level. Validator signature is also returned ([#2075](https://github.com/0xMiden/node/pull/2075)).
 

crates/store/src/server/ntx_builder.rs

@@ -255,7 +255,6 @@ impl ntx_builder_server::NtxBuilder for StoreApi {
         let asset_witnesses = self
             .state
             .get_vault_asset_witnesses(account_id, block_num, vault_keys)
-            .await
             .map_err(internal_error)?;
 
         // Convert AssetWitness to protobuf format by extracting witness data.
@@ -309,7 +308,6 @@ impl ntx_builder_server::NtxBuilder for StoreApi {
         let storage_witness = self
             .state
             .get_storage_map_witness(account_id, &slot_name, block_num, map_key)
-            .await
             .map_err(internal_error)?;
 
         // Convert StorageMapWitness to protobuf format by extracting witness data.

crates/store/src/state/apply_block.rs

@@ -123,17 +123,7 @@ impl State {
         // Awaiting the block saving task to complete without errors.
         block_save_task.await??;
 
-        // Scope to update the in-memory data.
-        async move {
-            // We need to hold the write lock here to prevent inconsistency between the in-memory
-            // state and the DB state. Thus, we need to wait for the DB update task to complete
-            // successfully.
-            let mut inner = self
-                .inner
-                .write()
-                .instrument(info_span!(target: COMPONENT, "acquire_inner_write_lock"))
-                .await;
-
+        self.with_inner_write_blocking(|inner| {
             // 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
@@ -154,8 +144,8 @@ impl State {
             // Await for successful commit of the DB transaction. If the commit fails, we mustn't
             // change in-memory state, so we return a block applying error and don't proceed with
             // in-memory updates.
-            db_update_task
-                .await?
+            tokio::runtime::Handle::current()
+                .block_on(db_update_task)?
                 .map_err(|err| ApplyBlockError::DbUpdateTaskFailed(err.as_report()))?;
 
             // Update the in-memory data structures after successful commit of the DB transaction
@@ -171,15 +161,11 @@ impl State {
             inner.blockchain.push(block_commitment);
 
             Ok(())
-        }
-        .in_current_span()
-        .await?;
+        })?;
 
-        self.forest
-            .write()
-            .instrument(info_span!(target: COMPONENT, "acquire_forest_write_lock"))
-            .await
-            .apply_block_updates(block_num, account_deltas)?;
+        self.with_forest_write_blocking(|forest| {
+            forest.apply_block_updates(block_num, account_deltas)
+        })?;
 
         // Push to cache and notify replica subscribers.
         self.block_cache.push(block_num, BlockNotification::new(block_num, cache_bytes));
@@ -237,78 +223,74 @@ impl State {
         header: &BlockHeader,
         body: &BlockBody,
     ) -> Result<(Word, NullifierMutationSet, Word, AccountMutationSet), ApplyBlockError> {
-        let inner = self
-            .inner
-            .read()
-            .instrument(info_span!(target: COMPONENT, "acquire_inner_read_lock"))
-            .await;
-
-        let block_num = header.block_num();
-
-        // nullifiers can be produced only once
-        let duplicate_nullifiers: Vec<_> = body
-            .created_nullifiers()
-            .iter()
-            .filter(|&nullifier| inner.nullifier_tree.get_block_num(nullifier).is_some())
-            .copied()
-            .collect();
-        if !duplicate_nullifiers.is_empty() {
-            return Err(InvalidBlockError::DuplicatedNullifiers(duplicate_nullifiers).into());
-        }
-
-        // new_block.chain_root must be equal to the chain MMR root prior to the update
-        let peaks = inner.blockchain.peaks();
-        if peaks.hash_peaks() != header.chain_commitment() {
-            return Err(InvalidBlockError::NewBlockInvalidChainCommitment.into());
-        }
+        self.with_inner_read_blocking(|inner| {
+            let block_num = header.block_num();
+
+            // nullifiers can be produced only once
+            let duplicate_nullifiers: Vec<_> = body
+                .created_nullifiers()
+                .iter()
+                .filter(|&nullifier| inner.nullifier_tree.get_block_num(nullifier).is_some())
+                .copied()
+                .collect();
+            if !duplicate_nullifiers.is_empty() {
+                return Err(InvalidBlockError::DuplicatedNullifiers(duplicate_nullifiers).into());
+            }
 
-        // compute update for nullifier tree
-        let nullifier_tree_update = inner
-            .nullifier_tree
-            .compute_mutations(
-                body.created_nullifiers().iter().map(|nullifier| (*nullifier, block_num)),
-            )
-            .map_err(InvalidBlockError::NewBlockNullifierAlreadySpent)?;
-
-        if nullifier_tree_update.as_mutation_set().root() != header.nullifier_root() {
-            // We do our best here to notify the serve routine, if it doesn't care (dropped the
-            // receiver) we can't do much.
-            let _ = self.termination_ask.try_send(ApplyBlockError::InvalidBlockError(
-                InvalidBlockError::NewBlockInvalidNullifierRoot,
-            ));
-            return Err(InvalidBlockError::NewBlockInvalidNullifierRoot.into());
-        }
+            // new_block.chain_root must be equal to the chain MMR root prior to the update
+            let peaks = inner.blockchain.peaks();
+            if peaks.hash_peaks() != header.chain_commitment() {
+                return Err(InvalidBlockError::NewBlockInvalidChainCommitment.into());
+            }
 
-        // compute update for account tree
-        let account_tree_update = inner
-            .account_tree
-            .compute_mutations(
-                body.updated_accounts()
-                    .iter()
-                    .map(|update| (update.account_id(), update.final_state_commitment())),
-            )
-            .map_err(|e| match e {
-                HistoricalError::AccountTreeError(err) => {
-                    InvalidBlockError::NewBlockDuplicateAccountIdPrefix(err)
-                },
-                HistoricalError::MerkleError(_) => {
-                    panic!("Unexpected MerkleError during account tree mutation computation")
-                },
-            })?;
+            // compute update for nullifier tree
+            let nullifier_tree_update = inner
+                .nullifier_tree
+                .compute_mutations(
+                    body.created_nullifiers().iter().map(|nullifier| (*nullifier, block_num)),
+                )
+                .map_err(InvalidBlockError::NewBlockNullifierAlreadySpent)?;
+
+            if nullifier_tree_update.as_mutation_set().root() != header.nullifier_root() {
+                // We do our best here to notify the serve routine, if it doesn't care (dropped the
+                // receiver) we can't do much.
+                let _ = self.termination_ask.try_send(ApplyBlockError::InvalidBlockError(
+                    InvalidBlockError::NewBlockInvalidNullifierRoot,
+                ));
+                return Err(InvalidBlockError::NewBlockInvalidNullifierRoot.into());
+            }
 
-        if account_tree_update.as_mutation_set().root() != header.account_root() {
-            let _ = self.termination_ask.try_send(ApplyBlockError::InvalidBlockError(
-                InvalidBlockError::NewBlockInvalidAccountRoot,
-            ));
-            return Err(InvalidBlockError::NewBlockInvalidAccountRoot.into());
-        }
+            // compute update for account tree
+            let account_tree_update = inner
+                .account_tree
+                .compute_mutations(
+                    body.updated_accounts()
+                        .iter()
+                        .map(|update| (update.account_id(), update.final_state_commitment())),
+                )
+                .map_err(|e| match e {
+                    HistoricalError::AccountTreeError(err) => {
+                        InvalidBlockError::NewBlockDuplicateAccountIdPrefix(err)
+                    },
+                    HistoricalError::MerkleError(_) => {
+                        panic!("Unexpected MerkleError during account tree mutation computation")
+                    },
+                })?;
+
+            if account_tree_update.as_mutation_set().root() != header.account_root() {
+                let _ = self.termination_ask.try_send(ApplyBlockError::InvalidBlockError(
+                    InvalidBlockError::NewBlockInvalidAccountRoot,
+                ));
+                return Err(InvalidBlockError::NewBlockInvalidAccountRoot.into());
+            }
 
-        Ok((
-            inner.nullifier_tree.root(),
-            nullifier_tree_update,
-            inner.account_tree.root_latest(),
-            account_tree_update,
-        ))
+            Ok((
+                inner.nullifier_tree.root(),
+                nullifier_tree_update,
+                inner.account_tree.root_latest(),
+                account_tree_update,
+            ))
+        })
     }
 
     /// Builds note records with inclusion proofs from the block body.