Add more metrics and logs to the witness generator

core/bin/zksync_witness_generator/src/witness_generator.rs

@@ -1,3 +1,4 @@
+use std::time::Instant;
 // Built-in
 use std::{thread, time};
 // External
@@ -74,6 +75,7 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
         &self,
         block_number: BlockNumber,
     ) -> Result<BlockInfo, anyhow::Error> {
+        let start = Instant::now();
         let mut storage = self.database.acquire_connection().await?;
         let mut transaction = storage.start_transaction().await?;
         let block = self
@@ -94,19 +96,25 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
             BlockInfo::NotReadyBlock
         };
         transaction.commit().await?;
+        metrics::histogram!("witness_generator", start.elapsed(), "stage" => "should_work_on_block");
         Ok(block_info)
     }
 
     async fn load_account_tree(
         &self,
         block: BlockNumber,
     ) -> Result<CircuitAccountTree, anyhow::Error> {
+        let fn_start = Instant::now();
+
         let mut storage = self.database.acquire_connection().await?;
+
+        let start = Instant::now();
         let mut circuit_account_tree = CircuitAccountTree::new(account_tree_depth());
+        let cache = self.database.load_account_tree_cache(&mut storage).await?;
+        metrics::histogram!("witness_generator", start.elapsed(), "stage" => "load_cache");
 
-        if let Some((cached_block, account_tree_cache)) =
-            self.database.load_account_tree_cache(&mut storage).await?
-        {
+        let start = Instant::now();
+        if let Some((cached_block, account_tree_cache)) = cache {
             let (_, accounts) = self
                 .database
                 .load_committed_state(&mut storage, Some(block))
@@ -116,6 +124,11 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
             }
             circuit_account_tree.set_internals(serde_json::from_value(account_tree_cache)?);
             if block != cached_block {
+                // There is no relevant cache, so we have to use some outdated cache and update the tree.
+                metrics::increment_counter!("witness_generator.cache_access", "type" => "miss");
+
+                vlog::info!("Reconstructing the cache for the block {} using the cached tree for the block {}", block, cached_block);
+
                 let (_, accounts) = self
                     .database
                     .load_committed_state(&mut storage, Some(block))
@@ -137,6 +150,9 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
                     }
                 }
                 circuit_account_tree.root_hash();
+                metrics::histogram!("witness_generator", start.elapsed(), "stage" => "recreate_tree_from_cache");
+
+                let start = Instant::now();
                 let account_tree_cache = circuit_account_tree.get_internals();
                 self.database
                     .store_account_tree_cache(
@@ -145,8 +161,13 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
                         serde_json::to_value(account_tree_cache)?,
                     )
                     .await?;
+                metrics::histogram!("witness_generator", start.elapsed(), "stage" => "store_cache");
+            } else {
+                // There exists a cache for the block we are interested in.
+                metrics::increment_counter!("witness_generator.cache_access", "type" => "hit");
             }
         } else {
+            // There are no caches at all.
             let (_, accounts) = self
                 .database
                 .load_committed_state(&mut storage, Some(block))
@@ -155,16 +176,21 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
                 circuit_account_tree.insert(*id, account.into());
             }
             circuit_account_tree.root_hash();
-            let account_tree_cache = circuit_account_tree.get_internals();
+
+            metrics::histogram!("witness_generator", start.elapsed(), "stage" => "recreate_tree_from_scratch");
+
+            let start = Instant::now();
+            let tree_cache = serde_json::to_value(circuit_account_tree.get_internals())?;
+            metrics::histogram!("witness_generator", start.elapsed(), "stage" => "serialize_cache");
+
+            let start = Instant::now();
             self.database
-                .store_account_tree_cache(
-                    &mut storage,
-                    block,
-                    serde_json::to_value(account_tree_cache)?,
-                )
+                .store_account_tree_cache(&mut storage, block, tree_cache)
                 .await?;
+            metrics::histogram!("witness_generator", start.elapsed(), "stage" => "store_cache");
         }
 
+        let start = Instant::now();
         if block != BlockNumber(0) {
             let storage_block = self
                 .database
@@ -177,34 +203,35 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
                 "account tree root hash restored incorrectly"
             );
         }
+        metrics::histogram!("witness_generator", start.elapsed(), "stage" => "ensure_root_hash");
 
+        metrics::histogram!("witness_generator", fn_start.elapsed(), "stage" => "load_account_tree");
         Ok(circuit_account_tree)
     }
 
     async fn prepare_witness_and_save_it(&self, block: Block) -> anyhow::Result<()> {
-        let timer = time::Instant::now();
+        let fn_start = Instant::now();
         let mut storage = self.database.acquire_connection().await?;
 
+        let start = Instant::now();
         let mut circuit_account_tree = self.load_account_tree(block.block_number - 1).await?;
-        vlog::trace!(
-            "Witness generator loading circuit account tree {}s",
-            timer.elapsed().as_secs()
-        );
+        metrics::histogram!("witness_generator", start.elapsed(), "stage" => "load_tree_full");
 
-        let timer = time::Instant::now();
+        let start = Instant::now();
         let witness: ProverData = build_block_witness(&mut circuit_account_tree, &block)?.into();
-        vlog::trace!(
-            "Witness generator witness build {}s",
-            timer.elapsed().as_secs()
-        );
+        metrics::histogram!("witness_generator", start.elapsed(), "stage" => "build_witness");
 
+        let start = Instant::now();
         self.database
             .store_witness(
                 &mut storage,
                 block.block_number,
                 serde_json::to_value(witness).expect("Witness serialize to json"),
             )
             .await?;
+        metrics::histogram!("witness_generator", start.elapsed(), "stage" => "store_witness");
+
+        metrics::histogram!("witness_generator", fn_start.elapsed(), "stage" => "prepare_witness_and_save_it");
         Ok(())
     }