Replace serde_json::Value for tree cache with string and add cache si…

…ze metric
huynhvantoanbsbc · Mar 23, 2022 · 256c065 · 256c065
1 parent 9da0da7
commit 256c065
Show file tree

Hide file tree

Showing 11 changed files with 27 additions and 37 deletions.
diff --git a/core/bin/data_restore/src/data_restore_driver.rs b/core/bin/data_restore/src/data_restore_driver.rs
@@ -256,7 +256,7 @@ impl<T: Transport> DataRestoreDriver<T> {
         interactor
             .store_tree_cache(
                 self.tree_state.block_number,
-                serde_json::to_value(tree_cache).expect("failed to serialize tree cache"),
+                serde_json::to_string(&tree_cache).expect("failed to serialize tree cache"),
             )
             .await;
     }

diff --git a/core/bin/data_restore/src/database_storage_interactor.rs b/core/bin/data_restore/src/database_storage_interactor.rs
@@ -381,11 +381,7 @@ impl<'a> DatabaseStorageInteractor<'a> {
         }
     }
 
-    pub async fn store_tree_cache(
-        &mut self,
-        block_number: BlockNumber,
-        tree_cache: serde_json::Value,
-    ) {
+    pub async fn store_tree_cache(&mut self, block_number: BlockNumber, tree_cache: String) {
         self.storage
             .chain()
             .block_schema()

diff --git a/core/bin/data_restore/src/inmemory_storage_interactor.rs b/core/bin/data_restore/src/inmemory_storage_interactor.rs
@@ -316,11 +316,7 @@ impl InMemoryStorageInteractor {
         None
     }
 
-    pub async fn store_tree_cache(
-        &mut self,
-        _block_number: BlockNumber,
-        _tree_cache: serde_json::Value,
-    ) {
+    pub async fn store_tree_cache(&mut self, _block_number: BlockNumber, _tree_cache: String) {
         // Inmemory storage doesn't support caching.
     }
 

diff --git a/core/bin/data_restore/src/storage_interactor.rs b/core/bin/data_restore/src/storage_interactor.rs
@@ -203,11 +203,7 @@ impl StorageInteractor<'_> {
     /// * `block_number` - The corresponding block number
     /// * `tree_cache` - Merkle tree cache
     ///
-    pub async fn store_tree_cache(
-        &mut self,
-        block_number: BlockNumber,
-        tree_cache: serde_json::Value,
-    ) {
+    pub async fn store_tree_cache(&mut self, block_number: BlockNumber, tree_cache: String) {
         storage_interact!(self.store_tree_cache(block_number, tree_cache))
     }
 

diff --git a/core/bin/zksync_core/src/state_keeper/state_restore/db/storage.rs b/core/bin/zksync_core/src/state_keeper/state_restore/db/storage.rs
@@ -102,8 +102,8 @@ impl<'a, 'b> StateRestoreDb for StateRestoreStorage<'a, 'b> {
         block: BlockNumber,
         account_tree_cache: SparseMerkleTreeSerializableCacheBN256,
     ) {
-        let encoded_tree_cache =
-            serde_json::to_value(account_tree_cache).expect("Unable to encode account tree cache");
+        let encoded_tree_cache = serde_json::to_string(&account_tree_cache)
+            .expect("Unable to encode account tree cache");
         self.storage
             .chain()
             .block_schema()

diff --git a/core/bin/zksync_witness_generator/src/database.rs b/core/bin/zksync_witness_generator/src/database.rs
@@ -254,7 +254,7 @@ impl DatabaseInterface for Database {
         &self,
         connection: &mut StorageProcessor<'_>,
         block: BlockNumber,
-        tree_cache: serde_json::Value,
+        tree_cache: String,
     ) -> anyhow::Result<()> {
         connection
             .chain()

diff --git a/core/bin/zksync_witness_generator/src/database_interface.rs b/core/bin/zksync_witness_generator/src/database_interface.rs
@@ -132,7 +132,7 @@ pub trait DatabaseInterface: Send + Sync + Clone + 'static {
         &self,
         connection: &mut StorageProcessor<'_>,
         block: BlockNumber,
-        tree_cache: serde_json::Value,
+        tree_cache: String,
     ) -> anyhow::Result<()>;
 
     async fn store_witness(

diff --git a/core/bin/zksync_witness_generator/src/tests/mock.rs b/core/bin/zksync_witness_generator/src/tests/mock.rs
@@ -343,13 +343,11 @@ impl DatabaseInterface for MockDatabase {
         &self,
         _: &mut StorageProcessor<'_>,
         block: BlockNumber,
-        tree_cache: serde_json::Value,
+        tree_cache: String,
     ) -> anyhow::Result<()> {
         if *block == 0 {
             return Ok(());
         }
-        let tree_cache =
-            serde_json::to_string(&tree_cache).expect("Failed to serialize Account Tree Cache");
 
         let mut account_tree_cache = self.account_tree_cache.write().await;
         *account_tree_cache = AccountTreeCache {

diff --git a/core/bin/zksync_witness_generator/src/witness_generator.rs b/core/bin/zksync_witness_generator/src/witness_generator.rs
@@ -153,13 +153,11 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
                 metrics::histogram!("witness_generator", start.elapsed(), "stage" => "recreate_tree_from_cache");
 
                 let start = Instant::now();
-                let account_tree_cache = circuit_account_tree.get_internals();
+                let tree_cache = serde_json::to_string(&circuit_account_tree.get_internals())?;
+                metrics::histogram!("tree_cache_size", tree_cache.len() as f64);
+
                 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");
             } else {
@@ -180,7 +178,8 @@ impl<DB: DatabaseInterface> WitnessGenerator<DB> {
             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())?;
+            let tree_cache = serde_json::to_string(&circuit_account_tree.get_internals())?;
+            metrics::histogram!("tree_cache_size", tree_cache.len() as f64);
             metrics::histogram!("witness_generator", start.elapsed(), "stage" => "serialize_cache");
 
             let start = Instant::now();

diff --git a/core/lib/storage/src/chain/block/mod.rs b/core/lib/storage/src/chain/block/mod.rs
@@ -1161,19 +1161,18 @@ impl<'a, 'c> BlockSchema<'a, 'c> {
         Ok(())
     }
 
-    /// Stores account tree cache for a block
+    /// Stores account tree cache for a block.
+    /// Expects `tree_cache` to be a valid encoded JSON.
     pub async fn store_account_tree_cache(
         &mut self,
         block: BlockNumber,
-        tree_cache: serde_json::Value,
+        tree_cache: String,
     ) -> QueryResult<()> {
         let start = Instant::now();
         if *block == 0 {
             return Ok(());
         }
 
-        let tree_cache_str =
-            serde_json::to_string(&tree_cache).expect("Failed to serialize Account Tree Cache");
         sqlx::query!(
             "
             INSERT INTO account_tree_cache (block, tree_cache)
@@ -1182,7 +1181,7 @@ impl<'a, 'c> BlockSchema<'a, 'c> {
             DO NOTHING
             ",
             *block as i64,
-            tree_cache_str,
+            tree_cache,
         )
         .execute(self.0.conn())
         .await?;

diff --git a/core/lib/storage/src/tests/chain/block.rs b/core/lib/storage/src/tests/chain/block.rs
@@ -1284,7 +1284,10 @@ async fn test_remove_new_account_tree_cache(mut storage: StorageProcessor<'_>) -
             ))
             .await?;
         BlockSchema(&mut storage)
-            .store_account_tree_cache(BlockNumber(block_number), serde_json::Value::default())
+            .store_account_tree_cache(
+                BlockNumber(block_number),
+                serde_json::Value::default().to_string(),
+            )
             .await?;
     }
 
@@ -1370,7 +1373,10 @@ async fn test_remove_old_account_tree_cache(mut storage: StorageProcessor<'_>) -
             ))
             .await?;
         BlockSchema(&mut storage)
-            .store_account_tree_cache(BlockNumber(block_number), serde_json::Value::default())
+            .store_account_tree_cache(
+                BlockNumber(block_number),
+                serde_json::Value::default().to_string(),
+            )
             .await?;
     }