feat: Lock-free apply_block refactor

Cargo.toml

@@ -66,6 +66,7 @@ miden-crypto = { version = "0.28" }
 
 # External dependencies
 anyhow             = { version = "1.0" }
+arc-swap           = { version = "1.7" }
 assert_matches     = { version = "1.5" }
 axum               = { version = "0.8" }
 backon             = { version = "1.6" }

bin/node/src/commands/modes.rs

@@ -12,7 +12,7 @@ use miden_node_proto::clients::{
     ValidatorClient,
 };
 use miden_node_rpc::{PreAuthSubmission, Rpc, RpcMode, SequencerInternal, ValidatorClients};
-use miden_node_store::State;
+use miden_node_store::{BlockWriter, ProofWriter, State, WriterTask};
 use miden_node_utils::clap::{GrpcOptionsInternal, duration_to_human_readable_string};
 use miden_node_utils::formatting::format_endpoint;
 use miden_node_utils::shutdown::CancellationToken;
@@ -57,11 +57,14 @@ impl SequencerCommand {
         let runtime = self.runtime.runtime_config(&self.store);
         self.block_producer.validate()?;
         let network_tx_auth = self.runtime.rpc.network_tx_auth()?;
-        let state = load_state(&runtime).await?;
+        let (state, block_writer, proof_writer, writer_task) =
+            load_state(&runtime, shutdown.clone()).await?;
         let _disk_monitor = state.spawn_disk_monitor(shutdown.clone());
 
         let sequencer = Sequencer {
-            store: Arc::clone(&state),
+            state: Arc::clone(&state),
+            block_writer,
+            proof_writer,
             validator_urls: self.external_services.validator_urls.clone(),
             validator_timeout: self.external_services.validator_timeout,
             batch_prover_url: self.block_producer.batch.prover_url,
@@ -75,24 +78,25 @@ impl SequencerCommand {
             batch_workers: self.block_producer.batch.workers,
         }
         .spawn(shutdown.clone())
-        .await
         .context("failed to spawn sequencer")?;
         let block_producer = sequencer.api();
 
         let rpc = Rpc {
             listener: bind_rpc(runtime.rpc_listen).await?,
-            store: state,
+            state,
             mode: RpcMode::sequencer(
                 block_producer.clone(),
                 self.external_services.validator_clients()?,
             ),
+            sync_writers: None,
             ntx_builder: Some(self.external_services.ntx_builder_client()?),
             grpc_options: runtime.external_grpc_options,
             network_tx_auth,
         };
         let mut tasks = Tasks::new();
         tasks.spawn("sequencer", sequencer.wait());
         tasks.spawn("RPC server", rpc.serve(shutdown.clone()));
+        tasks.spawn("store block writer", join_store_writer(writer_task));
         if let Some(internal_listen) = self.internal {
             let sequencer_internal = SequencerInternal {
                 listener: bind_rpc(internal_listen).await?,
@@ -240,19 +244,22 @@ impl FullNodeCommand {
         let source_rpc = self.sync.source_rpc_client()?;
         let pre_auth = self.pre_auth_submission()?;
         let network_tx_auth = self.runtime.rpc.network_tx_auth()?;
-        let state = load_state(&runtime).await?;
+        let (state, block_writer, proof_writer, writer_task) =
+            load_state(&runtime, shutdown.clone()).await?;
         let _disk_monitor = state.spawn_disk_monitor(shutdown.clone());
 
         let rpc = Rpc {
             listener: bind_rpc(runtime.rpc_listen).await?,
-            store: state,
+            state,
             mode: RpcMode::full_node(source_rpc, self.sync.readiness_threshold, pre_auth),
+            sync_writers: Some((block_writer, proof_writer)),
             ntx_builder: None,
             grpc_options: runtime.external_grpc_options,
             network_tx_auth,
         };
         let mut tasks = Tasks::new();
         tasks.spawn("RPC server", rpc.serve(shutdown.clone()));
+        tasks.spawn("store block writer", join_store_writer(writer_task));
 
         tasks.join_next_or_cancelled(shutdown).await
     }
@@ -326,16 +333,29 @@ impl SyncOptions {
     }
 }
 
-async fn load_state(runtime: &RuntimeConfig) -> anyhow::Result<Arc<State>> {
-    let state = State::load_with_database_options(
+async fn load_state(
+    runtime: &RuntimeConfig,
+    shutdown: CancellationToken,
+) -> anyhow::Result<(Arc<State>, BlockWriter, ProofWriter, WriterTask)> {
+    let loaded = State::load_with_database_options(
         &runtime.data_directory,
         runtime.storage_options.clone(),
         runtime.database_options,
+        shutdown,
     )
     .await
     .context("failed to load state")?;
 
-    Ok(Arc::new(state))
+    Ok(loaded.start())
+}
+
+/// Supervises the store's write worker task.
+///
+/// On shutdown the task-drain loop waits for the writer to finish any in-flight block write and
+/// close its storage; an early exit or panic surfaces through the task set like any other task
+/// failure.
+async fn join_store_writer(writer_task: WriterTask) -> anyhow::Result<()> {
+    writer_task.await.map_err(anyhow::Error::from)
 }
 
 async fn bind_rpc(listen: SocketAddr) -> anyhow::Result<TcpListener> {

bin/node/src/commands/recover.rs

@@ -5,8 +5,8 @@ use std::time::Duration;
 use anyhow::Context;
 use miden_node_proto::clients::{Builder, ValidatorClient};
 use miden_node_proto::generated::validator::BlockSubscriptionRequest;
-use miden_node_store::State;
-use miden_node_store::state::Finality;
+use miden_node_store::{BlockWriter, State, WriterTask};
+use miden_node_utils::shutdown::CancellationToken;
 use miden_protocol::block::{BlockNumber, SignedBlock};
 use miden_protocol::utils::serde::Deserializable;
 use tokio_stream::StreamExt;
@@ -36,21 +36,28 @@ pub struct RecoverCommand {
 
 impl RecoverCommand {
     pub async fn handle(self) -> anyhow::Result<()> {
-        let state = self.load_state().await?;
+        let (state, block_writer, writer_task) = self.load_state().await?;
         let validator = self.validator_client()?;
-        recover_from_validator(&state, validator).await
+        let result = recover_from_validator(&state, &block_writer, validator).await;
+        // Wait for the writer to drain and release the backing storage before the process exits.
+        block_writer.stop(writer_task).await;
+        result
     }
 
-    async fn load_state(&self) -> anyhow::Result<Arc<State>> {
-        let state = State::load_with_database_options(
+    async fn load_state(&self) -> anyhow::Result<(Arc<State>, BlockWriter, WriterTask)> {
+        // Recovery is not wired into the node's shutdown token; the writer exits once the
+        // `BlockWriter` (holding the only write handle) is dropped after recovery completes.
+        let loaded = State::load_with_database_options(
             &self.data_directory,
             self.store.storage.clone().into(),
             self.store.sqlite.database_options(),
+            CancellationToken::new(),
         )
         .await
         .context("failed to load state")?;
 
-        Ok(Arc::new(state))
+        let (state, block_writer, _proof_writer, writer_task) = loaded.start();
+        Ok((state, block_writer, writer_task))
     }
 
     fn validator_client(&self) -> anyhow::Result<ValidatorClient> {
@@ -66,7 +73,8 @@ impl RecoverCommand {
 
 /// Streams blocks from the validator into the local store until the chain tip is reached.
 async fn recover_from_validator(
-    state: &Arc<State>,
+    state: &State,
+    block_writer: &BlockWriter,
     mut validator: ValidatorClient,
 ) -> anyhow::Result<()> {
     // Capture the validator's chain tip as the recovery target. The validator's block stream
@@ -81,7 +89,7 @@ async fn recover_from_validator(
             .chain_tip,
     );
 
-    let local_tip = state.chain_tip(Finality::Committed).await;
+    let local_tip = state.committed_tip();
     if local_tip >= validator_tip {
         info!(
             target: LOG_TARGET,
@@ -111,7 +119,10 @@ async fn recover_from_validator(
         let block = SignedBlock::read_from_bytes(&event.block)
             .context("failed to deserialize block from validator")?;
         let block_num = block.header().block_num();
-        state.apply_block(block).await.context("failed to apply recovered block")?;
+        block_writer
+            .apply_block(block)
+            .await
+            .context("failed to apply recovered block")?;
         info!(target: LOG_TARGET, block_number = %block_num.as_u32(), "Applied recovered block");
 
         // Stop once we reach the tip captured at the start of recovery.
@@ -121,7 +132,7 @@ async fn recover_from_validator(
     }
 
     // The stream can end before reaching the tip if the validator restarts or drops the connection.
-    let final_tip = state.chain_tip(Finality::Committed).await;
+    let final_tip = state.committed_tip();
     anyhow::ensure!(
         final_tip >= validator_tip,
         "validator block stream ended at block {} before reaching the chain tip {}",