From eac95bdf10f84652997f88459e9c0680248610cc Mon Sep 17 00:00:00 2001 From: Conrado Gouvea Date: Tue, 10 Aug 2021 10:33:34 -0300 Subject: [PATCH] Cache note commitment tree roots (#2584) * Cache note commitment tree roots * Add comments to cached root fields * Apply suggestions from code review Co-authored-by: teor Co-authored-by: teor --- zebra-chain/src/orchard/tree.rs | 33 +++++++++++++++++- zebra-chain/src/sapling/tree.rs | 34 +++++++++++++++++-- zebra-state/src/constants.rs | 2 +- .../service/finalized_state/tests/vectors.rs | 22 +++++++++--- 4 files changed, 83 insertions(+), 8 deletions(-) diff --git a/zebra-chain/src/orchard/tree.rs b/zebra-chain/src/orchard/tree.rs index 2b9644aa..ba045e5f 100644 --- a/zebra-chain/src/orchard/tree.rs +++ b/zebra-chain/src/orchard/tree.rs @@ -15,6 +15,7 @@ #![allow(dead_code)] use std::{ + cell::Cell, convert::TryFrom, fmt, hash::{Hash, Hasher}, @@ -221,6 +222,24 @@ pub struct NoteCommitmentTree { /// has non-empty nodes. Upper (near root) empty nodes of the branch are not /// stored. inner: bridgetree::Frontier, + /// A cached root of the tree. + /// + /// Every time the root is computed by [`Self::root`] it is cached here, + /// and the cached value will be returned by [`Self::root`] until the tree is + /// changed by [`Self::append`]. This greatly increases performance + /// because it avoids recomputing the root when the tree does not change + /// between blocks. In the finalized state, the tree is read from + /// disk for every block processed, which would also require recomputing + /// the root even if it has not changed (note that the cached root is + /// serialized with the tree). This is particularly important since we decided + /// to instantiate the trees from the genesis block, for simplicity. + /// + /// [`Cell`] offers interior mutability (it works even with a non-mutable + /// reference to the tree) but it prevents the tree (and anything that uses it) + /// from being shared between threads. If this ever becomes an issue we can + /// leave caching to the callers (which requires much more code), or replace + /// `Cell` with `Arc>` (and be careful of deadlocks and async code.) + cached_root: Cell>, } impl NoteCommitmentTree { @@ -233,6 +252,8 @@ impl NoteCommitmentTree { /// Returns an error if the tree is full. pub fn append(&mut self, cm_x: pallas::Base) -> Result<(), NoteCommitmentTreeError> { if self.inner.append(&cm_x.into()) { + // Invalidate cached root + self.cached_root.replace(None); Ok(()) } else { Err(NoteCommitmentTreeError::FullTree) @@ -242,7 +263,16 @@ impl NoteCommitmentTree { /// Returns the current root of the tree, used as an anchor in Orchard /// shielded transactions. pub fn root(&self) -> Root { - Root(self.inner.root().0) + match self.cached_root.get() { + // Return cached root + Some(root) => root, + None => { + // Compute root and cache it + let root = Root(self.inner.root().0); + self.cached_root.replace(Some(root)); + root + } + } } /// Get the Pallas-based Sinsemilla hash / root node of this merkle tree of @@ -274,6 +304,7 @@ impl Default for NoteCommitmentTree { fn default() -> Self { Self { inner: bridgetree::Frontier::new(), + cached_root: Default::default(), } } } diff --git a/zebra-chain/src/sapling/tree.rs b/zebra-chain/src/sapling/tree.rs index ba7b3d59..493e5e4d 100644 --- a/zebra-chain/src/sapling/tree.rs +++ b/zebra-chain/src/sapling/tree.rs @@ -13,7 +13,7 @@ #![allow(clippy::unit_arg)] #![allow(dead_code)] -use std::fmt; +use std::{cell::Cell, fmt}; use bitvec::prelude::*; use incrementalmerkletree::{bridgetree, Frontier}; @@ -185,6 +185,24 @@ pub struct NoteCommitmentTree { /// has non-empty nodes. Upper (near root) empty nodes of the branch are not /// stored. inner: bridgetree::Frontier, + /// A cached root of the tree. + /// + /// Every time the root is computed by [`Self::root`] it is cached here, + /// and the cached value will be returned by [`Self::root`] until the tree is + /// changed by [`Self::append`]. This greatly increases performance + /// because it avoids recomputing the root when the tree does not change + /// between blocks. In the finalized state, the tree is read from + /// disk for every block processed, which would also require recomputing + /// the root even if it has not changed (note that the cached root is + /// serialized with the tree). This is particularly important since we decided + /// to instantiate the trees from the genesis block, for simplicity. + /// + /// [`Cell`] offers interior mutability (it works even with a non-mutable + /// reference to the tree) but it prevents the tree (and anything that uses it) + /// from being shared between threads. If this ever becomes an issue we can + /// leave caching to the callers (which requires much more code), or replace + /// `Cell` with `Arc>` (and be careful of deadlocks and async code.) + cached_root: Cell>, } impl NoteCommitmentTree { @@ -197,6 +215,8 @@ impl NoteCommitmentTree { /// Returns an error if the tree is full. pub fn append(&mut self, cm_u: jubjub::Fq) -> Result<(), NoteCommitmentTreeError> { if self.inner.append(&cm_u.into()) { + // Invalidate cached root + self.cached_root.replace(None); Ok(()) } else { Err(NoteCommitmentTreeError::FullTree) @@ -206,7 +226,16 @@ impl NoteCommitmentTree { /// Returns the current root of the tree, used as an anchor in Sapling /// shielded transactions. pub fn root(&self) -> Root { - Root(self.inner.root().0) + match self.cached_root.get() { + // Return cached root + Some(root) => root, + None => { + // Compute root and cache it + let root = Root(self.inner.root().0); + self.cached_root.replace(Some(root)); + root + } + } } /// Get the Jubjub-based Pedersen hash of root node of this merkle tree of @@ -238,6 +267,7 @@ impl Default for NoteCommitmentTree { fn default() -> Self { Self { inner: bridgetree::Frontier::new(), + cached_root: Default::default(), } } } diff --git a/zebra-state/src/constants.rs b/zebra-state/src/constants.rs index 81375c5b..a26e5862 100644 --- a/zebra-state/src/constants.rs +++ b/zebra-state/src/constants.rs @@ -25,7 +25,7 @@ pub const MIN_TRANSPARENT_COINBASE_MATURITY: u32 = 100; pub const MAX_BLOCK_REORG_HEIGHT: u32 = MIN_TRANSPARENT_COINBASE_MATURITY - 1; /// The database format version, incremented each time the database format changes. -pub const DATABASE_FORMAT_VERSION: u32 = 8; +pub const DATABASE_FORMAT_VERSION: u32 = 9; /// The maximum number of blocks to check for NU5 transactions, /// before we assume we are on a pre-NU5 legacy chain. diff --git a/zebra-state/src/service/finalized_state/tests/vectors.rs b/zebra-state/src/service/finalized_state/tests/vectors.rs index 6501f091..7e298706 100644 --- a/zebra-state/src/service/finalized_state/tests/vectors.rs +++ b/zebra-state/src/service/finalized_state/tests/vectors.rs @@ -18,17 +18,24 @@ fn sapling_note_commitment_tree_serialization() { "7c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c", ]; - for cm_u_hex in hex_commitments { + for (idx, cm_u_hex) in hex_commitments.iter().enumerate() { let bytes = <[u8; 32]>::from_hex(cm_u_hex).unwrap(); let cm_u = jubjub::Fq::from_bytes(&bytes).unwrap(); incremental_tree.append(cm_u).unwrap(); + if idx % 2 == 0 { + // Cache the root half of the time to make sure it works in both cases + let _ = incremental_tree.root(); + } } + // Make sure the last root is cached + let _ = incremental_tree.root(); + // This test vector was generated by the code itself. // The purpose of this test is to make sure the serialization format does // not change by accident. - let expected_serialized_tree_hex = "0102007c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c0162324ff2c329e99193a74d28a585a3c167a93bf41a255135529c913bd9b1e666"; + let expected_serialized_tree_hex = "0102007c3ea01a6e3a3d90cf59cd789e467044b5cd78eb2c84cc6816f960746d0e036c0162324ff2c329e99193a74d28a585a3c167a93bf41a255135529c913bd9b1e66601ddaa1ab86de5c153993414f34ba97e9674c459dfadde112b89eeeafa0e5a204c"; let serialized_tree = incremental_tree.as_bytes(); assert_eq!(hex::encode(&serialized_tree), expected_serialized_tree_hex); @@ -62,15 +69,22 @@ fn orchard_note_commitment_tree_serialization() { ], ]; - for cm_x_bytes in &commitments { + for (idx, cm_x_bytes) in commitments.iter().enumerate() { let cm_x = pallas::Base::from_bytes(cm_x_bytes).unwrap(); incremental_tree.append(cm_x).unwrap(); + if idx % 2 == 0 { + // Cache the root half of the time to make sure it works in both cases + let _ = incremental_tree.root(); + } } + // Make sure the last root is cached + let _ = incremental_tree.root(); + // This test vector was generated by the code itself. // The purpose of this test is to make sure the serialization format does // not change by accident. - let expected_serialized_tree_hex = "010200ee9488053a30c596b43014105d3477e6f578c89240d1d1ee1743b77bb6adc40a01a34b69a4e4d9ccf954d46e5da1004d361a5497f511aeb4d481d23c0be1778133"; + let expected_serialized_tree_hex = "010200ee9488053a30c596b43014105d3477e6f578c89240d1d1ee1743b77bb6adc40a01a34b69a4e4d9ccf954d46e5da1004d361a5497f511aeb4d481d23c0be177813301a0be6dab19bc2c65d8299258c16e14d48ec4d4959568c6412aa85763c222a702"; let serialized_tree = incremental_tree.as_bytes(); assert_eq!(hex::encode(&serialized_tree), expected_serialized_tree_hex);