feat(en): Support Merkle tree recovery with pruning enabled #3172
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -32,16 +32,14 @@ use std::{ | |
| }; | ||
|
|
||
| use anyhow::Context as _; | ||
| use async_trait::async_trait; | ||
| use futures::future; | ||
| use tokio::sync::{watch, Mutex, Semaphore}; | ||
| use zksync_dal::{Connection, ConnectionPool, Core, CoreDal}; | ||
| use zksync_health_check::HealthUpdater; | ||
| use zksync_merkle_tree::TreeEntry; | ||
| use zksync_shared_metrics::{SnapshotRecoveryStage, APP_METRICS}; | ||
| use zksync_types::{snapshots::uniform_hashed_keys_chunk, L1BatchNumber, L2BlockNumber, H256}; | ||
|
|
||
| use super::{ | ||
| helpers::{AsyncTree, AsyncTreeRecovery, GenericAsyncTree, MerkleTreeHealth}, | ||
|
|
@@ -54,12 +52,13 @@ mod tests; | |
|
|
||
| /// Handler of recovery life cycle events. This functionality is encapsulated in a trait to be able | ||
| /// to control recovery behavior in tests. | ||
| #[async_trait] | ||
| trait HandleRecoveryEvent: fmt::Debug + Send + Sync { | ||
| fn recovery_started(&mut self, _chunk_count: u64, _recovered_chunk_count: u64) { | ||
| // Default implementation does nothing | ||
| } | ||
|
|
||
| async fn chunk_recovered(&self) { | ||
| // Default implementation does nothing | ||
| } | ||
| } | ||
|
|
@@ -82,6 +81,7 @@ impl<'a> RecoveryHealthUpdater<'a> { | |
| } | ||
| } | ||
|
|
||
| #[async_trait] | ||
| impl HandleRecoveryEvent for RecoveryHealthUpdater<'_> { | ||
| fn recovery_started(&mut self, chunk_count: u64, recovered_chunk_count: u64) { | ||
| self.chunk_count = chunk_count; | ||
|
|
@@ -91,7 +91,7 @@ impl HandleRecoveryEvent for RecoveryHealthUpdater<'_> { | |
| .set(recovered_chunk_count); | ||
| } | ||
|
|
||
| async fn chunk_recovered(&self) { | ||
| let recovered_chunk_count = self.recovered_chunk_count.fetch_add(1, Ordering::SeqCst) + 1; | ||
| let chunks_left = self.chunk_count.saturating_sub(recovered_chunk_count); | ||
| tracing::info!( | ||
|
|
@@ -110,34 +110,68 @@ impl HandleRecoveryEvent for RecoveryHealthUpdater<'_> { | |
| } | ||
|
|
||
| #[derive(Debug, Clone, Copy)] | ||
| struct InitParameters { | ||
| l1_batch: L1BatchNumber, | ||
| l2_block: L2BlockNumber, | ||
| expected_root_hash: Option<H256>, | ||
| log_count: u64, | ||
| desired_chunk_size: u64, | ||
| } | ||
|
|
||
| impl InitParameters { | ||
| async fn new( | ||
| pool: &ConnectionPool<Core>, | ||
| config: &MetadataCalculatorRecoveryConfig, | ||
| ) -> anyhow::Result<Option<Self>> { | ||
| let mut storage = pool.connection_tagged("metadata_calculator").await?; | ||
| let recovery_status = storage | ||
| .snapshot_recovery_dal() | ||
| .get_applied_snapshot_status() | ||
| .await?; | ||
| let pruning_info = storage.pruning_dal().get_pruning_info().await?; | ||
|
|
||
| let (l1_batch, l2_block); | ||
| let mut expected_root_hash = None; | ||
| match (recovery_status, pruning_info.last_hard_pruned_l2_block) { | ||
| (Some(recovery), None) => { | ||
| tracing::warn!( | ||
| "Snapshot recovery {recovery:?} is present on the node, but pruning info is empty; assuming no pruning happened" | ||
| ); | ||
| l1_batch = recovery.l1_batch_number; | ||
| l2_block = recovery.l2_block_number; | ||
| expected_root_hash = Some(recovery.l1_batch_root_hash); | ||
| } | ||
| (Some(recovery), Some(pruned_l2_block)) => { | ||
| // We have both recovery and some pruning on top of it. | ||
| l2_block = pruned_l2_block.max(recovery.l2_block_number); | ||
| l1_batch = pruning_info | ||
| .last_hard_pruned_l1_batch | ||
| .with_context(|| format!("malformed pruning info: {pruning_info:?}"))?; | ||
| if l1_batch == recovery.l1_batch_number { | ||
| expected_root_hash = Some(recovery.l1_batch_root_hash); | ||
| } | ||
| } | ||
| (None, Some(pruned_l2_block)) => { | ||
| l2_block = pruned_l2_block; | ||
| l1_batch = pruning_info | ||
| .last_hard_pruned_l1_batch | ||
| .with_context(|| format!("malformed pruning info: {pruning_info:?}"))?; | ||
| } | ||
| (None, None) => return Ok(None), | ||
| }; | ||
|
|
||
| let log_count = storage | ||
| .storage_logs_dal() | ||
| .get_storage_logs_row_count(l2_block) | ||
| .await?; | ||
|
|
||
| Ok(Some(Self { | ||
| l1_batch, | ||
| l2_block, | ||
| expected_root_hash, | ||
| log_count, | ||
| desired_chunk_size: config.desired_chunk_size, | ||
| })) | ||
| } | ||
|
|
||
| fn chunk_count(&self) -> u64 { | ||
|
|
@@ -168,47 +202,44 @@ impl GenericAsyncTree { | |
| stop_receiver: &watch::Receiver<bool>, | ||
| ) -> anyhow::Result<Option<AsyncTree>> { | ||
| let started_at = Instant::now(); | ||
| let (tree, init_params) = match self { | ||
| Self::Ready(tree) => return Ok(Some(tree)), | ||
| Self::Recovering(tree) => { | ||
| let params = InitParameters::new(main_pool, config).await?.context( | ||
| "Merkle tree is recovering, but Postgres doesn't contain snapshot recovery information", | ||
| )?; | ||
| let recovered_version = tree.recovered_version(); | ||
| anyhow::ensure!( | ||
| u64::from(params.l1_batch.0) == recovered_version, | ||
| "Snapshot L1 batch in Postgres ({params:?}) differs from the recovered Merkle tree version \ | ||
| ({recovered_version})" | ||
| ); | ||
| tracing::info!("Resuming tree recovery with status: {params:?}"); | ||
| (tree, params) | ||
| } | ||
| Self::Empty { db, mode } => { | ||
| if let Some(params) = InitParameters::new(main_pool, config).await? { | ||
| tracing::info!("Starting Merkle tree recovery with status {params:?}"); | ||
| let l1_batch = params.l1_batch; | ||
| let tree = AsyncTreeRecovery::new(db, l1_batch.0.into(), mode, config)?; | ||
| (tree, params) | ||
| } else { | ||
| // Start the tree from scratch. The genesis block will be filled in `TreeUpdater::loop_updating_tree()`. | ||
| return Ok(Some(AsyncTree::new(db, mode)?)); | ||
| } | ||
| } | ||
| }; | ||
|
|
||
| tracing::debug!( | ||
| "Obtained recovery init parameters: {init_params:?} based on recovery configuration {config:?}" | ||
| ); | ||
| let recovery_options = RecoveryOptions { | ||
| chunk_count: init_params.chunk_count(), | ||
| concurrency_limit: recovery_pool.max_size() as usize, | ||
| events: Box::new(RecoveryHealthUpdater::new(health_updater)), | ||
| }; | ||
| let tree = tree | ||
| .recover(init_params, recovery_options, &recovery_pool, stop_receiver) | ||
| .await?; | ||
| if tree.is_some() { | ||
| // Only report latency if recovery wasn't canceled | ||
|
|
@@ -223,12 +254,12 @@ impl GenericAsyncTree { | |
| impl AsyncTreeRecovery { | ||
| async fn recover( | ||
| mut self, | ||
| init_params: InitParameters, | ||
| mut options: RecoveryOptions<'_>, | ||
| pool: &ConnectionPool<Core>, | ||
| stop_receiver: &watch::Receiver<bool>, | ||
| ) -> anyhow::Result<Option<AsyncTree>> { | ||
| self.ensure_desired_chunk_size(init_params.desired_chunk_size) | ||
| .await?; | ||
|
|
||
| let start_time = Instant::now(); | ||
|
|
@@ -237,13 +268,15 @@ impl AsyncTreeRecovery { | |
| .map(|chunk_id| uniform_hashed_keys_chunk(chunk_id, chunk_count)) | ||
| .collect(); | ||
| tracing::info!( | ||
| "Recovering Merkle tree from Postgres snapshot in {chunk_count} chunks with max concurrency {}. \ | ||
| Be aware that enabling node pruning during recovery will probably result in a recovery error; always disable pruning \ | ||
| until recovery is complete", | ||
| options.concurrency_limit | ||
| ); | ||
|
|
||
| let mut storage = pool.connection_tagged("metadata_calculator").await?; | ||
| let remaining_chunks = self | ||
| .filter_chunks(&mut storage, init_params.l2_block, &chunks) | ||
| .await?; | ||
| drop(storage); | ||
| options | ||
|
|
@@ -261,9 +294,10 @@ impl AsyncTreeRecovery { | |
| .acquire() | ||
| .await | ||
| .context("semaphore is never closed")?; | ||
| if Self::recover_key_chunk(&tree, init_params.l2_block, chunk, pool, stop_receiver) | ||
| .await? | ||
| { | ||
| options.events.chunk_recovered().await; | ||
| } | ||
| anyhow::Ok(()) | ||
| }); | ||
|
|
@@ -279,13 +313,18 @@ impl AsyncTreeRecovery { | |
|
|
||
| let finalize_latency = RECOVERY_METRICS.latency[&RecoveryStage::Finalize].start(); | ||
| let actual_root_hash = tree.root_hash().await; | ||
| if let Some(expected_root_hash) = init_params.expected_root_hash { | ||
| anyhow::ensure!( | ||
| actual_root_hash == expected_root_hash, | ||
| "Root hash of recovered tree {actual_root_hash:?} differs from expected root hash {expected_root_hash:?}" | ||
| ); | ||
| } | ||
|
|
||
| // Check pruning info one last time before finalizing the tree. | ||
| let mut storage = pool.connection_tagged("metadata_calculator").await?; | ||
| Self::check_pruning_info(&mut storage, init_params.l2_block).await?; | ||
| drop(storage); | ||
|
|
||
| let tree = tree.finalize().await?; | ||
| finalize_latency.observe(); | ||
| tracing::info!( | ||
|
|
@@ -340,6 +379,21 @@ impl AsyncTreeRecovery { | |
| Ok(output) | ||
| } | ||
|
|
||
| async fn check_pruning_info( | ||
|
There was a problem hiding this comment.
There was a problem hiding this comment. WDYM by " if pruning is configured to never happen in practice"? There was a problem hiding this comment. For example, |
||
| storage: &mut Connection<'_, Core>, | ||
| snapshot_l2_block: L2BlockNumber, | ||
| ) -> anyhow::Result<()> { | ||
| let pruning_info = storage.pruning_dal().get_pruning_info().await?; | ||
| if let Some(last_hard_pruned_l2_block) = pruning_info.last_hard_pruned_l2_block { | ||
| anyhow::ensure!( | ||
| last_hard_pruned_l2_block == snapshot_l2_block, | ||
| "Additional data was pruned compared to tree recovery L2 block #{snapshot_l2_block}: {pruning_info:?}. \ | ||
| Continuing recovery is impossible; to recover the tree, drop its RocksDB directory, stop pruning and restart recovery" | ||
| ); | ||
| } | ||
| Ok(()) | ||
| } | ||
|
|
||
| /// Returns `Ok(true)` if the chunk was recovered, `Ok(false)` if the recovery process was interrupted. | ||
| async fn recover_key_chunk( | ||
| tree: &Mutex<AsyncTreeRecovery>, | ||
|
|
@@ -363,7 +417,9 @@ impl AsyncTreeRecovery { | |
| .storage_logs_dal() | ||
| .get_tree_entries_for_l2_block(snapshot_l2_block, key_chunk.clone()) | ||
| .await?; | ||
| Self::check_pruning_info(&mut storage, snapshot_l2_block).await?; | ||
| drop(storage); | ||
|
|
||
| let entries_latency = entries_latency.observe(); | ||
| tracing::debug!( | ||
| "Loaded {} entries for chunk {key_chunk:?} in {entries_latency:?}", | ||
|
|
@@ -414,13 +470,3 @@ impl AsyncTreeRecovery { | |
| Ok(true) | ||
| } | ||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
It would be beneficial to record the state root hash of the last hard-pruned L1 batch somewhere (in pruning logs?) so that it can be checked here. If this sounds OK, I'd prefer to implement it in a follow-up PR.