Skip to main content

miden_node_store/state/
lifecycle.rs

1//! Store lifecycle: loading the state, starting its write worker, and stopping the store.
2
3use std::num::NonZeroUsize;
4use std::path::Path;
5use std::sync::Arc;
6use std::sync::atomic::AtomicUsize;
7
8use arc_swap::ArcSwap;
9use miden_node_utils::ErrorReport;
10use miden_node_utils::clap::StorageOptions;
11use miden_node_utils::shutdown::CancellationToken;
12use miden_node_utils::tracing::miden_instrument;
13use miden_protocol::block::BlockNumber;
14use tokio::sync::{mpsc, watch};
15
16use crate::account_state_forest::AccountStateForestBackend;
17use crate::accounts::AccountTreeWithHistory;
18use crate::blocks::BlockStore;
19use crate::db::Db;
20use crate::errors::StateInitializationError;
21use crate::proven_tip::ProvenTipWriter;
22use crate::state::loader::{
23    ACCOUNT_STATE_FOREST_STORAGE_DIR,
24    ACCOUNT_TREE_STORAGE_DIR,
25    AccountForestLoader,
26    NULLIFIER_TREE_STORAGE_DIR,
27    TreeStorage,
28    TreeStorageLoader,
29    load_mmr,
30    verify_account_state_forest_consistency,
31    verify_tree_consistency,
32};
33use crate::state::writer::{WriteRequest, WriteWorker, WriterTask};
34use crate::state::{
35    BlockCache,
36    BlockWriter,
37    ProofCache,
38    ProofWriter,
39    SnapshotGuard,
40    State,
41    StateSnapshot,
42};
43use crate::{COMPONENT, DataDirectory, DatabaseOptions};
44
45/// Number of recent committed blocks held in the in-memory cache for replica subscriptions.
46const BLOCK_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(512).unwrap();
47
48/// Number of recent block proofs held in the in-memory cache for replica subscriptions.
49const PROOF_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(512).unwrap();
50
51// LOADED STATE
52// ================================================================================================
53
54/// A loaded store state whose write worker has not been started yet.
55///
56/// Returned by [`State::load`]. [`Self::start`] spawns the writer and yields the read-only
57/// [`State`] together with the write capabilities; since this is the only way to obtain a
58/// [`BlockWriter`], [`BlockWriter::apply_block`] can always make progress.
59#[must_use = "call `start` to spawn the write worker and obtain the state"]
60pub struct LoadedState {
61    state: State,
62    writer: WriteWorker,
63    write_tx: mpsc::Sender<WriteRequest>,
64}
65
66impl LoadedState {
67    /// Spawns the write worker onto the current runtime and returns the read-only state together
68    /// with the write capabilities and the writer task's handle.
69    ///
70    /// [`Arc<State>`] is the read-only view shared with every component that queries or
71    /// subscribes; it exposes no mutating methods. [`BlockWriter`] and [`ProofWriter`] are the
72    /// only handles able to mutate the store — hand them to the single task driving each write
73    /// path (block production or sync, and proof scheduling or sync respectively). The
74    /// capabilities expose no read access: tasks that both read and write receive the
75    /// [`Arc<State>`] alongside their capability.
76    ///
77    /// The writer exits once `shutdown` is cancelled or the [`BlockWriter`] (holding the only
78    /// request sender) is dropped — an in-flight block write always completes first. Awaiting the
79    /// returned handle after either event guarantees the writer has released the tree storage it
80    /// owns; a join error carries a writer panic.
81    ///
82    /// Callers without a token to cancel should pass [`CancellationToken::new`] and stop the store
83    /// via [`BlockWriter::stop`] instead of dropping and joining by hand.
84    pub fn start(
85        self,
86        shutdown: CancellationToken,
87    ) -> (Arc<State>, BlockWriter, ProofWriter, WriterTask) {
88        let writer_task = tokio::spawn(self.writer.run(shutdown));
89        let state = Arc::new(self.state);
90        let block_writer = BlockWriter {
91            block_store: Arc::clone(&state.block_store),
92            write_tx: self.write_tx,
93        };
94        let proof_writer = ProofWriter { state: Arc::clone(&state) };
95        (state, block_writer, proof_writer, WriterTask(writer_task))
96    }
97}
98
99// LOAD
100// ================================================================================================
101
102impl State {
103    /// Loads the state from the data directory.
104    ///
105    /// The loaded state owns all store data structures and exposes subscription methods for
106    /// sequencer and replica tasks. Call [`LoadedState::start`] on the result to spawn the block
107    /// writer and obtain the usable [`State`].
108    #[miden_instrument(
109        target = COMPONENT,
110    )]
111    pub async fn load(
112        data_path: &Path,
113        storage_options: StorageOptions,
114    ) -> Result<LoadedState, StateInitializationError> {
115        Self::load_with_database_options(data_path, storage_options, DatabaseOptions::default())
116            .await
117    }
118
119    /// Loads the state from the data directory using explicit database options.
120    ///
121    /// The loaded state owns all store data structures and exposes subscription methods for
122    /// sequencer and replica tasks. Call [`LoadedState::start`] on the result to spawn the block
123    /// writer and obtain the usable [`State`].
124    #[miden_instrument(
125        target = COMPONENT,
126    )]
127    pub async fn load_with_database_options(
128        data_path: &Path,
129        storage_options: StorageOptions,
130        database_options: DatabaseOptions,
131    ) -> Result<LoadedState, StateInitializationError> {
132        let data_directory = DataDirectory::load(data_path.to_path_buf())
133            .map_err(StateInitializationError::DataDirectoryLoadError)?;
134
135        let block_store = Arc::new(
136            BlockStore::load(data_directory.block_store_dir())
137                .map_err(StateInitializationError::BlockStoreLoadError)?,
138        );
139
140        let database_filepath = data_directory.database_path();
141        let mut db = Db::load_with_pool_size(
142            database_filepath.clone(),
143            database_options.connection_pool_size,
144        )
145        .await
146        .map_err(StateInitializationError::DatabaseLoadError)?;
147
148        let blockchain = load_mmr(&mut db).await?;
149        let latest_block_num = blockchain.chain_tip().unwrap_or(BlockNumber::GENESIS);
150
151        let apply_block_thread_priority = storage_options.apply_block_thread_priority;
152
153        #[cfg(feature = "rocksdb")]
154        let (account_storage_config, nullifier_storage_config, forest_storage_config) = (
155            storage_options.account_tree.into(),
156            storage_options.nullifier_tree.into(),
157            storage_options.account_state_forest.into(),
158        );
159        #[cfg(not(feature = "rocksdb"))]
160        let (account_storage_config, nullifier_storage_config, forest_storage_config) =
161            ((), (), ());
162        let account_storage =
163            TreeStorage::create(data_path, &account_storage_config, ACCOUNT_TREE_STORAGE_DIR)?;
164        let account_tree = account_storage.load_account_tree(&mut db).await?;
165
166        let nullifier_storage =
167            TreeStorage::create(data_path, &nullifier_storage_config, NULLIFIER_TREE_STORAGE_DIR)?;
168        let nullifier_tree = nullifier_storage.load_nullifier_tree(&mut db).await?;
169
170        // Verify that tree roots match the expected roots from the database. This catches any
171        // divergence between persistent storage and the database caused by corruption or incomplete
172        // shutdown.
173        verify_tree_consistency(account_tree.root(), nullifier_tree.root(), &mut db).await?;
174
175        let account_tree = AccountTreeWithHistory::new(account_tree, latest_block_num);
176
177        let forest_backend = AccountStateForestBackend::create(
178            data_path,
179            &forest_storage_config,
180            ACCOUNT_STATE_FOREST_STORAGE_DIR,
181        )?;
182        let forest = forest_backend.load_account_state_forest(&mut db, latest_block_num).await?;
183        verify_account_state_forest_consistency(&forest, &mut db).await?;
184
185        let db = Arc::new(db);
186
187        // Initialize the proven tip from the block store.
188        let proven_tip_init = block_store
189            .load_proven_tip()
190            .map_err(StateInitializationError::ProvenTipLoadError)?;
191        let (proven_tip, _rx) = ProvenTipWriter::new(proven_tip_init);
192
193        // Committed-tip watch: fires after each successful apply_block.
194        let (committed_tip_tx, _rx) = watch::channel(latest_block_num);
195        let committed_tip_tx = Arc::new(committed_tip_tx);
196
197        let block_cache = BlockCache::new(BLOCK_CACHE_CAPACITY);
198        let proof_cache = ProofCache::new(PROOF_CACHE_CAPACITY);
199
200        // Shared counter of live snapshot generations, for observability.
201        let snapshots_live = Arc::new(AtomicUsize::new(0));
202
203        // Create the initial snapshot from reader views of the just-loaded trees.
204        let initial_snapshot = Arc::new(StateSnapshot::new(
205            nullifier_tree
206                .reader()
207                .map_err(|e| StateInitializationError::NullifierTreeIoError(e.as_report()))?,
208            blockchain.clone(),
209            account_tree.reader(),
210            forest
211                .reader()
212                .map_err(|e| StateInitializationError::AccountStateForestIoError(e.as_report()))?,
213            SnapshotGuard::new(Arc::clone(&snapshots_live), latest_block_num),
214        ));
215        let latest_snapshot = Arc::new(ArcSwap::from(initial_snapshot));
216
217        // Assemble the write worker. It owns the writable trees and processes write requests
218        // serially, publishing a new snapshot after each committed block. The caller runs it; it
219        // exits when the shutdown token is cancelled or the `BlockWriter` (holding the only request
220        // sender) is dropped.
221        let (write_tx, write_rx) = mpsc::channel(1);
222        let block_writer = WriteWorker::new(
223            Arc::clone(&db),
224            Arc::clone(&block_store),
225            Arc::clone(&latest_snapshot),
226            Arc::clone(&committed_tip_tx),
227            block_cache.clone(),
228            write_rx,
229            nullifier_tree,
230            account_tree,
231            blockchain,
232            forest,
233            snapshots_live,
234            apply_block_thread_priority,
235        );
236        let state = Self {
237            data_directory: data_path.to_path_buf(),
238            db,
239            block_store,
240            latest_snapshot,
241            proven_tip,
242            committed_tip_tx,
243            block_cache,
244            proof_cache,
245        };
246
247        Ok(LoadedState { state, writer: block_writer, write_tx })
248    }
249
250    /// Loads the state with default options and starts its write worker, detaching the worker
251    /// task.
252    ///
253    /// Test-only helper for tests in sibling crates that don't manage the writer's lifecycle:
254    /// the detached writer exits once the returned [`BlockWriter`] is dropped. Hidden from public
255    /// docs and not part of the stable API.
256    ///
257    /// # Panics
258    ///
259    /// Panics if the state fails to load.
260    #[doc(hidden)]
261    pub async fn for_tests(data_path: &Path) -> (Arc<Self>, BlockWriter, ProofWriter) {
262        let (state, block_writer, proof_writer, _writer_task) =
263            Self::load(data_path, StorageOptions::default())
264                .await
265                .expect("state should load")
266                .start(CancellationToken::new());
267        (state, block_writer, proof_writer)
268    }
269}