Move RewardMerkleTreeV2 storage to memory #3908
Conversation
Summary of ChangesHello @ss-es, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly refactors the storage and access mechanism for RewardMerkleTreeV2 data. The core change involves moving the storage of these Merkle trees from the SQL database to an in-memory, thread-safe registry. This change aims to improve performance by providing faster access to reward account proofs, as the data will be readily available in memory rather than requiring database lookups. The in-memory trees are populated during state updates for relevant epoch versions, and a configurable history of these trees can be maintained. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the storage mechanism for RewardMerkleTreeV2, moving it from persistent database storage to an in-memory registry. This change simplifies the data access logic for RewardMerkleTreeV2 proofs and allows for more efficient retrieval. However, there are a couple of areas that could be improved regarding default retention policy and potential performance implications for batch operations.
sequencer/src/api.rs
Outdated
| let n = std::env::var("REWARD_MERKLE_TREE_V2_RETAIN_HISTORY") | ||
| .ok() | ||
| .and_then(|n| n.parse::<u64>().ok()) | ||
| .unwrap_or(0); |
There was a problem hiding this comment.
The environment variable REWARD_MERKLE_TREE_V2_RETAIN_HISTORY defaults to 0 if not set or unparseable. This means that by default, no historical RewardMerkleTreeV2 instances will be retained in memory. This might lead to unexpected behavior for historical queries if not explicitly configured, as the name "retain history" suggests some level of retention. Consider setting a more sensible default (e.g., 1 or a small number) or adding a clear warning in the documentation about this default behavior.
let n = std::env::var("REWARD_MERKLE_TREE_V2_RETAIN_HISTORY")
.ok()
.and_then(|n| n.parse::<u64>().ok())
.unwrap_or(10); // Default to retaining 10 historical trees
sequencer/src/api/sql.rs
Outdated
| for account in accounts { | ||
| // Clone things we will need in the closure | ||
| let db_clone = db.clone(); | ||
| let account_clone = *account; | ||
| let header_height = header.height(); | ||
|
|
||
| // Create the closure that will get the path for the account | ||
| let func = async move { | ||
| // Open a new transaction | ||
| let mut tx = db_clone | ||
| .read() | ||
| .await | ||
| .with_context(|| "failed to open read transaction")?; | ||
|
|
||
| // Get the path for the account | ||
| let proof = tx | ||
| .get_path( | ||
| Snapshot::<SeqTypes, RewardMerkleTreeV2, { RewardMerkleTreeV2::ARITY }>::Index( | ||
| header_height, | ||
| ), | ||
| account_clone, | ||
| ) | ||
| .await | ||
| .with_context(|| { | ||
| format!( | ||
| "failed to get path for v2 reward account {account_clone:?}; height \ | ||
| {height}" | ||
| ) | ||
| })?; | ||
|
|
||
| Ok::<_, anyhow::Error>(proof) | ||
| }; | ||
|
|
||
| // Spawn the task | ||
| let id = join_set.spawn(func).id(); | ||
|
|
||
| // Add the task ID to the account map | ||
| task_id_to_account.insert(id, account); | ||
| } | ||
|
|
||
| // Wait for each task to complete | ||
| while let Some(result) = join_set.join_next_with_id().await { | ||
| // Get the inner result (past the join error) | ||
| let (id, result) = result.with_context(|| "failed to join task")?; | ||
|
|
||
| // Get the proof from the result | ||
| let proof = result?; | ||
|
|
||
| // Get the account from the task ID to account map | ||
| let account = task_id_to_account | ||
| .remove(&id) | ||
| .with_context(|| "task ID for spawned task not found")?; | ||
|
|
||
| match proof.proof.first().with_context(|| { | ||
| format!("empty proof for v2 reward account {account:?}; height {height}") | ||
| })? { | ||
| MerkleNode::Leaf { pos, elem, .. } => { | ||
| snapshot.remember(*pos, *elem, proof)?; | ||
| }, | ||
| MerkleNode::Empty => { | ||
| snapshot.non_membership_remember(*account, proof)?; | ||
| }, | ||
| _ => { | ||
| bail!("invalid proof for v2 reward account {account:?}; height {height}"); | ||
| match db | ||
| .load_v2_reward_account_proof(height, *account) | ||
| .await? | ||
| .proof | ||
| .proof | ||
| { | ||
| RewardMerkleProofV2::Presence(pf) | RewardMerkleProofV2::Absence(pf) => { | ||
| snapshot.insert_path(*account, &pf)?; | ||
| }, | ||
| } |
There was a problem hiding this comment.
The previous implementation of load_v2_reward_accounts used a BoundedJoinSet to fetch account paths concurrently. The new implementation iterates and awaits each load_v2_reward_account_proof call sequentially. If the accounts slice can be large, this sequential processing might introduce a performance bottleneck. Consider reintroducing concurrency (e.g., using futures::stream::iter().map().buffer_unordered()) to maintain performance for batch operations.
for account in accounts {
let (proof, _) = db
.load_v2_reward_account_proof(height, *account)
.await?
.proof;
match proof {
RewardMerkleProofV2::Presence(pf) | RewardMerkleProofV2::Absence(pf) => {
snapshot.insert_path(*account, &pf)?;
}
}
}
sequencer/src/api.rs
Outdated
| #[derive(Serialize, Deserialize)] | ||
| pub(crate) enum RewardMerkleTreeV2Data { | ||
| Tree(RewardMerkleTreeV2), | ||
| Balances(Vec<(RewardAccountV2, RewardAmount)>), | ||
| } |
There was a problem hiding this comment.
Could use some comments explaining why we have the two separate cases
| self.persist_tree(height, serialization).await?; | ||
| } | ||
|
|
||
| if height.is_multiple_of(30) { |
There was a problem hiding this comment.
Make 30 a configurable parameter?
ss-es
left a comment
ss-es
left a comment
There was a problem hiding this comment.
Update: proof generation is now decoupled from the state storage loop. specifically, no part of the logic which generates and persists proofs can cause update_state_storage to fail
because the reward merkle proofs are transient and only serve the reward claim contract, it's not consequential if this fails. we end up retrying it as part of the normal state update loop within a minute or so anyway
Uh oh!
There was an error while loading. Please reload this page.