Migrate to latest shardtree to improve performance

nuttycom · Jul 18, 2023 · df04203 · df04203
1 parent 843930a
commit df04203
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 15 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,6 +19,6 @@ panic = 'abort'
 codegen-units = 1
 
 [patch.crates-io]
-incrementalmerkletree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "67111e29403c33f2e36d6924167f1d5f04ad0fc2" }
-shardtree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "67111e29403c33f2e36d6924167f1d5f04ad0fc2" }
+incrementalmerkletree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "bae25ad89c0c192bee625252d2d419bd56638e48" }
+shardtree = { git = "https://github.com/zcash/incrementalmerkletree.git", rev = "bae25ad89c0c192bee625252d2d419bd56638e48" }
 orchard = { git = "https://github.com/zcash/orchard.git", rev = "6ef89d5f154de2cf7b7dd87edb8d8c49158beebb" }
diff --git a/zcash_client_sqlite/src/serialization.rs b/zcash_client_sqlite/src/serialization.rs
@@ -4,7 +4,7 @@ use byteorder::{ReadBytesExt, WriteBytesExt};
 use core::ops::Deref;
 use shardtree::{Node, PrunableTree, RetentionFlags, Tree};
 use std::io::{self, Read, Write};
-use std::rc::Rc;
+use std::sync::Arc;
 use zcash_encoding::Optional;
 use zcash_primitives::merkle_tree::HashSer;
 
@@ -53,7 +53,7 @@ pub fn write_shard<H: HashSer, W: Write>(writer: &mut W, tree: &PrunableTree<H>)
 fn read_shard_v1<H: HashSer, R: Read>(mut reader: &mut R) -> io::Result<PrunableTree<H>> {
     match reader.read_u8()? {
         PARENT_TAG => {
-            let ann = Optional::read(&mut reader, <H as HashSer>::read)?.map(Rc::new);
+            let ann = Optional::read(&mut reader, <H as HashSer>::read)?.map(Arc::new);
             let left = read_shard_v1(reader)?;
             let right = read_shard_v1(reader)?;
             Ok(Tree::parent(ann, left, right))

diff --git a/zcash_client_sqlite/src/wallet/commitment_tree.rs b/zcash_client_sqlite/src/wallet/commitment_tree.rs
@@ -4,7 +4,7 @@ use std::{
     collections::BTreeSet,
     io::{self, Cursor},
     marker::PhantomData,
-    rc::Rc,
+    sync::Arc,
 };
 use zcash_client_backend::data_api::chain::CommitmentTreeRoot;
 
@@ -281,7 +281,7 @@ pub(crate) fn get_shard<H: HashSer>(
         let located_tree = LocatedPrunableTree::from_parts(shard_root_addr, shard_tree);
         if let Some(root_hash_data) = root_hash {
             let root_hash = H::read(Cursor::new(root_hash_data)).map_err(Either::Left)?;
-            Ok(located_tree.reannotate_root(Some(Rc::new(root_hash))))
+            Ok(located_tree.reannotate_root(Some(Arc::new(root_hash))))
         } else {
             Ok(located_tree)
         }

diff --git a/zcash_client_sqlite/src/wallet/init/migrations/shardtree_support.rs b/zcash_client_sqlite/src/wallet/init/migrations/shardtree_support.rs
@@ -8,7 +8,7 @@ use incrementalmerkletree::Retention;
 use rusqlite::{self, named_params, params};
 use schemer;
 use schemer_rusqlite::RusqliteMigration;
-use shardtree::ShardTree;
+use shardtree::{caching::CachingShardStore, ShardTree, ShardTreeError};
 use tracing::{debug, trace};
 use uuid::Uuid;
 
@@ -108,6 +108,7 @@ impl RusqliteMigration for Migration {
                 transaction,
                 SAPLING_TABLES_PREFIX,
             )?;
+        let shard_store = CachingShardStore::load(shard_store).map_err(ShardTreeError::Storage)?;
         let mut shard_tree: ShardTree<
             _,
             { sapling::NOTE_COMMITMENT_TREE_DEPTH },
@@ -154,13 +155,19 @@ impl RusqliteMigration for Migration {
                         frontier = ?nonempty_frontier,
                         "Inserting frontier nodes",
                     );
-                    shard_tree.insert_frontier_nodes(
-                        nonempty_frontier.clone(),
-                        Retention::Checkpoint {
-                            id: BlockHeight::from(block_height),
-                            is_marked: false,
-                        },
-                    )?;
+                    shard_tree
+                        .insert_frontier_nodes(
+                            nonempty_frontier.clone(),
+                            Retention::Checkpoint {
+                                id: BlockHeight::from(block_height),
+                                is_marked: false,
+                            },
+                        )
+                        .map_err(|e| match e {
+                            ShardTreeError::Query(e) => ShardTreeError::Query(e),
+                            ShardTreeError::Insert(e) => ShardTreeError::Insert(e),
+                            ShardTreeError::Storage(_) => unreachable!(),
+                        })?
                 }
             }
         }
@@ -203,10 +210,21 @@ impl RusqliteMigration for Migration {
                     updated_note_positions.insert(witnessed_position);
                 }
 
-                shard_tree.insert_witness_nodes(witness, BlockHeight::from(block_height))?;
+                shard_tree
+                    .insert_witness_nodes(witness, BlockHeight::from(block_height))
+                    .map_err(|e| match e {
+                        ShardTreeError::Query(e) => ShardTreeError::Query(e),
+                        ShardTreeError::Insert(e) => ShardTreeError::Insert(e),
+                        ShardTreeError::Storage(_) => unreachable!(),
+                    })?;
             }
         }
 
+        shard_tree
+            .into_store()
+            .flush()
+            .map_err(ShardTreeError::Storage)?;
+
         // Establish the scan queue & wallet history table.
         // block_range_end is exclusive.
         debug!("Creating table for scan queue");