Skip to content

Commit

Permalink
cleanup(state): Refactor subtree read functions to accept ranges and …
Browse files Browse the repository at this point in the history 
…check memory first (#7739)

* refactors sapling_subtrees to use ranges

* refactor orchard_subtrees() to use ranges

* Restores check for inconsistent subtree roots

* updates vectors tests

* Applies some suggestions from code review and adds `read::tree::subtrees` function

* updates correctness comment & db read method names

* adds test_subtrees

* tests in-memory reads

* test that subtree read fns work right with excluded start bound

* applies suggestions from code review

* Applies suggestions from code review.

* Updates docs applying suggestion from code review

* adds new arg to zs_range_iter calls
  • Loading branch information
@arya2
arya2 authored Oct 20, 2023
1 parent 9c993f8 commit 06a2983
Show file tree
Hide file tree
Showing 5 changed files with 265 additions and 305 deletions.
55 changes: 43 additions & 12 deletions zebra-state/src/service.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ use zebra_chain::{
block::{self, CountedHeader, HeightDiff},
diagnostic::{task::WaitForPanics, CodeTimer},
parameters::{Network, NetworkUpgrade},
subtree::NoteCommitmentSubtreeIndex,
};

use crate::{
Expand Down Expand Up @@ -1507,14 +1508,29 @@ impl Service<ReadRequest> for ReadStateService {

tokio::task::spawn_blocking(move || {
span.in_scope(move || {
let end_index = limit
.and_then(|limit| start_index.0.checked_add(limit.0))
.map(NoteCommitmentSubtreeIndex);

let sapling_subtrees = state.non_finalized_state_receiver.with_watch_data(
|non_finalized_state| {
read::sapling_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index,
limit,
)
if let Some(end_index) = end_index {
read::sapling_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index..end_index,
)
} else {
// If there is no end bound, just return all the trees.
// If the end bound would overflow, just returns all the trees, because that's what
// `zcashd` does. (It never calculates an end bound, so it just keeps iterating until
// the trees run out.)
read::sapling_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index..,
)
}
},
);

Expand All @@ -1532,14 +1548,29 @@ impl Service<ReadRequest> for ReadStateService {

tokio::task::spawn_blocking(move || {
span.in_scope(move || {
let end_index = limit
.and_then(|limit| start_index.0.checked_add(limit.0))
.map(NoteCommitmentSubtreeIndex);

let orchard_subtrees = state.non_finalized_state_receiver.with_watch_data(
|non_finalized_state| {
read::orchard_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index,
limit,
)
if let Some(end_index) = end_index {
read::orchard_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index..end_index,
)
} else {
// If there is no end bound, just return all the trees.
// If the end bound would overflow, just returns all the trees, because that's what
// `zcashd` does. (It never calculates an end bound, so it just keeps iterating until
// the trees run out.)
read::orchard_subtrees(
non_finalized_state.best_chain(),
&state.db,
start_index..,
)
}
},
);

Expand Down
116 changes: 14 additions & 102 deletions zebra-state/src/service/finalized_state/zebra_db/shielded.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -245,64 +245,20 @@ impl ZebraDb {
Some(subtree_data.with_index(index))
}

/// Returns a list of Sapling [`NoteCommitmentSubtree`]s starting at `start_index`.
/// If `limit` is provided, the list is limited to `limit` entries.
///
/// If there is no subtree at `start_index`, the returned list is empty.
/// Otherwise, subtrees are continuous up to the finalized tip.
///
/// # Correctness
///
/// This method is specifically designed for the `z_getsubtreesbyindex` state request.
/// It might not work for other RPCs or state checks.
pub fn sapling_subtree_list_by_index_for_rpc(
/// Returns a list of Sapling [`NoteCommitmentSubtree`]s in the provided range.
#[allow(clippy::unwrap_in_result)]
pub fn sapling_subtree_list_by_index_range(
&self,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>> {
let sapling_subtrees = self
.db
.cf_handle("sapling_note_commitment_subtree")
.unwrap();

// Calculate the end bound, checking for overflow.
let exclusive_end_bound: Option<NoteCommitmentSubtreeIndex> = limit
.and_then(|limit| start_index.0.checked_add(limit.0))
.map(NoteCommitmentSubtreeIndex);

let list: BTreeMap<
NoteCommitmentSubtreeIndex,
NoteCommitmentSubtreeData<sapling::tree::Node>,
>;

if let Some(exclusive_end_bound) = exclusive_end_bound {
list = self
.db
.zs_range_iter(&sapling_subtrees, start_index..exclusive_end_bound, false)
.collect();
} else {
// If there is no end bound, just return all the trees.
// If the end bound would overflow, just returns all the trees, because that's what
// `zcashd` does. (It never calculates an end bound, so it just keeps iterating until
// the trees run out.)
list = self
.db
.zs_range_iter(&sapling_subtrees, start_index.., false)
.collect();
}

// Make sure the amount of retrieved subtrees does not exceed the given limit.
#[cfg(debug_assertions)]
if let Some(limit) = limit {
assert!(list.len() <= limit.0.into());
}

// Check that we got the start subtree.
if list.get(&start_index).is_some() {
list
} else {
BTreeMap::new()
}
self.db
.zs_range_iter(&sapling_subtrees, range, false)
.collect()
}

/// Get the sapling note commitment subtress for the finalized tip.
Expand Down Expand Up @@ -415,64 +371,20 @@ impl ZebraDb {
Some(subtree_data.with_index(index))
}

/// Returns a list of Orchard [`NoteCommitmentSubtree`]s starting at `start_index`.
/// If `limit` is provided, the list is limited to `limit` entries.
///
/// If there is no subtree at `start_index`, the returned list is empty.
/// Otherwise, subtrees are continuous up to the finalized tip.
///
/// # Correctness
///
/// This method is specifically designed for the `z_getsubtreesbyindex` state request.
/// It might not work for other RPCs or state checks.
pub fn orchard_subtree_list_by_index_for_rpc(
/// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range.
#[allow(clippy::unwrap_in_result)]
pub fn orchard_subtree_list_by_index_range(
&self,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
let orchard_subtrees = self
.db
.cf_handle("orchard_note_commitment_subtree")
.unwrap();

// Calculate the end bound, checking for overflow.
let exclusive_end_bound: Option<NoteCommitmentSubtreeIndex> = limit
.and_then(|limit| start_index.0.checked_add(limit.0))
.map(NoteCommitmentSubtreeIndex);

let list: BTreeMap<
NoteCommitmentSubtreeIndex,
NoteCommitmentSubtreeData<orchard::tree::Node>,
>;

if let Some(exclusive_end_bound) = exclusive_end_bound {
list = self
.db
.zs_range_iter(&orchard_subtrees, start_index..exclusive_end_bound, false)
.collect();
} else {
// If there is no end bound, just return all the trees.
// If the end bound would overflow, just returns all the trees, because that's what
// `zcashd` does. (It never calculates an end bound, so it just keeps iterating until
// the trees run out.)
list = self
.db
.zs_range_iter(&orchard_subtrees, start_index.., false)
.collect();
}

// Make sure the amount of retrieved subtrees does not exceed the given limit.
#[cfg(debug_assertions)]
if let Some(limit) = limit {
assert!(list.len() <= limit.0.into());
}

// Check that we got the start subtree.
if list.get(&start_index).is_some() {
list
} else {
BTreeMap::new()
}
self.db
.zs_range_iter(&orchard_subtrees, range, false)
.collect()
}

/// Get the orchard note commitment subtress for the finalized tip.
Expand Down
28 changes: 6 additions & 22 deletions zebra-state/src/service/non_finalized_state/chain.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -698,8 +698,7 @@ impl Chain {
.map(|(index, subtree)| subtree.with_index(*index))
}

/// Returns a list of Sapling [`NoteCommitmentSubtree`]s at or after `start_index`.
/// If `limit` is provided, the list is limited to `limit` entries.
/// Returns a list of Sapling [`NoteCommitmentSubtree`]s in the provided range.
///
/// Unlike the finalized state and `ReadRequest::SaplingSubtrees`, the returned subtrees
/// can start after `start_index`. These subtrees are continuous up to the tip.
Expand All @@ -709,17 +708,10 @@ impl Chain {
/// finalized updates.
pub fn sapling_subtrees_in_range(
&self,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<sapling::tree::Node>> {
let limit = limit
.map(|limit| usize::from(limit.0))
.unwrap_or(usize::MAX);

// Since we're working in memory, it's ok to iterate through the whole range here.
self.sapling_subtrees
.range(start_index..)
.take(limit)
.range(range)
.map(|(index, subtree)| (*index, *subtree))
.collect()
}
Expand Down Expand Up @@ -910,8 +902,7 @@ impl Chain {
.map(|(index, subtree)| subtree.with_index(*index))
}

/// Returns a list of Orchard [`NoteCommitmentSubtree`]s at or after `start_index`.
/// If `limit` is provided, the list is limited to `limit` entries.
/// Returns a list of Orchard [`NoteCommitmentSubtree`]s in the provided range.
///
/// Unlike the finalized state and `ReadRequest::OrchardSubtrees`, the returned subtrees
/// can start after `start_index`. These subtrees are continuous up to the tip.
Expand All @@ -921,17 +912,10 @@ impl Chain {
/// finalized updates.
pub fn orchard_subtrees_in_range(
&self,
start_index: NoteCommitmentSubtreeIndex,
limit: Option<NoteCommitmentSubtreeIndex>,
range: impl std::ops::RangeBounds<NoteCommitmentSubtreeIndex>,
) -> BTreeMap<NoteCommitmentSubtreeIndex, NoteCommitmentSubtreeData<orchard::tree::Node>> {
let limit = limit
.map(|limit| usize::from(limit.0))
.unwrap_or(usize::MAX);

// Since we're working in memory, it's ok to iterate through the whole range here.
self.orchard_subtrees
.range(start_index..)
.take(limit)
.range(range)
.map(|(index, subtree)| (*index, *subtree))
.collect()
}
Expand Down
Loading

0 comments on commit 06a2983

Please sign in to comment.