Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

cleanup(state): Refactor subtree read functions to accept ranges and check memory first #7739

Merged
merged 15 commits into from
Oct 20, 2023
Merged

cleanup(state): Refactor subtree read functions to accept ranges and check memory first #7739

Show file tree
Hide file tree
Changes from 12 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
837061f
refactors sapling_subtrees to use ranges
arya2 Oct 12, 2023
a7edaa9
refactor orchard_subtrees() to use ranges
arya2 Oct 12, 2023
f1aaf2b
Restores check for inconsistent subtree roots
arya2 Oct 13, 2023
5b55813
updates vectors tests
arya2 Oct 13, 2023
ffc4f75
Applies some suggestions from code review and adds `read::tree::subtr…
arya2 Oct 18, 2023
0d7b13b
updates correctness comment & db read method names
arya2 Oct 18, 2023
8fe86fe
adds test_subtrees
arya2 Oct 19, 2023
4ae71e5
tests in-memory reads
arya2 Oct 19, 2023
7ec6f21
test that subtree read fns work right with excluded start bound
arya2 Oct 19, 2023
4cd27ab
Merge branch 'main' into cleanup-read-subtree
arya2 Oct 19, 2023
9737d22
applies suggestions from code review
arya2 Oct 19, 2023
668462d
Applies suggestions from code review.
arya2 Oct 20, 2023
3416c11
Updates docs applying suggestion from code review
arya2 Oct 20, 2023
7db03b1
Merge branch 'main' into cleanup-read-subtree
arya2 Oct 20, 2023
d396369
adds new arg to zs_range_iter calls
arya2 Oct 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 43 additions & 12 deletions zebra-state/src/service.rs
View file
Open in desktop
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))
arya2 marked this conversation as resolved.
Show resolved Hide resolved
.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
112 changes: 10 additions & 102 deletions zebra-state/src/service/finalized_state/zebra_db/shielded.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -245,64 +245,18 @@ 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)
.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..)
.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).collect()
}

/// Get the sapling note commitment subtress for the finalized tip.
Expand Down Expand Up @@ -415,64 +369,18 @@ 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)
.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..)
.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).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
Open in desktop
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
Loading