1use 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
45const BLOCK_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(512).unwrap();
47
48const PROOF_CACHE_CAPACITY: NonZeroUsize = NonZeroUsize::new(512).unwrap();
50
51#[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 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
99impl State {
103 #[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 #[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_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 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 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 let snapshots_live = Arc::new(AtomicUsize::new(0));
202
203 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 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 #[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}