diff --git a/.github/workflows/forester-tests.yml b/.github/workflows/forester-tests.yml index d078ea52e8..ee6ebe5032 100644 --- a/.github/workflows/forester-tests.yml +++ b/.github/workflows/forester-tests.yml @@ -30,6 +30,11 @@ concurrency: env: RUST_BACKTRACE: "1" RUSTFLAGS: "--cfg tokio_unstable -D warnings" + TEST_MODE: "local" + TEST_V1_STATE: "true" + TEST_V2_STATE: "true" + TEST_V1_ADDRESS: "true" + TEST_V2_ADDRESS: "true" jobs: test: @@ -38,41 +43,17 @@ jobs: test-name: [ { - name: "e2e", - command: "test_state_indexer_async_batched", + name: "e2e (legacy)", + command: "test_e2e_v1", timeout: 60, needs-test-program: false, }, { - name: "address-batched", - command: "test_address_batched", + name: "e2e", + command: "test_e2e_v2", timeout: 60, needs-test-program: true, }, - { - name: "state-batched", - command: "test_state_batched", - timeout: 60, - needs-test-program: false, - }, - { - name: "state-photon-batched", - command: "test_state_indexer_batched", - timeout: 60, - needs-test-program: false, - }, - { - name: "2-foresters", - command: "test_epoch_monitor_with_2_foresters", - timeout: 60, - needs-test-program: false, - }, - { - name: "double-registration", - command: "test_epoch_double_registration", - timeout: 60, - needs-test-program: false, - }, ] name: test-${{ matrix.test-name.name }} runs-on: warp-ubuntu-latest-x64-4x @@ -117,6 +98,7 @@ jobs: run: | source ./scripts/devenv.sh cargo test-sbf -p create-address-test-program + - name: Run ${{ matrix.test-name.name }} tests run: | source ./scripts/devenv.sh diff --git a/Cargo.lock b/Cargo.lock index 03057b4669..0f822051d7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -783,6 +783,28 @@ dependencies = [ "pin-project-lite", ] +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.103", +] + [[package]] name = "async-trait" version = "0.1.88" @@ -2104,6 +2126,7 @@ dependencies = [ "account-compression", "anchor-lang", "anyhow", + "async-stream", "async-trait", "bb8", "borsh 0.10.4", @@ -2159,6 +2182,7 @@ version = "2.0.0" dependencies = [ "account-compression", "anchor-lang", + "async-stream", "async-trait", "bb8", "futures", diff --git a/cli/src/utils/constants.ts b/cli/src/utils/constants.ts index f883829def..81aaa478d3 100644 --- a/cli/src/utils/constants.ts +++ b/cli/src/utils/constants.ts @@ -24,7 +24,7 @@ export const PHOTON_VERSION = "0.51.0"; // Set these to override Photon requirements with a specific git commit: export const USE_PHOTON_FROM_GIT = true; // If true, will show git install command instead of crates.io. export const PHOTON_GIT_REPO = "https://github.com/lightprotocol/photon.git"; -export const PHOTON_GIT_COMMIT = "6ee3c027226ab9c90dc9d16691cdf76dd2f29dbf"; // If empty, will use main branch. +export const PHOTON_GIT_COMMIT = "c938ee83ad1b34abc389943334627a899da72953"; // If empty, will use main branch. export const LIGHT_PROTOCOL_PROGRAMS_DIR_ENV = "LIGHT_PROTOCOL_PROGRAMS_DIR"; export const BASE_PATH = "../../bin/"; diff --git a/forester-utils/Cargo.toml b/forester-utils/Cargo.toml index 6faa56da05..cc7e1af82c 100644 --- a/forester-utils/Cargo.toml +++ b/forester-utils/Cargo.toml @@ -34,6 +34,7 @@ account-compression = { workspace = true, features = ["cpi"] } tokio = { workspace = true } futures = { workspace = true } +async-stream = "0.3" anchor-lang = { workspace = true } diff --git a/forester-utils/src/error.rs b/forester-utils/src/error.rs index caec758b06..6ec96efdf8 100644 --- a/forester-utils/src/error.rs +++ b/forester-utils/src/error.rs @@ -1,5 +1,9 @@ +use light_batched_merkle_tree::errors::BatchedMerkleTreeError; +use light_hasher::HasherError; use thiserror::Error; +use crate::rpc_pool::PoolError; + #[derive(Error, Debug)] pub enum ForesterUtilsError { #[error("parse error: {0:?}")] @@ -12,4 +16,18 @@ pub enum ForesterUtilsError { Indexer(String), #[error("invalid slot number")] InvalidSlotNumber, + #[error("Hasher error: {0}")] + Hasher(#[from] HasherError), + + #[error("Account zero-copy error: {0}")] + AccountZeroCopy(String), + + #[error("light client error: {0}")] + LightClient(#[from] light_client::rpc::RpcError), + + #[error("batched merkle tree error: {0}")] + BatchedMerkleTree(#[from] BatchedMerkleTreeError), + + #[error("pool error: {0}")] + Pool(#[from] PoolError), } diff --git a/forester-utils/src/instructions/address_batch_update.rs b/forester-utils/src/instructions/address_batch_update.rs index 1a148a1997..b95901a2aa 100644 --- a/forester-utils/src/instructions/address_batch_update.rs +++ b/forester-utils/src/instructions/address_batch_update.rs @@ -1,13 +1,13 @@ -use std::time::Duration; +use std::{pin::Pin, sync::Arc, time::Duration}; use account_compression::processor::initialize_address_merkle_tree::Pubkey; -use futures::future; +use async_stream::stream; +use futures::{ + stream::{FuturesOrdered, Stream}, + StreamExt, +}; use light_batched_merkle_tree::{ - constants::DEFAULT_BATCH_ADDRESS_TREE_HEIGHT, - merkle_tree::{ - BatchedMerkleTreeAccount, InstructionDataAddressAppendInputs, - InstructionDataBatchNullifyInputs, - }, + constants::DEFAULT_BATCH_ADDRESS_TREE_HEIGHT, merkle_tree::InstructionDataAddressAppendInputs, }; use light_client::{indexer::Indexer, rpc::Rpc}; use light_compressed_account::{ @@ -18,101 +18,207 @@ use light_prover_client::{ proof_client::ProofClient, proof_types::batch_address_append::get_batch_address_append_circuit_inputs, }; -use light_sparse_merkle_tree::{ - changelog::ChangelogEntry, indexed_changelog::IndexedChangelogEntry, SparseMerkleTree, -}; +use light_sparse_merkle_tree::SparseMerkleTree; +use tokio::sync::Mutex; use tracing::{debug, error, info, warn}; -use crate::{error::ForesterUtilsError, utils::wait_for_indexer}; +use crate::{error::ForesterUtilsError, rpc_pool::SolanaRpcPool, utils::wait_for_indexer}; + +const MAX_PHOTON_ELEMENTS_PER_CALL: usize = 500; + +pub struct AddressUpdateConfig +where + R: Rpc + Send + Sync, + I: Indexer + Send, +{ + pub rpc_pool: Arc>, + pub indexer: Arc>, + pub merkle_tree_pubkey: Pubkey, + pub prover_url: String, + pub polling_interval: Duration, + pub max_wait_time: Duration, + pub ixs_per_tx: usize, +} -pub async fn create_batch_update_address_tree_instruction_data( - rpc: &mut R, - indexer: &mut I, - merkle_tree_pubkey: &Pubkey, +#[allow(clippy::too_many_arguments)] +fn stream_instruction_data<'a, I>( + indexer: Arc>, + merkle_tree_pubkey: Pubkey, prover_url: String, polling_interval: Duration, max_wait_time: Duration, -) -> Result<(Vec, u16), ForesterUtilsError> + leaves_hash_chains: Vec<[u8; 32]>, + start_index: u64, + zkp_batch_size: u16, + mut current_root: [u8; 32], + yield_batch_size: usize, +) -> impl Stream, ForesterUtilsError>> + Send + 'a where - R: Rpc, - I: Indexer, + I: Indexer + Send + 'a, { - info!("Creating batch update address tree instruction data"); - - let mut merkle_tree_account = rpc - .get_account(*merkle_tree_pubkey) - .await - .map_err(|e| { - error!("Failed to get account data from rpc: {:?}", e); - ForesterUtilsError::Rpc("Failed to get account data".into()) - })? - .unwrap(); - - let (leaves_hash_chains, start_index, current_root, batch_size) = { - let merkle_tree = BatchedMerkleTreeAccount::address_from_bytes( - merkle_tree_account.data.as_mut_slice(), - &(*merkle_tree_pubkey).into(), - ) - .unwrap(); - - let full_batch_index = merkle_tree.queue_batches.pending_batch_index; - let batch = &merkle_tree.queue_batches.batches[full_batch_index as usize]; + stream! { + let proof_client = Arc::new(ProofClient::with_config(prover_url, polling_interval, max_wait_time)); + let max_zkp_batches_per_call = calculate_max_zkp_batches_per_call(zkp_batch_size); + let total_chunks = leaves_hash_chains.len().div_ceil(max_zkp_batches_per_call); + + for chunk_idx in 0..total_chunks { + let mut indexer_guard = indexer.lock().await; + let chunk_start = chunk_idx * max_zkp_batches_per_call; + let chunk_end = std::cmp::min(chunk_start + max_zkp_batches_per_call, leaves_hash_chains.len()); + let chunk_hash_chains = &leaves_hash_chains[chunk_start..chunk_end]; + + let elements_for_chunk = chunk_hash_chains.len() * zkp_batch_size as usize; + let processed_items_offset = chunk_start * zkp_batch_size as usize; + + let indexer_update_info = match indexer_guard + .get_address_queue_with_proofs( + &merkle_tree_pubkey, + elements_for_chunk as u16, + Some(processed_items_offset as u64), + None, + ) + .await { + Ok(info) => info, + Err(e) => { + yield Err(ForesterUtilsError::Indexer(format!("Failed to get address queue with proofs: {}", e))); + return; + } + }; + + if chunk_idx == 0 { + if let Some(first_proof) = indexer_update_info.value.non_inclusion_proofs.first() { + if first_proof.root != current_root { + warn!("Indexer root does not match on-chain root"); + yield Err(ForesterUtilsError::Indexer("Indexer root does not match on-chain root".into())); + return; + } + } else { + yield Err(ForesterUtilsError::Indexer("No non-inclusion proofs found in indexer response".into())); + return; + } + } - let mut hash_chains = Vec::new(); - let zkp_batch_index = batch.get_num_inserted_zkps(); - let current_zkp_batch_index = batch.get_current_zkp_batch_index(); + let (all_inputs, new_current_root) = match get_all_circuit_inputs_for_chunk( + chunk_hash_chains, + &indexer_update_info, + zkp_batch_size, + chunk_start, + start_index, + current_root, + ) { + Ok((inputs, new_root)) => (inputs, new_root), + Err(e) => { + yield Err(e); + return; + } + }; + current_root = new_current_root; + + info!("Generating {} ZK proofs with hybrid approach for chunk {}", all_inputs.len(), chunk_idx + 1); + + let mut futures_ordered = FuturesOrdered::new(); + let mut proof_buffer = Vec::new(); + let mut pending_count = 0; + + for (i, inputs) in all_inputs.into_iter().enumerate() { + let client = Arc::clone(&proof_client); + futures_ordered.push_back(async move { + let result = client.generate_batch_address_append_proof(inputs).await; + (i, result) + }); + pending_count += 1; + + if pending_count >= yield_batch_size { + for _ in 0..yield_batch_size.min(pending_count) { + if let Some((idx, result)) = futures_ordered.next().await { + match result { + Ok((compressed_proof, new_root)) => { + let instruction_data = InstructionDataAddressAppendInputs { + new_root, + compressed_proof: CompressedProof { + a: compressed_proof.a, + b: compressed_proof.b, + c: compressed_proof.c, + }, + }; + proof_buffer.push(instruction_data); + }, + Err(e) => { + error!("Address proof failed to generate at index {}: {:?}", idx, e); + yield Err(ForesterUtilsError::Prover(format!( + "Address proof generation failed at batch {} in chunk {}: {}", + idx, chunk_idx, e + ))); + return; + } + } + pending_count -= 1; + } + } + + if !proof_buffer.is_empty() { + yield Ok(proof_buffer.clone()); + proof_buffer.clear(); + } + } + } - debug!( - "Full batch index: {}, inserted ZKPs: {}, current ZKP index: {}, ready for insertion: {}", - full_batch_index, zkp_batch_index, current_zkp_batch_index, current_zkp_batch_index - zkp_batch_index - ); + while let Some((idx, result)) = futures_ordered.next().await { + match result { + Ok((compressed_proof, new_root)) => { + let instruction_data = InstructionDataAddressAppendInputs { + new_root, + compressed_proof: CompressedProof { + a: compressed_proof.a, + b: compressed_proof.b, + c: compressed_proof.c, + }, + }; + proof_buffer.push(instruction_data); + + if proof_buffer.len() >= yield_batch_size { + yield Ok(proof_buffer.clone()); + proof_buffer.clear(); + } + }, + Err(e) => { + error!("Address proof failed to generate at index {}: {:?}", idx, e); + yield Err(ForesterUtilsError::Prover(format!( + "Address proof generation failed at batch {} in chunk {}: {}", + idx, chunk_idx, e + ))); + return; + } + } + } - for i in zkp_batch_index..current_zkp_batch_index { - hash_chains.push(merkle_tree.hash_chain_stores[full_batch_index as usize][i as usize]); + if !proof_buffer.is_empty() { + yield Ok(proof_buffer); + } } - - let start_index = merkle_tree.next_index; - let current_root = *merkle_tree.root_history.last().unwrap(); - let zkp_batch_size = batch.zkp_batch_size as u16; - - (hash_chains, start_index, current_root, zkp_batch_size) - }; - - if leaves_hash_chains.is_empty() { - debug!("No hash chains to process"); - return Ok((Vec::new(), batch_size)); } +} - wait_for_indexer(rpc, indexer).await?; - - let total_elements = batch_size as usize * leaves_hash_chains.len(); - debug!("Requesting {} total elements from indexer", total_elements); - - let indexer_update_info = indexer - .get_address_queue_with_proofs(merkle_tree_pubkey, total_elements as u16, None, None) - .await - .map_err(|e| { - error!("Failed to get batch address update info: {:?}", e); - ForesterUtilsError::Indexer("Failed to get batch address update info".into()) - })?; - debug!("indexer_update_info {:?}", indexer_update_info); - let indexer_root = indexer_update_info - .value - .non_inclusion_proofs - .first() - .unwrap() - .root; - - if indexer_root != current_root { - warn!("Indexer root does not match on-chain root"); - warn!("Indexer root: {:?}", indexer_root); - warn!("On-chain root: {:?}", current_root); - - return Err(ForesterUtilsError::Indexer( - "Indexer root does not match on-chain root".into(), - )); - } +fn calculate_max_zkp_batches_per_call(batch_size: u16) -> usize { + std::cmp::max(1, MAX_PHOTON_ELEMENTS_PER_CALL / batch_size as usize) +} +fn get_all_circuit_inputs_for_chunk( + chunk_hash_chains: &[[u8; 32]], + indexer_update_info: &light_client::indexer::Response< + light_client::indexer::BatchAddressUpdateIndexerResponse, + >, + batch_size: u16, + chunk_start_idx: usize, + global_start_index: u64, + mut current_root: [u8; 32], +) -> Result< + ( + Vec, + [u8; 32], + ), + ForesterUtilsError, +> { let subtrees_array: [[u8; 32]; DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize] = indexer_update_info .value @@ -123,75 +229,32 @@ where ForesterUtilsError::Prover("Failed to convert subtrees to array".into()) })?; - let mut sparse_merkle_tree = SparseMerkleTree::< - Poseidon, - { DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize }, - >::new(subtrees_array, start_index as usize); - - let all_addresses = indexer_update_info - .value - .addresses - .iter() - .map(|x| x.address) - .collect::>(); - - debug!("Got {} addresses from indexer", all_addresses.len()); - - let mut all_inputs = Vec::new(); - let mut current_root = current_root; - - let mut changelog: Vec> = - Vec::new(); - let mut indexed_changelog: Vec< - IndexedChangelogEntry, - > = Vec::new(); - - for (batch_idx, leaves_hash_chain) in leaves_hash_chains.iter().enumerate() { - debug!( - "Preparing circuit inputs for batch {} with root {:?}", - batch_idx, current_root + let mut sparse_merkle_tree = + SparseMerkleTree::::new( + subtrees_array, + global_start_index as usize + (chunk_start_idx * batch_size as usize), ); - let start_addr_idx = batch_idx * batch_size as usize; - let end_addr_idx = start_addr_idx + batch_size as usize; - - if end_addr_idx > all_addresses.len() { - error!( - "Not enough addresses from indexer. Expected at least {}, got {}", - end_addr_idx, - all_addresses.len() - ); - return Err(ForesterUtilsError::Indexer( - "Not enough addresses from indexer".into(), - )); - } - - let batch_addresses = all_addresses[start_addr_idx..end_addr_idx].to_vec(); - - let start_proof_idx = batch_idx * batch_size as usize; - let end_proof_idx = start_proof_idx + batch_size as usize; - - if end_proof_idx > indexer_update_info.value.non_inclusion_proofs.len() { - error!( - "Not enough proofs from indexer. Expected at least {}, got {}", - end_proof_idx, - indexer_update_info.value.non_inclusion_proofs.len() - ); - return Err(ForesterUtilsError::Indexer( - "Not enough proofs from indexer".into(), - )); - } - - let batch_proofs = - &indexer_update_info.value.non_inclusion_proofs[start_proof_idx..end_proof_idx]; + let mut all_inputs = Vec::new(); + let mut changelog = Vec::new(); + let mut indexed_changelog = Vec::new(); + + for (batch_idx, leaves_hash_chain) in chunk_hash_chains.iter().enumerate() { + let start_idx = batch_idx * batch_size as usize; + let end_idx = start_idx + batch_size as usize; + let batch_addresses: Vec<[u8; 32]> = indexer_update_info.value.addresses + [start_idx..end_idx] + .iter() + .map(|x| x.address) + .collect(); let mut low_element_values = Vec::new(); + let mut low_element_next_values = Vec::new(); let mut low_element_indices = Vec::new(); let mut low_element_next_indices = Vec::new(); - let mut low_element_next_values = Vec::new(); - let mut low_element_proofs: Vec> = Vec::new(); + let mut low_element_proofs = Vec::new(); - for proof in batch_proofs { + for proof in &indexer_update_info.value.non_inclusion_proofs[start_idx..end_idx] { low_element_values.push(proof.low_address_value); low_element_indices.push(proof.low_address_index as usize); low_element_next_indices.push(proof.low_address_next_index as usize); @@ -199,34 +262,17 @@ where low_element_proofs.push(proof.low_address_proof.to_vec()); } - let addresses_hashchain = create_hash_chain_from_slice(batch_addresses.as_slice()) - .map_err(|e| { - error!("Failed to create hash chain from addresses: {:?}", e); - ForesterUtilsError::Prover("Failed to create hash chain from addresses".into()) - })?; - - if addresses_hashchain != *leaves_hash_chain { - error!( - "Addresses hash chain does not match leaves hash chain for batch {}", - batch_idx - ); - error!("Addresses hash chain: {:?}", addresses_hashchain); - error!("Leaves hash chain: {:?}", leaves_hash_chain); + if create_hash_chain_from_slice(&batch_addresses)? != *leaves_hash_chain { return Err(ForesterUtilsError::Prover( - "Addresses hash chain does not match leaves hash chain".into(), + "Addresses hash chain does not match".into(), )); } - let adjusted_start_index = start_index as usize + (batch_idx * batch_size as usize); - - debug!( - "Batch {} using root {:?}, start index {}", - batch_idx, current_root, adjusted_start_index - ); + let adjusted_start_index = global_start_index as usize + + (chunk_start_idx * batch_size as usize) + + (batch_idx * batch_size as usize); - let inputs = get_batch_address_append_circuit_inputs::< - { DEFAULT_BATCH_ADDRESS_TREE_HEIGHT as usize }, - >( + let inputs = get_batch_address_append_circuit_inputs( adjusted_start_index, current_root, low_element_values, @@ -241,53 +287,66 @@ where &mut changelog, &mut indexed_changelog, ) - .map_err(|e| { - error!( - "Failed to get circuit inputs for batch {}: {:?}", - batch_idx, e - ); - ForesterUtilsError::Prover(format!( - "Failed to get circuit inputs for batch {}: {}", - batch_idx, e - )) - })?; - - current_root = bigint_to_be_bytes_array::<32>(&inputs.new_root).unwrap(); - debug!("Updated root after batch {}: {:?}", batch_idx, current_root); + .map_err(|e| ForesterUtilsError::Prover(format!("Failed to get circuit inputs: {}", e)))?; + + current_root = bigint_to_be_bytes_array::<32>(&inputs.new_root)?; all_inputs.push(inputs); } - info!("Generating {} ZK proofs asynchronously", all_inputs.len()); - let proof_client = ProofClient::with_config(prover_url, polling_interval, max_wait_time); - let proof_futures = all_inputs - .into_iter() - .map(|inputs| proof_client.generate_batch_address_append_proof(inputs)); - let proof_results = future::join_all(proof_futures).await; - - let mut instruction_data_vec = Vec::new(); - for (i, proof_result) in proof_results.into_iter().enumerate() { - match proof_result { - Ok((compressed_proof, new_root)) => { - debug!("Successfully generated proof for batch {}", i); - instruction_data_vec.push(InstructionDataAddressAppendInputs { - new_root, - compressed_proof: CompressedProof { - a: compressed_proof.a, - b: compressed_proof.b, - c: compressed_proof.c, - }, - }); - } - Err(e) => { - error!("Failed to generate proof for batch {}: {:?}", i, e); - return Err(ForesterUtilsError::Prover(e.to_string())); - } - } + Ok((all_inputs, current_root)) +} + +pub async fn get_address_update_instruction_stream<'a, R, I>( + config: AddressUpdateConfig, + merkle_tree_data: crate::ParsedMerkleTreeData, +) -> Result< + ( + Pin< + Box< + dyn Stream< + Item = Result, ForesterUtilsError>, + > + Send + + 'a, + >, + >, + u16, + ), + ForesterUtilsError, +> +where + R: Rpc + Send + Sync + 'a, + I: Indexer + Send + 'a, +{ + let rpc = config.rpc_pool.get_connection().await?; + let indexer_guard = config.indexer.lock().await; + wait_for_indexer(&*rpc, &*indexer_guard).await?; + drop(rpc); + drop(indexer_guard); + + let (current_root, leaves_hash_chains, start_index, zkp_batch_size) = ( + merkle_tree_data.current_root, + merkle_tree_data.leaves_hash_chains, + merkle_tree_data.next_index, + merkle_tree_data.zkp_batch_size, + ); + + if leaves_hash_chains.is_empty() { + debug!("No hash chains to process for address update, returning empty stream."); + return Ok((Box::pin(futures::stream::empty()), zkp_batch_size)); } - info!( - "Successfully generated {} instruction data entries", - instruction_data_vec.len() + let stream = stream_instruction_data( + config.indexer, + config.merkle_tree_pubkey, + config.prover_url, + config.polling_interval, + config.max_wait_time, + leaves_hash_chains, + start_index, + zkp_batch_size, + current_root, + config.ixs_per_tx, ); - Ok((instruction_data_vec, batch_size)) + + Ok((Box::pin(stream), zkp_batch_size)) } diff --git a/forester-utils/src/instructions/state_batch_append.rs b/forester-utils/src/instructions/state_batch_append.rs index 9929739300..1a931c19f1 100644 --- a/forester-utils/src/instructions/state_batch_append.rs +++ b/forester-utils/src/instructions/state_batch_append.rs @@ -1,10 +1,13 @@ -use std::{sync::Arc, time::Duration}; +use std::{pin::Pin, sync::Arc, time::Duration}; use account_compression::processor::initialize_address_merkle_tree::Pubkey; +use async_stream::stream; +use futures::{ + stream::{FuturesOrdered, Stream}, + StreamExt, +}; use light_batched_merkle_tree::{ - constants::DEFAULT_BATCH_STATE_TREE_HEIGHT, - merkle_tree::{BatchedMerkleTreeAccount, InstructionDataBatchAppendInputs}, - queue::BatchedQueueAccount, + constants::DEFAULT_BATCH_STATE_TREE_HEIGHT, merkle_tree::InstructionDataBatchAppendInputs, }; use light_client::{indexer::Indexer, rpc::Rpc}; use light_compressed_account::instruction_data::compressed_proof::CompressedProof; @@ -15,238 +18,201 @@ use light_prover_client::{ proof_types::batch_append::{get_batch_append_inputs, BatchAppendsCircuitInputs}, }; use light_sparse_merkle_tree::changelog::ChangelogEntry; -use tracing::{error, trace}; - -use crate::{error::ForesterUtilsError, utils::wait_for_indexer}; - -pub async fn create_append_batch_ix_data( - rpc: &mut R, - indexer: &mut I, - merkle_tree_pubkey: Pubkey, - output_queue_pubkey: Pubkey, - prover_url: String, - polling_interval: Duration, - max_wait_time: Duration, -) -> Result, ForesterUtilsError> { - trace!("Creating append batch instruction data"); - - let (merkle_tree_next_index, current_root, root_history) = - get_merkle_tree_metadata(rpc, merkle_tree_pubkey).await?; - - trace!( - "merkle_tree_next_index: {:?} current_root: {:?}", - merkle_tree_next_index, - current_root - ); - - // Get output queue metadata and hash chains - let (zkp_batch_size, leaves_hash_chains) = - get_output_queue_metadata(rpc, output_queue_pubkey).await?; - - if leaves_hash_chains.is_empty() { - trace!("No hash chains to process"); - return Ok(Vec::new()); - } - - wait_for_indexer(rpc, indexer).await?; - - let total_elements = zkp_batch_size as usize * leaves_hash_chains.len(); - let offset = merkle_tree_next_index; - - let queue_elements = indexer - .get_queue_elements( - merkle_tree_pubkey.to_bytes(), - QueueType::OutputStateV2, - total_elements as u16, - Some(offset), - None, - ) - .await - .map_err(|e| { - error!("Failed to get queue elements from indexer: {:?}", e); - ForesterUtilsError::Indexer("Failed to get queue elements".into()) - })? - .value - .items; - - trace!("Got {} queue elements in total", queue_elements.len()); - - if queue_elements.len() != total_elements { - return Err(ForesterUtilsError::Indexer(format!( - "Expected {} elements, got {}", - total_elements, - queue_elements.len() - ))); - } - let indexer_root = queue_elements.first().unwrap().root; - debug_assert_eq!( - indexer_root, current_root, - "root_history: {:?}", - root_history - ); +use tokio::sync::Mutex; +use tracing::trace; - let mut current_root = current_root; - let mut all_changelogs: Vec> = - Vec::new(); - let mut proof_futures = Vec::new(); - let proof_client = Arc::new(ProofClient::with_config( - prover_url.clone(), - polling_interval, - max_wait_time, - )); - - for (batch_idx, leaves_hash_chain) in leaves_hash_chains.iter().enumerate() { - let start_idx = batch_idx * zkp_batch_size as usize; - let end_idx = start_idx + zkp_batch_size as usize; - let batch_elements = &queue_elements[start_idx..end_idx]; - - trace!( - "Processing batch {}: index range {}-{}", - batch_idx, - start_idx, - end_idx - ); - - let old_leaves = batch_elements - .iter() - .map(|x| x.leaf) - .collect::>(); - - let leaves = batch_elements - .iter() - .map(|x| x.account_hash) - .collect::>(); - - let merkle_proofs = batch_elements - .iter() - .map(|x| x.proof.clone()) - .collect::>>(); - - let adjusted_start_index = - merkle_tree_next_index as u32 + (batch_idx * zkp_batch_size as usize) as u32; - - let (circuit_inputs, batch_changelogs) = get_batch_append_inputs::<32>( - current_root, - adjusted_start_index, - leaves, - *leaves_hash_chain, - old_leaves, - merkle_proofs, - zkp_batch_size as u32, - all_changelogs.as_slice(), - ) - .map_err(|e| { - error!("Failed to get circuit inputs: {:?}", e); - ForesterUtilsError::Prover("Failed to get circuit inputs".into()) - })?; - - current_root = - bigint_to_be_bytes_array::<32>(&circuit_inputs.new_root.to_biguint().unwrap()).unwrap(); - all_changelogs.extend(batch_changelogs); - - let client = Arc::clone(&proof_client); - let proof_future = generate_zkp_proof(circuit_inputs, client); - - proof_futures.push(proof_future); - } - - let proof_results = futures::future::join_all(proof_futures).await; - let mut instruction_data_vec = Vec::new(); - - for (i, proof_result) in proof_results.into_iter().enumerate() { - match proof_result { - Ok((proof, new_root)) => { - trace!("Successfully generated proof for batch {}", i); - instruction_data_vec.push(InstructionDataBatchAppendInputs { - new_root, - compressed_proof: proof, - }); - } - Err(e) => { - error!("Failed to generate proof for batch {}: {:?}", i, e); - return Err(e); - } - } - } +use crate::{ + error::ForesterUtilsError, rpc_pool::SolanaRpcPool, utils::wait_for_indexer, + ParsedMerkleTreeData, ParsedQueueData, +}; - Ok(instruction_data_vec) -} async fn generate_zkp_proof( circuit_inputs: BatchAppendsCircuitInputs, proof_client: Arc, -) -> Result<(CompressedProof, [u8; 32]), ForesterUtilsError> { +) -> Result { let (proof, new_root) = proof_client .generate_batch_append_proof(circuit_inputs) .await .map_err(|e| ForesterUtilsError::Prover(e.to_string()))?; - Ok(( - CompressedProof { + Ok(InstructionDataBatchAppendInputs { + new_root, + compressed_proof: CompressedProof { a: proof.a, b: proof.b, c: proof.c, }, - new_root, - )) + }) } -/// Get metadata from the Merkle tree account -async fn get_merkle_tree_metadata( - rpc: &mut impl Rpc, +#[allow(clippy::too_many_arguments)] +pub async fn get_append_instruction_stream<'a, R, I>( + rpc_pool: Arc>, + indexer: Arc>, merkle_tree_pubkey: Pubkey, -) -> Result<(u64, [u8; 32], Vec<[u8; 32]>), ForesterUtilsError> { - let mut merkle_tree_account = rpc - .get_account(merkle_tree_pubkey) - .await - .map_err(|e| ForesterUtilsError::Rpc(format!("Failed to get merkle tree account: {}", e)))? - .ok_or_else(|| ForesterUtilsError::Rpc("Merkle tree account not found".into()))?; - - let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( - merkle_tree_account.data.as_mut_slice(), - &merkle_tree_pubkey.into(), - ) - .map_err(|e| ForesterUtilsError::Rpc(format!("Failed to parse merkle tree: {}", e)))?; - - Ok(( - merkle_tree.next_index, - *merkle_tree.root_history.last().unwrap(), - merkle_tree.root_history.to_vec(), - )) -} + prover_url: String, + polling_interval: Duration, + max_wait_time: Duration, + merkle_tree_data: ParsedMerkleTreeData, + output_queue_data: ParsedQueueData, + yield_batch_size: usize, +) -> Result< + ( + Pin< + Box< + dyn Stream, ForesterUtilsError>> + + Send + + 'a, + >, + >, + u16, + ), + ForesterUtilsError, +> +where + R: Rpc + Send + Sync + 'a, + I: Indexer + Send + 'a, +{ + trace!("Initializing append batch instruction stream with parsed data"); + + let (indexer_guard, rpc_result) = tokio::join!(indexer.lock(), rpc_pool.get_connection()); + let rpc = rpc_result?; + + let (merkle_tree_next_index, mut current_root, _) = ( + merkle_tree_data.next_index, + merkle_tree_data.current_root, + merkle_tree_data.root_history, + ); + let (zkp_batch_size, leaves_hash_chains) = ( + output_queue_data.zkp_batch_size, + output_queue_data.leaves_hash_chains, + ); -/// Get metadata and hash chains from the output queue -async fn get_output_queue_metadata( - rpc: &mut impl Rpc, - output_queue_pubkey: Pubkey, -) -> Result<(u16, Vec<[u8; 32]>), ForesterUtilsError> { - let mut output_queue_account = rpc - .get_account(output_queue_pubkey) - .await - .map_err(|e| ForesterUtilsError::Rpc(format!("Failed to get output queue account: {}", e)))? - .ok_or_else(|| ForesterUtilsError::Rpc("Output queue account not found".into()))?; - - let output_queue = - BatchedQueueAccount::output_from_bytes(output_queue_account.data.as_mut_slice()) - .map_err(|e| ForesterUtilsError::Rpc(format!("Failed to parse output queue: {}", e)))?; - - let full_batch_index = output_queue.batch_metadata.pending_batch_index; - let zkp_batch_size = output_queue.batch_metadata.zkp_batch_size; - let batch = &output_queue.batch_metadata.batches[full_batch_index as usize]; - let num_inserted_zkps = batch.get_num_inserted_zkps(); - - // Get all remaining hash chains for the batch - let mut leaves_hash_chains = Vec::new(); - for i in num_inserted_zkps..batch.get_current_zkp_batch_index() { - leaves_hash_chains - .push(output_queue.hash_chain_stores[full_batch_index as usize][i as usize]); + if leaves_hash_chains.is_empty() { + trace!("No hash chains to process, returning empty stream."); + return Ok((Box::pin(futures::stream::empty()), zkp_batch_size)); } - trace!( - "ZKP batch size: {}, inserted ZKPs: {}, current ZKP index: {}, ready for insertion: {}", - zkp_batch_size, - num_inserted_zkps, - batch.get_current_zkp_batch_index(), - leaves_hash_chains.len() - ); + wait_for_indexer(&*rpc, &*indexer_guard).await?; + drop(rpc); + drop(indexer_guard); + + let stream = stream! { + let total_elements = zkp_batch_size as usize * leaves_hash_chains.len(); + let offset = merkle_tree_next_index; + + let queue_elements = { + let mut indexer_guard = indexer.lock().await; + + match indexer_guard + .get_queue_elements( + merkle_tree_pubkey.to_bytes(), + QueueType::OutputStateV2, + total_elements as u16, + Some(offset), + None, + ) + .await { + Ok(res) => res.value.items, + Err(e) => { + yield Err(ForesterUtilsError::Indexer(format!("Failed to get queue elements: {}", e))); + return; + } + } + }; + + if queue_elements.len() != total_elements { + yield Err(ForesterUtilsError::Indexer(format!( + "Expected {} elements, got {}", + total_elements, + queue_elements.len() + ))); + return; + } + + if let Some(first_element) = queue_elements.first() { + if first_element.root != current_root { + yield Err(ForesterUtilsError::Indexer("Root mismatch between indexer and on-chain state".into())); + return; + } + } + + let mut all_changelogs: Vec> = Vec::new(); + let proof_client = Arc::new(ProofClient::with_config(prover_url.clone(), polling_interval, max_wait_time)); + let mut futures_ordered = FuturesOrdered::new(); + let mut pending_count = 0; + + let mut proof_buffer = Vec::new(); + + for (batch_idx, leaves_hash_chain) in leaves_hash_chains.iter().enumerate() { + let start_idx = batch_idx * zkp_batch_size as usize; + let end_idx = start_idx + zkp_batch_size as usize; + let batch_elements = &queue_elements[start_idx..end_idx]; + + let old_leaves: Vec<[u8; 32]> = batch_elements.iter().map(|x| x.leaf).collect(); + let leaves: Vec<[u8; 32]> = batch_elements.iter().map(|x| x.account_hash).collect(); + let merkle_proofs: Vec> = batch_elements.iter().map(|x| x.proof.clone()).collect(); + let adjusted_start_index = merkle_tree_next_index as u32 + (batch_idx * zkp_batch_size as usize) as u32; + + let (circuit_inputs, batch_changelogs) = match get_batch_append_inputs::<32>( + current_root, adjusted_start_index, leaves, *leaves_hash_chain, old_leaves, merkle_proofs, zkp_batch_size as u32, &all_changelogs, + ) { + Ok(inputs) => inputs, + Err(e) => { + yield Err(ForesterUtilsError::Prover(format!("Failed to get circuit inputs: {}", e))); + return; + } + }; + + current_root = bigint_to_be_bytes_array::<32>(&circuit_inputs.new_root.to_biguint().unwrap()).unwrap(); + all_changelogs.extend(batch_changelogs); + + let client = Arc::clone(&proof_client); + futures_ordered.push_back(generate_zkp_proof(circuit_inputs, client)); + pending_count += 1; + + while pending_count >= yield_batch_size { + for _ in 0..yield_batch_size.min(pending_count) { + if let Some(result) = futures_ordered.next().await { + match result { + Ok(proof) => proof_buffer.push(proof), + Err(e) => { + yield Err(e); + return; + } + } + pending_count -= 1; + } + } + + if !proof_buffer.is_empty() { + yield Ok(proof_buffer.clone()); + proof_buffer.clear(); + } + } + } + + while let Some(result) = futures_ordered.next().await { + match result { + Ok(proof) => { + proof_buffer.push(proof); + + if proof_buffer.len() >= yield_batch_size { + yield Ok(proof_buffer.clone()); + proof_buffer.clear(); + } + }, + Err(e) => { + yield Err(e); + return; + } + } + } + + // Yield any remaining proofs + if !proof_buffer.is_empty() { + yield Ok(proof_buffer); + } + }; - Ok((zkp_batch_size as u16, leaves_hash_chains)) + Ok((Box::pin(stream), zkp_batch_size)) } diff --git a/forester-utils/src/instructions/state_batch_nullify.rs b/forester-utils/src/instructions/state_batch_nullify.rs index 4bfe0004d3..ecf4ed6e13 100644 --- a/forester-utils/src/instructions/state_batch_nullify.rs +++ b/forester-utils/src/instructions/state_batch_nullify.rs @@ -1,260 +1,223 @@ -use std::{sync::Arc, time::Duration}; +use std::{pin::Pin, sync::Arc, time::Duration}; use account_compression::processor::initialize_address_merkle_tree::Pubkey; +use async_stream::stream; +use futures::{ + stream::{FuturesOrdered, Stream}, + StreamExt, +}; use light_batched_merkle_tree::{ - constants::DEFAULT_BATCH_STATE_TREE_HEIGHT, - merkle_tree::{BatchedMerkleTreeAccount, InstructionDataBatchNullifyInputs}, + constants::DEFAULT_BATCH_STATE_TREE_HEIGHT, merkle_tree::InstructionDataBatchNullifyInputs, }; use light_client::{indexer::Indexer, rpc::Rpc}; use light_compressed_account::instruction_data::compressed_proof::CompressedProof; -use light_hasher::{bigint::bigint_to_be_bytes_array, Hasher, Poseidon}; +use light_hasher::bigint::bigint_to_be_bytes_array; use light_merkle_tree_metadata::QueueType; use light_prover_client::{ proof_client::ProofClient, proof_types::batch_update::{get_batch_update_inputs, BatchUpdateCircuitInputs}, }; -use tracing::{error, trace}; +use tokio::sync::Mutex; +use tracing::{debug, trace}; -use crate::{error::ForesterUtilsError, utils::wait_for_indexer}; +use crate::{ + error::ForesterUtilsError, rpc_pool::SolanaRpcPool, utils::wait_for_indexer, + ParsedMerkleTreeData, +}; -pub async fn create_nullify_batch_ix_data( - rpc: &mut R, - indexer: &mut I, +async fn generate_nullify_zkp_proof( + inputs: BatchUpdateCircuitInputs, + proof_client: Arc, +) -> Result { + let (proof, new_root) = proof_client + .generate_batch_update_proof(inputs) + .await + .map_err(|e| ForesterUtilsError::Prover(e.to_string()))?; + Ok(InstructionDataBatchNullifyInputs { + new_root, + compressed_proof: CompressedProof { + a: proof.a, + b: proof.b, + c: proof.c, + }, + }) +} + +#[allow(clippy::too_many_arguments)] +pub async fn get_nullify_instruction_stream<'a, R, I>( + rpc_pool: Arc>, + indexer: Arc>, merkle_tree_pubkey: Pubkey, prover_url: String, polling_interval: Duration, max_wait_time: Duration, -) -> Result, ForesterUtilsError> { - trace!("create_multiple_nullify_batch_ix_data"); - // Get the tree information and find out how many ZKP batches need processing - let ( - batch_idx, - zkp_batch_size, - num_inserted_zkps, - num_ready_zkps, - old_root, - root_history, - leaves_hash_chains, - ) = { - let mut account = rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap(); - let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( - account.data.as_mut_slice(), - &merkle_tree_pubkey.into(), - ) - .unwrap(); - - trace!("queue_batches: {:?}", merkle_tree.queue_batches); - - let batch_idx = merkle_tree.queue_batches.pending_batch_index as usize; - let zkp_size = merkle_tree.queue_batches.zkp_batch_size; - let batch = &merkle_tree.queue_batches.batches[batch_idx]; - let num_inserted_zkps = batch.get_num_inserted_zkps(); - let num_current_zkp = batch.get_current_zkp_batch_index(); - let num_ready_zkps = num_current_zkp.saturating_sub(num_inserted_zkps); - - let mut leaves_hash_chains = Vec::new(); - for i in num_inserted_zkps..num_current_zkp { - leaves_hash_chains.push(merkle_tree.hash_chain_stores[batch_idx][i as usize]); - } - - let root = *merkle_tree.root_history.last().unwrap(); - let root_history = merkle_tree.root_history.to_vec(); - - ( - batch_idx, - zkp_size as u16, - num_inserted_zkps, - num_ready_zkps, - root, - root_history, - leaves_hash_chains, - ) - }; - - trace!( - "batch_idx: {}, zkp_batch_size: {}, num_inserted_zkps: {}, num_ready_zkps: {}, leaves_hash_chains: {:?}", - batch_idx, zkp_batch_size, num_inserted_zkps, num_ready_zkps, leaves_hash_chains.len() + merkle_tree_data: ParsedMerkleTreeData, + yield_batch_size: usize, +) -> Result< + ( + Pin< + Box< + dyn Stream< + Item = Result, ForesterUtilsError>, + > + Send + + 'a, + >, + >, + u16, + ), + ForesterUtilsError, +> +where + R: Rpc + Send + Sync + 'a, + I: Indexer + Send + 'a, +{ + let rpc = rpc_pool.get_connection().await?; + + let (mut current_root, leaves_hash_chains, num_inserted_zkps, zkp_batch_size) = ( + merkle_tree_data.current_root, + merkle_tree_data.leaves_hash_chains, + merkle_tree_data.num_inserted_zkps, + merkle_tree_data.zkp_batch_size, ); if leaves_hash_chains.is_empty() { - return Ok(Vec::new()); - } - - wait_for_indexer(rpc, indexer).await?; - - let current_slot = rpc.get_slot().await.unwrap(); - trace!("current_slot: {}", current_slot); - - let total_elements = zkp_batch_size as usize * leaves_hash_chains.len(); - let offset = num_inserted_zkps * zkp_batch_size as u64; - - trace!( - "Requesting {} total elements with offset {}", - total_elements, - offset - ); - - let all_queue_elements = indexer - .get_queue_elements( - merkle_tree_pubkey.to_bytes(), - QueueType::InputStateV2, - total_elements as u16, - Some(offset), - None, - ) - .await - .map_err(|e| { - error!( - "create_multiple_nullify_batch_ix_data: failed to get queue elements from indexer: {:?}", - e - ); - ForesterUtilsError::Indexer("Failed to get queue elements".into()) - })?.value.items; - - trace!("Got {} queue elements in total", all_queue_elements.len()); - if all_queue_elements.len() != total_elements { - return Err(ForesterUtilsError::Indexer(format!( - "Expected {} elements, got {}", - total_elements, - all_queue_elements.len() - ))); + debug!("No hash chains to process for nullification, returning empty stream."); + return Ok((Box::pin(futures::stream::empty()), zkp_batch_size)); } - let indexer_root = all_queue_elements.first().unwrap().root; - debug_assert_eq!( - indexer_root, old_root, - "Root mismatch. Expected: {:?}, Got: {:?}. Root history: {:?}", - old_root, indexer_root, root_history - ); - - let mut all_changelogs = Vec::new(); - let mut proof_futures = Vec::new(); - let proof_client = Arc::new(ProofClient::with_config( - prover_url.clone(), - polling_interval, - max_wait_time, - )); - - let mut current_root = old_root; - - for (batch_offset, leaves_hash_chain) in leaves_hash_chains.iter().enumerate() { - let start_idx = batch_offset * zkp_batch_size as usize; - let end_idx = start_idx + zkp_batch_size as usize; - let batch_elements = &all_queue_elements[start_idx..end_idx]; - - trace!( - "Processing batch {} with offset {}-{}", - batch_offset, - start_idx, - end_idx - ); + let indexer_guard = indexer.lock().await; + wait_for_indexer(&*rpc, &*indexer_guard).await?; + drop(rpc); + drop(indexer_guard); + + let stream = stream! { + let total_elements = zkp_batch_size as usize * leaves_hash_chains.len(); + let offset = num_inserted_zkps * zkp_batch_size as u64; + + trace!("Requesting {} total elements with offset {}", total_elements, offset); + + let all_queue_elements = { + let mut indexer_guard = indexer.lock().await; + indexer_guard + .get_queue_elements( + merkle_tree_pubkey.to_bytes(), + QueueType::InputStateV2, + total_elements as u16, + Some(offset), + None, + ) + .await + }; - // Process this batch's data - let mut leaves = Vec::new(); - let mut tx_hashes = Vec::new(); - let mut old_leaves = Vec::new(); - let mut path_indices = Vec::new(); - let mut merkle_proofs = Vec::new(); - let mut nullifiers = Vec::new(); + let all_queue_elements = match all_queue_elements { + Ok(res) => res.value.items, + Err(e) => { + yield Err(ForesterUtilsError::Indexer(format!("Failed to get queue elements: {}", e))); + return; + } + }; + + trace!("Got {} queue elements in total", all_queue_elements.len()); + if all_queue_elements.len() != total_elements { + yield Err(ForesterUtilsError::Indexer(format!( + "Expected {} elements, got {}", + total_elements, all_queue_elements.len() + ))); + return; + } - for (i, leaf_info) in batch_elements.iter().enumerate() { - let global_leaf_index = start_idx + i; - trace!( - "Element {}: local index={}, global index={}, reported index={}", - i, - i, - global_leaf_index, - leaf_info.leaf_index - ); + if let Some(first_element) = all_queue_elements.first() { + if first_element.root != current_root { + yield Err(ForesterUtilsError::Indexer("Root mismatch between indexer and on-chain state".into())); + return; + } + } - path_indices.push(leaf_info.leaf_index as u32); - leaves.push(leaf_info.account_hash); - old_leaves.push(leaf_info.leaf); - merkle_proofs.push(leaf_info.proof.clone()); + let mut all_changelogs = Vec::new(); + let proof_client = Arc::new(ProofClient::with_config(prover_url.clone(), polling_interval, max_wait_time)); + let mut futures_ordered = FuturesOrdered::new(); + let mut pending_count = 0; + + let mut proof_buffer = Vec::new(); + + for (batch_offset, leaves_hash_chain) in leaves_hash_chains.iter().enumerate() { + let start_idx = batch_offset * zkp_batch_size as usize; + let end_idx = start_idx + zkp_batch_size as usize; + let batch_elements = &all_queue_elements[start_idx..end_idx]; + + let mut leaves = Vec::new(); + let mut tx_hashes = Vec::new(); + let mut old_leaves = Vec::new(); + let mut path_indices = Vec::new(); + let mut merkle_proofs = Vec::new(); + + for leaf_info in batch_elements.iter() { + path_indices.push(leaf_info.leaf_index as u32); + leaves.push(leaf_info.account_hash); + old_leaves.push(leaf_info.leaf); + merkle_proofs.push(leaf_info.proof.clone()); + tx_hashes.push(leaf_info.tx_hash.ok_or_else(|| ForesterUtilsError::Indexer(format!("Missing tx_hash for leaf index {}", leaf_info.leaf_index)))?); + } - // Make sure tx_hash exists - let tx_hash = match leaf_info.tx_hash { - Some(hash) => hash, - None => { - return Err(ForesterUtilsError::Indexer(format!( - "Missing tx_hash for leaf index {}", - leaf_info.leaf_index - ))) + let (circuit_inputs, batch_changelog) = match get_batch_update_inputs::<{ DEFAULT_BATCH_STATE_TREE_HEIGHT as usize }>( + current_root, tx_hashes, leaves, *leaves_hash_chain, old_leaves, merkle_proofs, path_indices, zkp_batch_size as u32, &all_changelogs, + ) { + Ok(inputs) => inputs, + Err(e) => { + yield Err(ForesterUtilsError::Prover(format!("Failed to get batch update inputs: {}", e))); + return; } }; - tx_hashes.push(tx_hash); + all_changelogs.extend(batch_changelog); + current_root = bigint_to_be_bytes_array::<32>(&circuit_inputs.new_root.to_biguint().unwrap()).unwrap(); + + let client = Arc::clone(&proof_client); + futures_ordered.push_back(generate_nullify_zkp_proof(circuit_inputs, client)); + pending_count += 1; + + while pending_count >= yield_batch_size { + for _ in 0..yield_batch_size.min(pending_count) { + if let Some(result) = futures_ordered.next().await { + match result { + Ok(proof) => proof_buffer.push(proof), + Err(e) => { + yield Err(e); + return; + } + } + pending_count -= 1; + } + } - let index_bytes = leaf_info.leaf_index.to_be_bytes(); - let nullifier = - Poseidon::hashv(&[&leaf_info.account_hash, &index_bytes, &tx_hash]).unwrap(); - nullifiers.push(nullifier); + if !proof_buffer.is_empty() { + yield Ok(proof_buffer.clone()); + proof_buffer.clear(); + } + } } - let (circuit_inputs, batch_changelog) = - get_batch_update_inputs::<{ DEFAULT_BATCH_STATE_TREE_HEIGHT as usize }>( - current_root, - tx_hashes.clone(), - leaves.clone(), - *leaves_hash_chain, - old_leaves.clone(), - merkle_proofs.clone(), - path_indices.clone(), - zkp_batch_size as u32, - &all_changelogs, - ) - .map_err(|e| { - error!("Failed to get batch update inputs: {:?}", e); - ForesterUtilsError::Prover("Failed to get batch update inputs".into()) - })?; - - all_changelogs.extend(batch_changelog); - current_root = - bigint_to_be_bytes_array::<32>(&circuit_inputs.new_root.to_biguint().unwrap()) - .map_err(|_| { - ForesterUtilsError::Prover("Failed to convert new root to bytes".into()) - })?; - - let client = Arc::clone(&proof_client); - let proof_future = generate_nullify_zkp_proof(circuit_inputs, client); - proof_futures.push(proof_future); - } - - let proof_results = futures::future::join_all(proof_futures).await; - let mut instruction_data_vec = Vec::new(); - - for (i, proof_result) in proof_results.into_iter().enumerate() { - match proof_result { - Ok((proof, new_root)) => { - trace!("Successfully generated proof for batch {}", i); - instruction_data_vec.push(InstructionDataBatchNullifyInputs { - new_root, - compressed_proof: proof, - }); - } - Err(e) => { - error!("Failed to generate proof for batch {}: {:?}", i, e); - return Err(e); + while let Some(result) = futures_ordered.next().await { + match result { + Ok(proof) => { + proof_buffer.push(proof); + + if proof_buffer.len() >= yield_batch_size { + yield Ok(proof_buffer.clone()); + proof_buffer.clear(); + } + }, + Err(e) => { + yield Err(e); + return; + } } } - } - Ok(instruction_data_vec) -} -async fn generate_nullify_zkp_proof( - inputs: BatchUpdateCircuitInputs, - proof_client: Arc, -) -> Result<(CompressedProof, [u8; 32]), ForesterUtilsError> { - let (proof, new_root) = proof_client - .generate_batch_update_proof(inputs) - .await - .map_err(|e| ForesterUtilsError::Prover(e.to_string()))?; - Ok(( - CompressedProof { - a: proof.a, - b: proof.b, - c: proof.c, - }, - new_root, - )) + if !proof_buffer.is_empty() { + yield Ok(proof_buffer); + } + }; + + Ok((Box::pin(stream), zkp_batch_size)) } diff --git a/forester-utils/src/lib.rs b/forester-utils/src/lib.rs index 206033cb64..646d17662e 100644 --- a/forester-utils/src/lib.rs +++ b/forester-utils/src/lib.rs @@ -1,9 +1,32 @@ pub mod account_zero_copy; pub mod address_merkle_tree_config; -mod error; +pub mod error; pub mod forester_epoch; pub mod instructions; pub mod rate_limiter; pub mod registry; pub mod rpc_pool; pub mod utils; + +/// Parsed merkle tree data extracted from account +#[derive(Debug, Clone)] +pub struct ParsedMerkleTreeData { + pub next_index: u64, + pub current_root: [u8; 32], + pub root_history: Vec<[u8; 32]>, + pub zkp_batch_size: u16, + pub pending_batch_index: u32, + pub num_inserted_zkps: u64, + pub current_zkp_batch_index: u64, + pub leaves_hash_chains: Vec<[u8; 32]>, +} + +/// Parsed output queue data extracted from account +#[derive(Debug, Clone)] +pub struct ParsedQueueData { + pub zkp_batch_size: u16, + pub pending_batch_index: u32, + pub num_inserted_zkps: u64, + pub current_zkp_batch_index: u64, + pub leaves_hash_chains: Vec<[u8; 32]>, +} diff --git a/forester-utils/src/rpc_pool.rs b/forester-utils/src/rpc_pool.rs index cc9037ff30..42438aea1f 100644 --- a/forester-utils/src/rpc_pool.rs +++ b/forester-utils/src/rpc_pool.rs @@ -27,6 +27,7 @@ pub enum PoolError { pub struct SolanaConnectionManager { url: String, photon_url: Option, + api_key: Option, commitment: CommitmentConfig, // TODO: implement Rpc for SolanaConnectionManager and rate limit requests. _rpc_rate_limiter: Option, @@ -38,6 +39,7 @@ impl SolanaConnectionManager { pub fn new( url: String, photon_url: Option, + api_key: Option, commitment: CommitmentConfig, rpc_rate_limiter: Option, send_tx_rate_limiter: Option, @@ -45,6 +47,7 @@ impl SolanaConnectionManager { Self { url, photon_url, + api_key, commitment, _rpc_rate_limiter: rpc_rate_limiter, _send_tx_rate_limiter: send_tx_rate_limiter, @@ -64,6 +67,7 @@ impl bb8::ManageConnection for SolanaConnectionManager { photon_url: self.photon_url.clone(), commitment_config: Some(self.commitment), fetch_active_tree: false, + api_key: self.api_key.clone(), }; Ok(R::new(config).await?) @@ -90,7 +94,7 @@ pub struct SolanaRpcPool { pub struct SolanaRpcPoolBuilder { url: Option, photon_url: Option, - + api_key: Option, commitment: Option, max_size: u32, @@ -116,6 +120,7 @@ impl SolanaRpcPoolBuilder { Self { url: None, photon_url: None, + api_key: None, commitment: None, max_size: 50, connection_timeout_secs: 15, @@ -195,6 +200,7 @@ impl SolanaRpcPoolBuilder { let manager = SolanaConnectionManager::new( url, self.photon_url, + self.api_key, commitment, self.rpc_rate_limiter, self.send_tx_rate_limiter, diff --git a/forester-utils/src/utils.rs b/forester-utils/src/utils.rs index 8640d36f53..c99a8bb14b 100644 --- a/forester-utils/src/utils.rs +++ b/forester-utils/src/utils.rs @@ -29,7 +29,7 @@ pub async fn airdrop_lamports( } pub async fn wait_for_indexer( - rpc: &mut R, + rpc: &R, indexer: &I, ) -> Result<(), ForesterUtilsError> { let rpc_slot = rpc diff --git a/forester/Cargo.toml b/forester/Cargo.toml index 65c8f81081..79601efe11 100644 --- a/forester/Cargo.toml +++ b/forester/Cargo.toml @@ -32,6 +32,7 @@ serde = { version = "1.0", features = ["derive"] } tokio = { version = "1", features = ["full"] } reqwest = { workspace = true, features = ["json", "rustls-tls", "blocking"] } futures = "0.3.31" +async-stream = "0.3" thiserror = { workspace = true } borsh = { workspace = true } bs58 = "0.5.1" @@ -52,7 +53,7 @@ num-bigint = { workspace = true } [dev-dependencies] serial_test = { workspace = true } -light-prover-client = { workspace = true , features =[ "devenv"]} +light-prover-client = { workspace = true, features = ["devenv"] } light-test-utils = { workspace = true } light-program-test = { workspace = true, features = ["devenv"] } light-batched-merkle-tree = { workspace = true, features = ["test-only"] } diff --git a/forester/README.md b/forester/README.md index c603d72663..38bf3c2161 100644 --- a/forester/README.md +++ b/forester/README.md @@ -2,7 +2,7 @@ ## Description -Forester is a service for nullifying the state and address merkle trees. +Forester is a service for nullifying the state and address merkle trees. It subscribes to the nullifier queue and nullifies merkle tree leaves. ## Commands @@ -91,4 +91,78 @@ All configuration options can be set using environment variables with the `FORES ```bash export FORESTER_RPC_URL="your-rpc-url-here" -``` \ No newline at end of file +``` + +### Test Environment Variables + +The following environment variables are used for running the e2e_v2 tests: + +#### Test Mode + +- `TEST_MODE` - Specifies whether to run tests on local validator or devnet (values: `local` or `devnet`, default: `devnet`) + +#### Test Feature Flags + +Control which test scenarios to run (all default to `true`): + +- `TEST_V1_STATE` - Enable/disable V1 state tree testing (`true`/`false`) +- `TEST_V2_STATE` - Enable/disable V2 state tree testing (`true`/`false`) +- `TEST_V1_ADDRESS` - Enable/disable V1 address tree testing (`true`/`false`) +- `TEST_V2_ADDRESS` - Enable/disable V2 address tree testing (`true`/`false`) + +#### Required for Devnet mode: + +- `PHOTON_RPC_URL` - Photon RPC endpoint URL +- `PHOTON_WSS_RPC_URL` - Photon WebSocket RPC endpoint URL +- `PHOTON_INDEXER_URL` - Photon indexer endpoint URL +- `PHOTON_PROVER_URL` - Photon prover endpoint URL +- `PHOTON_API_KEY` - Photon API key for authentication + +#### Required for both modes: + +- `FORESTER_KEYPAIR` - Keypair for testing (supports both base58 format and byte array format like `[1,2,3,...]`) + +#### Example configurations: + +**Local validator mode with all tests:** +```bash +export TEST_MODE="local" +export FORESTER_KEYPAIR="your-base58-encoded-keypair" +# OR using byte array format: +# export FORESTER_KEYPAIR="[1,2,3,...]" +``` + +**Local validator mode with only V1 tests:** +```bash +export TEST_MODE="local" +export TEST_V1_STATE="true" +export TEST_V2_STATE="false" +export TEST_V1_ADDRESS="true" +export TEST_V2_ADDRESS="false" +export FORESTER_KEYPAIR="your-base58-encoded-keypair" +``` + +**Devnet mode with only V2 tests:** +```bash +export TEST_MODE="devnet" +export TEST_V1_STATE="false" +export TEST_V2_STATE="true" +export TEST_V1_ADDRESS="false" +export TEST_V2_ADDRESS="true" +export PHOTON_RPC_URL="https://devnet.helius-rpc.com/?api-key=your-key" +export PHOTON_WSS_RPC_URL="wss://devnet.helius-rpc.com/?api-key=your-key" +export PHOTON_INDEXER_URL="https://devnet.helius-rpc.com" +export PHOTON_PROVER_URL="https://devnet.helius-rpc.com" +export PHOTON_API_KEY="your-api-key" +export FORESTER_KEYPAIR="your-base58-encoded-keypair" +``` + +When running in local mode, the test will: +- Spawn a local validator +- Start a local prover service +- Use predefined local URLs (localhost:8899 for RPC, localhost:8784 for indexer, etc.) + +The test will automatically: +- Skip minting tokens for disabled test types +- Skip executing transactions for disabled test types +- Skip root verification for disabled test types diff --git a/forester/package.json b/forester/package.json index 6fb31b8527..ad3a801ce6 100644 --- a/forester/package.json +++ b/forester/package.json @@ -5,13 +5,11 @@ "scripts": { "build": "cargo build", "test": "RUSTFLAGS=\"--cfg tokio_unstable -D warnings\" cargo test --package forester -- --nocapture", - "test-state-batched-local": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_state_batched -- --nocapture", - "test-state-batched-indexer": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_state_indexer_batched -- --nocapture", - "test-state-batched-indexer-async": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_state_indexer_async_batched -- --nocapture", - "test-fetch-root": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_state_indexer_fetch_root -- --nocapture", + "e2e-test": "source .env && RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_e2e_v2 -- --nocapture", "test-address-batched-local": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_address_batched -- --nocapture", - "test-e2e-legacy-local": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_epoch_monitor_with_2_foresters -- --nocapture", + "e2e-v1": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_e2e_v1 -- --nocapture", "test-address-v2": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_create_v2_address -- --nocapture", + "test-state-batched-indexer-async": "RUST_LOG=forester=debug,forester_utils=debug cargo test --package forester test_state_indexer_async_batched -- --nocapture", "docker:build": "docker build --tag forester -f Dockerfile .." }, "devDependencies": { diff --git a/forester/scripts/compare_performance.py b/forester/scripts/compare_performance.py new file mode 100755 index 0000000000..35034aff17 --- /dev/null +++ b/forester/scripts/compare_performance.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python3 +""" +Performance Comparison Script for Forester Logs +Compares queue processing performance between old and new forester versions. +""" + +import re +import sys +import argparse +from datetime import datetime +from collections import defaultdict +from typing import Dict, List, Tuple, Optional +import statistics + +class PerformanceAnalyzer: + def __init__(self): + # ANSI color removal + self.ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') + + # Patterns + self.timestamp_pattern = re.compile(r'(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+Z)') + self.queue_metric_pattern = re.compile(r'QUEUE_METRIC: (queue_empty|queue_has_elements) tree_type=(\S+) tree=(\S+)') + self.operation_start_pattern = re.compile(r'V2_TPS_METRIC: operation_start tree_type=(\w+)') + self.operation_complete_pattern = re.compile(r'V2_TPS_METRIC: operation_complete.*?duration_ms=(\d+).*?items_processed=(\d+)') + self.transaction_sent_pattern = re.compile(r'V2_TPS_METRIC: transaction_sent.*?tx_duration_ms=(\d+)') + + def clean_line(self, line: str) -> str: + return self.ansi_escape.sub('', line) + + def parse_timestamp(self, line: str) -> Optional[datetime]: + timestamp_match = self.timestamp_pattern.search(line) + if timestamp_match: + return datetime.fromisoformat(timestamp_match.group(1).replace('Z', '+00:00')) + return None + + def analyze_log(self, filename: str) -> Dict: + """Comprehensive analysis of a log file.""" + results = { + 'filename': filename, + 'queue_events': [], + 'operations': [], + 'transactions': [], + 'queue_emptying_times': [], + 'queue_response_times': [], + 'processing_rates': [], + 'transaction_durations': [] + } + + with open(filename, 'r') as f: + current_operation = None + + for line in f: + clean_line = self.clean_line(line) + timestamp = self.parse_timestamp(clean_line) + + if not timestamp: + continue + + # Parse queue metrics + if 'QUEUE_METRIC:' in clean_line: + queue_match = self.queue_metric_pattern.search(clean_line) + if queue_match: + state = queue_match.group(1) + results['queue_events'].append((timestamp, state)) + + # Parse operation start + elif 'operation_start' in clean_line: + start_match = self.operation_start_pattern.search(clean_line) + if start_match: + current_operation = { + 'start_time': timestamp, + 'tree_type': start_match.group(1) + } + + # Parse operation complete + elif 'operation_complete' in clean_line and current_operation: + complete_match = self.operation_complete_pattern.search(clean_line) + if complete_match: + duration_ms = int(complete_match.group(1)) + items_processed = int(complete_match.group(2)) + + operation = { + 'start_time': current_operation['start_time'], + 'end_time': timestamp, + 'duration_ms': duration_ms, + 'items_processed': items_processed, + 'tree_type': current_operation['tree_type'], + 'processing_rate': items_processed / (duration_ms / 1000) if duration_ms > 0 else 0 + } + + results['operations'].append(operation) + results['processing_rates'].append(operation['processing_rate']) + current_operation = None + + # Parse transaction sent + elif 'transaction_sent' in clean_line: + tx_match = self.transaction_sent_pattern.search(clean_line) + if tx_match: + tx_duration = int(tx_match.group(1)) + results['transactions'].append({ + 'timestamp': timestamp, + 'duration_ms': tx_duration + }) + results['transaction_durations'].append(tx_duration) + + # Calculate queue metrics + self._calculate_queue_metrics(results) + + return results + + def _calculate_queue_metrics(self, results: Dict): + """Calculate queue emptying and response times.""" + events = results['queue_events'] + + for i in range(1, len(events)): + prev_time, prev_state = events[i-1] + curr_time, curr_state = events[i] + + duration = (curr_time - prev_time).total_seconds() + + # Queue emptying time: has_elements -> empty + if prev_state == 'queue_has_elements' and curr_state == 'queue_empty': + results['queue_emptying_times'].append(duration) + + # Response time: empty -> has_elements (filter out immediate responses) + elif prev_state == 'queue_empty' and curr_state == 'queue_has_elements': + if duration > 0.01: # Filter immediate responses + results['queue_response_times'].append(duration) + + def generate_stats(self, data: List[float], name: str) -> Dict: + """Generate statistics for a dataset.""" + if not data: + return {'name': name, 'count': 0} + + return { + 'name': name, + 'count': len(data), + 'min': min(data), + 'max': max(data), + 'mean': statistics.mean(data), + 'median': statistics.median(data), + 'std_dev': statistics.stdev(data) if len(data) > 1 else 0 + } + + def print_stats(self, stats: Dict, unit: str = ""): + """Print statistics in a formatted way.""" + if stats['count'] == 0: + print(f" {stats['name']}: No data") + return + + print(f" {stats['name']}:") + print(f" Count: {stats['count']}") + print(f" Min: {stats['min']:.2f}{unit}") + print(f" Max: {stats['max']:.2f}{unit}") + print(f" Mean: {stats['mean']:.2f}{unit}") + print(f" Median: {stats['median']:.2f}{unit}") + if stats['count'] > 1: + print(f" Std Dev: {stats['std_dev']:.2f}{unit}") + + def compare_stats(self, old_stats: Dict, new_stats: Dict, unit: str = "") -> Dict: + """Compare two statistics and return improvement metrics.""" + if old_stats['count'] == 0 or new_stats['count'] == 0: + return {'valid': False} + + mean_improvement = ((old_stats['mean'] - new_stats['mean']) / old_stats['mean']) * 100 + median_improvement = ((old_stats['median'] - new_stats['median']) / old_stats['median']) * 100 + + return { + 'valid': True, + 'mean_improvement': mean_improvement, + 'median_improvement': median_improvement, + 'old_mean': old_stats['mean'], + 'new_mean': new_stats['mean'], + 'old_median': old_stats['median'], + 'new_median': new_stats['median'], + 'unit': unit + } + + def print_comparison(self, comparison: Dict, metric_name: str): + """Print comparison results.""" + if not comparison['valid']: + print(f" {metric_name}: Insufficient data for comparison") + return + + unit = comparison['unit'] + print(f" {metric_name}:") + print(f" Mean: {comparison['old_mean']:.2f}{unit} → {comparison['new_mean']:.2f}{unit} ({comparison['mean_improvement']:+.1f}%)") + print(f" Median: {comparison['old_median']:.2f}{unit} → {comparison['new_median']:.2f}{unit} ({comparison['median_improvement']:+.1f}%)") + + def analyze_and_compare(self, old_file: str, new_file: str): + """Main analysis and comparison function.""" + print("FORESTER PERFORMANCE COMPARISON") + print("=" * 60) + print() + + # Analyze both files + print("Analyzing log files...") + old_results = self.analyze_log(old_file) + new_results = self.analyze_log(new_file) + + print(f"Old version: {old_file}") + print(f"New version: {new_file}") + print() + + # Overall summary + print("OVERALL SUMMARY") + print("-" * 40) + print(f"Old version: {len(old_results['operations'])} operations, {len(old_results['transactions'])} transactions") + print(f"New version: {len(new_results['operations'])} operations, {len(new_results['transactions'])} transactions") + print() + + # Queue Performance Analysis + print("QUEUE PERFORMANCE ANALYSIS") + print("-" * 40) + + # Queue emptying times + old_emptying = self.generate_stats(old_results['queue_emptying_times'], "Queue Emptying Time") + new_emptying = self.generate_stats(new_results['queue_emptying_times'], "Queue Emptying Time") + + print("Old Version:") + self.print_stats(old_emptying, "s") + print() + print("New Version:") + self.print_stats(new_emptying, "s") + print() + + emptying_comparison = self.compare_stats(old_emptying, new_emptying, "s") + print("COMPARISON - Queue Emptying:") + self.print_comparison(emptying_comparison, "Queue Emptying Time") + print() + + # Response times + old_response = self.generate_stats(old_results['queue_response_times'], "Response Time") + new_response = self.generate_stats(new_results['queue_response_times'], "Response Time") + + response_comparison = self.compare_stats(old_response, new_response, "s") + print("COMPARISON - Response Time:") + self.print_comparison(response_comparison, "Response Time to New Work") + print() + + # Transaction Performance Analysis + print("TRANSACTION PERFORMANCE ANALYSIS") + print("-" * 40) + + old_tx = self.generate_stats(old_results['transaction_durations'], "Transaction Duration") + new_tx = self.generate_stats(new_results['transaction_durations'], "Transaction Duration") + + tx_comparison = self.compare_stats(old_tx, new_tx, "ms") + print("COMPARISON - Transaction Duration:") + self.print_comparison(tx_comparison, "Individual Transaction Time") + print() + + # Processing Rate Analysis + print("PROCESSING RATE ANALYSIS") + print("-" * 40) + + old_rate = self.generate_stats(old_results['processing_rates'], "Processing Rate") + new_rate = self.generate_stats(new_results['processing_rates'], "Processing Rate") + + rate_comparison = self.compare_stats(old_rate, new_rate, " items/sec") + print("COMPARISON - Processing Rate:") + self.print_comparison(rate_comparison, "Items Processing Rate") + print() + + # Key Insights + print("KEY INSIGHTS") + print("-" * 40) + + insights = [] + + if emptying_comparison['valid']: + if emptying_comparison['mean_improvement'] > 0: + insights.append(f"✅ Queue emptying is {emptying_comparison['mean_improvement']:.1f}% faster") + else: + insights.append(f"⚠️ Queue emptying is {abs(emptying_comparison['mean_improvement']):.1f}% slower") + + if response_comparison['valid']: + if response_comparison['mean_improvement'] > 0: + insights.append(f"✅ Response to new work is {response_comparison['mean_improvement']:.1f}% faster") + else: + insights.append(f"⚠️ Response to new work is {abs(response_comparison['mean_improvement']):.1f}% slower") + + if tx_comparison['valid']: + if tx_comparison['median_improvement'] > 0: + insights.append(f"✅ Individual transactions are {tx_comparison['median_improvement']:.1f}% faster") + else: + insights.append(f"⚠️ Individual transactions are {abs(tx_comparison['median_improvement']):.1f}% slower") + + if rate_comparison['valid']: + if rate_comparison['mean_improvement'] > 0: + insights.append(f"✅ Processing rate improved by {rate_comparison['mean_improvement']:.1f}%") + else: + insights.append(f"⚠️ Processing rate decreased by {abs(rate_comparison['mean_improvement']):.1f}%") + + for insight in insights: + print(f" {insight}") + + if not insights: + print(" No significant performance differences detected") + + print() + print("=" * 60) + +def main(): + parser = argparse.ArgumentParser(description='Compare forester performance between two log files') + parser.add_argument('old_log', help='Path to old version log file') + parser.add_argument('new_log', help='Path to new version log file') + + args = parser.parse_args() + + analyzer = PerformanceAnalyzer() + try: + analyzer.analyze_and_compare(args.old_log, args.new_log) + except FileNotFoundError as e: + print(f"Error: {e}") + sys.exit(1) + except Exception as e: + print(f"Error analyzing logs: {e}") + sys.exit(1) + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/forester/scripts/v2_stats.py b/forester/scripts/v2_stats.py new file mode 100755 index 0000000000..a0c8961b1a --- /dev/null +++ b/forester/scripts/v2_stats.py @@ -0,0 +1,272 @@ +#!/usr/bin/env python3 + +import re +import sys +from datetime import datetime, timedelta +from collections import defaultdict +from typing import Dict, List, Tuple, Optional +import argparse +import statistics +import json + +class V2TpsAnalyzer: + def __init__(self): + # ANSI color removal + self.ansi_escape = re.compile(r'\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])') + + self.v1_operation_start_pattern = re.compile( + r'V1_TPS_METRIC: operation_start tree_type=(\w+) tree=(\S+) epoch=(\d+)' + ) + self.v1_operation_complete_pattern = re.compile( + r'V1_TPS_METRIC: operation_complete tree_type=(\w+) tree=(\S+) epoch=(\d+) transactions=(\d+) duration_ms=(\d+) tps=([\d.]+)' + ) + self.v2_operation_start_pattern = re.compile( + r'V2_TPS_METRIC: operation_start tree_type=(\w+) (?:operation=(\w+) )?tree=(\S+) epoch=(\d+)' + ) + self.v2_operation_complete_pattern = re.compile( + r'V2_TPS_METRIC: operation_complete tree_type=(\w+) (?:operation=(\w+) )?tree=(\S+) epoch=(\d+) zkp_batches=(\d+) transactions=(\d+) instructions=(\d+) duration_ms=(\d+) tps=([\d.]+) ips=([\d.]+)(?:\s+items_processed=(\d+))?' + ) + self.v2_transaction_sent_pattern = re.compile( + r'V2_TPS_METRIC: transaction_sent tree_type=(\w+) (?:operation=(\w+) )?tree=(\S+) tx_num=(\d+) signature=(\S+) instructions=(\d+) tx_duration_ms=(\d+)' + ) + + # Data storage + self.operations: List[Dict] = [] + self.transactions: List[Dict] = [] + self.operation_summaries: List[Dict] = [] + + def clean_line(self, line: str) -> str: + """Remove ANSI color codes.""" + return self.ansi_escape.sub('', line) + + def parse_timestamp(self, line: str) -> Optional[datetime]: + """Extract timestamp from log line.""" + timestamp_match = re.search(r'(\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}\.\d+Z)', line) + if timestamp_match: + return datetime.fromisoformat(timestamp_match.group(1).replace('Z', '+00:00')) + return None + + def parse_log_line(self, line: str) -> None: + """Parse a single log line for V1/V2 TPS metrics.""" + clean_line = self.clean_line(line) + timestamp = self.parse_timestamp(clean_line) + + if not timestamp: + return + + # Parse V1 operation start + v1_start_match = self.v1_operation_start_pattern.search(clean_line) + if v1_start_match: + self.operations.append({ + 'type': 'start', + 'version': 'V1', + 'timestamp': timestamp, + 'tree_type': v1_start_match.group(1), + 'tree': v1_start_match.group(2), + 'epoch': int(v1_start_match.group(3)) + }) + return + + # Parse V1 operation complete + v1_complete_match = self.v1_operation_complete_pattern.search(clean_line) + if v1_complete_match: + self.operation_summaries.append({ + 'version': 'V1', + 'timestamp': timestamp, + 'tree_type': v1_complete_match.group(1), + 'tree': v1_complete_match.group(2), + 'epoch': int(v1_complete_match.group(3)), + 'transactions': int(v1_complete_match.group(4)), + 'duration_ms': int(v1_complete_match.group(5)), + 'tps': float(v1_complete_match.group(6)), + 'zkp_batches': 0, # V1 doesn't have zkp batches + 'instructions': int(v1_complete_match.group(4)), # For V1, instructions = transactions + 'ips': float(v1_complete_match.group(6)), # For V1, ips = tps + 'items_processed': 0 + }) + return + + # Parse V2 operation start + v2_start_match = self.v2_operation_start_pattern.search(clean_line) + if v2_start_match: + self.operations.append({ + 'type': 'start', + 'version': 'V2', + 'timestamp': timestamp, + 'tree_type': v2_start_match.group(1), + 'operation': v2_start_match.group(2) or 'batch', + 'tree': v2_start_match.group(3), + 'epoch': int(v2_start_match.group(4)) + }) + return + + # Parse V2 operation complete + v2_complete_match = self.v2_operation_complete_pattern.search(clean_line) + if v2_complete_match: + self.operation_summaries.append({ + 'version': 'V2', + 'timestamp': timestamp, + 'tree_type': v2_complete_match.group(1), + 'operation': v2_complete_match.group(2) or 'batch', + 'tree': v2_complete_match.group(3), + 'epoch': int(v2_complete_match.group(4)), + 'zkp_batches': int(v2_complete_match.group(5)), + 'transactions': int(v2_complete_match.group(6)), + 'instructions': int(v2_complete_match.group(7)), + 'duration_ms': int(v2_complete_match.group(8)), + 'tps': float(v2_complete_match.group(9)), + 'ips': float(v2_complete_match.group(10)), + 'items_processed': int(v2_complete_match.group(11)) if v2_complete_match.group(11) else 0 + }) + return + + # Parse V2 transaction sent + v2_tx_match = self.v2_transaction_sent_pattern.search(clean_line) + if v2_tx_match: + self.transactions.append({ + 'version': 'V2', + 'timestamp': timestamp, + 'tree_type': v2_tx_match.group(1), + 'operation': v2_tx_match.group(2) or 'batch', + 'tree': v2_tx_match.group(3), + 'tx_num': int(v2_tx_match.group(4)), + 'signature': v2_tx_match.group(5), + 'instructions': int(v2_tx_match.group(6)), + 'tx_duration_ms': int(v2_tx_match.group(7)) + }) + + def print_summary_stats(self) -> None: + """Print high-level summary statistics.""" + print("\n" + "="*80) + print("FORESTER PERFORMANCE ANALYSIS REPORT (V1 & V2)") + print("="*80) + + if not self.operation_summaries: + print("No TPS metrics found in logs") + return + + print(f"\nSUMMARY:") + print(f" Total operations analyzed: {len(self.operation_summaries)}") + + # Count total transactions from operation summaries + total_txs_from_ops = sum(op.get('transactions', 0) for op in self.operation_summaries) + print(f" Total transactions (from operations): {total_txs_from_ops}") + print(f" Total transaction events logged: {len(self.transactions)}") + + # Time span + if self.operation_summaries: + start_time = min(op['timestamp'] for op in self.operation_summaries) + end_time = max(op['timestamp'] for op in self.operation_summaries) + time_span = (end_time - start_time).total_seconds() + print(f" Analysis time span: {time_span:.1f}s ({time_span/60:.1f} minutes)") + + def print_tree_type_analysis(self) -> None: + """Analyze performance by tree type.""" + print("\n## PERFORMANCE BY TREE TYPE") + print("-" * 60) + print("\nNOTE: V1 and V2 use different transaction models:") + print(" V1: 1 tree update = 1 transaction (~1 slot/400ms latency)") + print(" V2: 10+ tree updates = 1 transaction (multi-slot batching + ZKP generation)") + print(" ") + print(" TPS comparison is misleading - V2 optimizes for cost efficiency, not transaction count.") + print(" Focus on 'Items Processed Per Second' and 'Total items processed' for V2.") + print(" V2's higher latency is architectural (batching) not a performance issue.") + print() + + tree_type_stats = defaultdict(lambda: { + 'operations': [], + 'total_transactions': 0, + 'total_instructions': 0, + 'total_zkp_batches': 0, + 'total_duration_ms': 0, + 'tps_values': [], + 'ips_values': [], + 'items_processed': 0 + }) + + for op in self.operation_summaries: + stats = tree_type_stats[op['tree_type']] + stats['operations'].append(op) + stats['total_transactions'] += op['transactions'] + stats['total_instructions'] += op['instructions'] + stats['total_zkp_batches'] += op['zkp_batches'] + stats['total_duration_ms'] += op['duration_ms'] + if op['tps'] > 0: + stats['tps_values'].append(op['tps']) + if op['ips'] > 0: + stats['ips_values'].append(op['ips']) + stats['items_processed'] += op['items_processed'] + + for tree_type, stats in sorted(tree_type_stats.items()): + print(f"\n{tree_type}:") + print(f" Operations: {len(stats['operations'])}") + print(f" Total transactions: {stats['total_transactions']}") + print(f" Total instructions: {stats['total_instructions']}") + print(f" Total ZKP batches: {stats['total_zkp_batches']}") + print(f" Total items processed: {stats['items_processed']}") + print(f" Total processing time: {stats['total_duration_ms']/1000:.2f}s") + + if stats['tps_values']: + print(f" TPS - Min: {min(stats['tps_values']):.2f}, Max: {max(stats['tps_values']):.2f}, Mean: {statistics.mean(stats['tps_values']):.2f}") + if stats['ips_values']: + print(f" IPS - Min: {min(stats['ips_values']):.2f}, Max: {max(stats['ips_values']):.2f}, Mean: {statistics.mean(stats['ips_values']):.2f}") + + # Calculate aggregate rates + if stats['total_duration_ms'] > 0: + aggregate_tps = stats['total_transactions'] / (stats['total_duration_ms'] / 1000) + aggregate_ips = stats['total_instructions'] / (stats['total_duration_ms'] / 1000) + print(f" Aggregate TPS: {aggregate_tps:.2f}") + print(f" Aggregate IPS: {aggregate_ips:.2f}") + + # For V2 trees, show Items Processed Per Second (more meaningful than TPS) + if 'V2' in tree_type and stats['items_processed'] > 0: + items_per_second = stats['items_processed'] / (stats['total_duration_ms'] / 1000) + print(f" *** Items Processed Per Second (IPPS): {items_per_second:.2f} ***") + print(f" ^ This is the meaningful throughput metric for V2 (actual tree updates/sec)") + + # Show batching efficiency + if stats['total_zkp_batches'] > 0: + avg_items_per_batch = stats['items_processed'] / stats['total_zkp_batches'] + print(f" Avg items per ZKP batch: {avg_items_per_batch:.1f}") + + def generate_report(self) -> None: + """Generate comprehensive TPS analysis report.""" + self.print_summary_stats() + self.print_tree_type_analysis() + print("\n" + "="*80) + +def main(): + parser = argparse.ArgumentParser(description='Analyze forester performance metrics (V1 & V2) - Focus on IPPS for V2') + parser.add_argument('logfile', nargs='?', default='-', help='Log file to analyze') + parser.add_argument('--tree-type', help='Filter to specific tree type') + + args = parser.parse_args() + + analyzer = V2TpsAnalyzer() + + # Read and parse log file + if args.logfile == '-': + log_file = sys.stdin + else: + log_file = open(args.logfile, 'r') + + try: + for line in log_file: + if 'TPS_METRIC' not in line: # Match both V1 and V2 + continue + + analyzer.parse_log_line(line) + finally: + if args.logfile != '-': + log_file.close() + + # Apply filters + if args.tree_type: + analyzer.operation_summaries = [op for op in analyzer.operation_summaries if op['tree_type'] == args.tree_type] + analyzer.transactions = [tx for tx in analyzer.transactions if tx['tree_type'] == args.tree_type] + + # Generate report + analyzer.generate_report() + +if __name__ == '__main__': + main() \ No newline at end of file diff --git a/forester/src/cli.rs b/forester/src/cli.rs index 6987504b94..1ed573038c 100644 --- a/forester/src/cli.rs +++ b/forester/src/cli.rs @@ -56,7 +56,7 @@ pub struct StartArgs { #[arg(long, env = "FORESTER_LEGACY_XS_PER_TX", default_value = "1")] pub legacy_ixs_per_tx: usize, - #[arg(long, env = "FORESTER_BATCH_IXS_PER_TX", default_value = "1")] + #[arg(long, env = "FORESTER_BATCH_IXS_PER_TX", default_value = "4")] pub batch_ixs_per_tx: usize, #[arg( @@ -88,7 +88,7 @@ pub struct StartArgs { #[arg(long, env = "FORESTER_ENABLE_PRIORITY_FEES", default_value = "false")] pub enable_priority_fees: bool, - #[arg(long, env = "FORESTER_RPC_POOL_SIZE", default_value = "50")] + #[arg(long, env = "FORESTER_RPC_POOL_SIZE", default_value = "100")] pub rpc_pool_size: u32, #[arg( diff --git a/forester/src/config.rs b/forester/src/config.rs index 595327d269..8317aeae15 100644 --- a/forester/src/config.rs +++ b/forester/src/config.rs @@ -153,7 +153,7 @@ impl Default for TransactionConfig { fn default() -> Self { Self { legacy_ixs_per_tx: 1, - batch_ixs_per_tx: 4, + batch_ixs_per_tx: 3, max_concurrent_batches: 20, cu_limit: 1_000_000, enable_priority_fees: false, diff --git a/forester/src/epoch_manager.rs b/forester/src/epoch_manager.rs index a25b30893b..829f489ab8 100644 --- a/forester/src/epoch_manager.rs +++ b/forester/src/epoch_manager.rs @@ -40,7 +40,6 @@ use crate::{ errors::{ ChannelError, ForesterError, InitializationError, RegistrationError, WorkReportError, }, - indexer_type::{rollover_address_merkle_tree, rollover_state_merkle_tree, IndexerType}, metrics::{push_metrics, queue_metric_update, update_forester_sol_balance}, pagerduty::send_pagerduty_alert, processor::{ @@ -53,7 +52,9 @@ use crate::{ v2::{process_batched_operations, BatchContext}, }, queue_helpers::QueueItemData, - rollover::is_tree_ready_for_rollover, + rollover::{ + is_tree_ready_for_rollover, rollover_address_merkle_tree, rollover_state_merkle_tree, + }, slot_tracker::{slot_duration, wait_until_slot_reached, SlotTracker}, tree_data_sync::fetch_trees, tree_finder::TreeFinder, @@ -89,7 +90,7 @@ pub enum MerkleProofType { } #[derive(Debug)] -pub struct EpochManager { +pub struct EpochManager { config: Arc, protocol_config: Arc, rpc_pool: Arc>, @@ -123,7 +124,7 @@ impl Clone for EpochManager { } } -impl + 'static> EpochManager { +impl EpochManager { #[allow(clippy::too_many_arguments)] pub async fn new( config: Arc, @@ -315,33 +316,66 @@ impl + 'static> EpochManager { "last_epoch: {:?}, current_epoch: {:?}, slot: {:?}", last_epoch, current_epoch, slot ); + if last_epoch.is_none_or(|last| current_epoch > last) { debug!("New epoch detected: {}", current_epoch); let phases = get_epoch_phases(&self.protocol_config, current_epoch); if slot < phases.registration.end { + debug!("Sending current epoch {} for processing", current_epoch); tx.send(current_epoch).await?; last_epoch = Some(current_epoch); } } let next_epoch = current_epoch + 1; - let next_phases = get_epoch_phases(&self.protocol_config, next_epoch); - let mut rpc = self.rpc_pool.get_connection().await?; - let slots_to_wait = next_phases.registration.start.saturating_sub(slot); - debug!( - "Waiting for epoch {} registration phase to start. Current slot: {}, Registration phase start slot: {}, Slots to wait: {}", - next_epoch, slot, next_phases.registration.start, slots_to_wait - ); + if last_epoch.is_none_or(|last| next_epoch > last) { + let next_phases = get_epoch_phases(&self.protocol_config, next_epoch); + + // If the next epoch's registration phase has started, send it immediately + if slot >= next_phases.registration.start && slot < next_phases.registration.end { + debug!( + "Next epoch {} registration phase already started, sending for processing", + next_epoch + ); + tx.send(next_epoch).await?; + last_epoch = Some(next_epoch); + continue; // Check for further epochs immediately + } - if let Err(e) = wait_until_slot_reached( - &mut *rpc, - &self.slot_tracker, - next_phases.registration.start, - ) - .await - { - error!("Error waiting for next registration phase: {:?}", e); - continue; + // Otherwise, wait for the next epoch's registration phase to start + let mut rpc = self.rpc_pool.get_connection().await?; + let slots_to_wait = next_phases.registration.start.saturating_sub(slot); + debug!( + "Waiting for epoch {} registration phase to start. Current slot: {}, Registration phase start slot: {}, Slots to wait: {}", + next_epoch, slot, next_phases.registration.start, slots_to_wait + ); + + if let Err(e) = wait_until_slot_reached( + &mut *rpc, + &self.slot_tracker, + next_phases.registration.start, + ) + .await + { + error!("Error waiting for next registration phase: {:?}", e); + continue; + } + + debug!( + "Next epoch {} registration phase started, sending for processing", + next_epoch + ); + if let Err(e) = tx.send(next_epoch).await { + error!( + "Failed to send next epoch {} for processing: {:?}", + next_epoch, e + ); + continue; + } + last_epoch = Some(next_epoch); + } else { + // we've already sent the next epoch, wait a bit before checking again + tokio::time::sleep(Duration::from_secs(10)).await; } } } @@ -480,11 +514,11 @@ impl + 'static> EpochManager { let rpc = LightClient::new(LightClientConfig { url: self.config.external_services.rpc_url.to_string(), photon_url: self.config.external_services.indexer_url.clone(), + api_key: self.config.external_services.photon_api_key.clone(), commitment_config: None, fetch_active_tree: false, }) - .await - .unwrap(); + .await?; let slot = rpc.get_slot().await?; let phases = get_epoch_phases(&self.protocol_config, epoch); @@ -509,7 +543,6 @@ impl + 'static> EpochManager { e ); if attempt < max_retries - 1 { - tokio::task::yield_now().await; sleep(retry_delay).await; } else { if let Err(alert_err) = send_pagerduty_alert( @@ -548,12 +581,12 @@ impl + 'static> EpochManager { info!("Registering for epoch: {}", epoch); let mut rpc = LightClient::new(LightClientConfig { url: self.config.external_services.rpc_url.to_string(), - photon_url: None, + photon_url: self.config.external_services.indexer_url.clone(), + api_key: self.config.external_services.photon_api_key.clone(), commitment_config: None, fetch_active_tree: false, }) - .await - .unwrap(); + .await?; let slot = rpc.get_slot().await?; let phases = get_epoch_phases(&self.protocol_config, epoch); @@ -699,10 +732,23 @@ impl + 'static> EpochManager { &self, epoch_info: &ForesterEpochInfo, ) -> Result { - info!("Waiting for active phase"); - let mut rpc = self.rpc_pool.get_connection().await?; let active_phase_start_slot = epoch_info.epoch.phases.active.start; + let current_slot = self.slot_tracker.estimated_current_slot(); + + if current_slot >= active_phase_start_slot { + info!( + "Active phase has already started. Current slot: {}. Active phase start slot: {}", + current_slot, active_phase_start_slot + ); + } else { + let waiting_slots = active_phase_start_slot - current_slot; + let waiting_secs = waiting_slots / 2; + info!("Waiting for active phase to start. Current slot: {}. Active phase start slot: {}. Waiting time: ~ {} seconds", + current_slot, + active_phase_start_slot, + waiting_secs); + } wait_until_slot_reached(&mut *rpc, &self.slot_tracker, active_phase_start_slot).await?; let forester_epoch_pda_pubkey = get_forester_epoch_pda_from_authority( @@ -902,7 +948,6 @@ impl + 'static> EpochManager { break 'outer_slot_loop; } - tokio::task::yield_now().await; current_slot = self.slot_tracker.estimated_current_slot(); } @@ -990,11 +1035,13 @@ impl + 'static> EpochManager { break 'inner_processing_loop; } }; + if items_processed_this_iteration > 0 { + debug!( + "Processed {} items in slot {:?}", + items_processed_this_iteration, forester_slot_details.slot + ); + } - debug!( - "Processed {} items in slot {:?}", - items_processed_this_iteration, forester_slot_details.slot - ); self.update_metrics_and_counts( epoch_info.epoch, items_processed_this_iteration, @@ -1004,6 +1051,9 @@ impl + 'static> EpochManager { push_metrics(&self.config.external_services.pushgateway_url).await?; estimated_slot = self.slot_tracker.estimated_current_slot(); + + // Add polling interval to reduce RPC pressure and improve response time + tokio::time::sleep(std::time::Duration::from_millis(25)).await; } Ok(()) } @@ -1055,6 +1105,10 @@ impl + 'static> EpochManager { ) -> Result { match tree_accounts.tree_type { TreeType::StateV1 | TreeType::AddressV1 => { + info!( + "Processing V1 tree: {} (type: {:?}, epoch: {})", + tree_accounts.merkle_tree, tree_accounts.tree_type, epoch_info.epoch + ); self.process_v1( epoch_info, epoch_pda, @@ -1213,12 +1267,12 @@ impl + 'static> EpochManager { info!("Reporting work"); let mut rpc = LightClient::new(LightClientConfig { url: self.config.external_services.rpc_url.to_string(), - photon_url: None, + photon_url: self.config.external_services.indexer_url.clone(), + api_key: self.config.external_services.photon_api_key.clone(), commitment_config: None, fetch_active_tree: false, }) - .await - .unwrap(); + .await?; let forester_epoch_pda_pubkey = get_forester_epoch_pda_from_authority( &self.config.derivation_pubkey, @@ -1305,7 +1359,6 @@ impl + 'static> EpochManager { rollover_address_merkle_tree( self.config.clone(), &mut *rpc, - self.indexer.clone(), tree_account, current_epoch, ) @@ -1315,7 +1368,6 @@ impl + 'static> EpochManager { rollover_state_merkle_tree( self.config.clone(), &mut *rpc, - self.indexer.clone(), tree_account, current_epoch, ) @@ -1363,7 +1415,7 @@ fn calculate_remaining_time_or_default( fields(forester = %config.payer_keypair.pubkey()) )] #[allow(clippy::too_many_arguments)] -pub async fn run_service + 'static>( +pub async fn run_service( config: Arc, protocol_config: Arc, rpc_pool: Arc>, @@ -1444,7 +1496,6 @@ pub async fn run_service + 'static>( retry_count += 1; if retry_count < config.retry_config.max_retries { debug!("Retrying in {:?}", retry_delay); - tokio::task::yield_now().await; sleep(retry_delay).await; retry_delay = std::cmp::min(retry_delay * 2, MAX_RETRY_DELAY); } else { diff --git a/forester/src/errors.rs b/forester/src/errors.rs index e1ca231da5..b0388859e9 100644 --- a/forester/src/errors.rs +++ b/forester/src/errors.rs @@ -9,8 +9,6 @@ use solana_program::{program_error::ProgramError, pubkey::Pubkey}; use thiserror::Error; use tracing::{info, warn}; -use crate::processor::v2::BatchProcessError; - #[derive(Error, Debug)] pub enum ForesterError { #[error("Element is not eligible for foresting")] @@ -31,9 +29,6 @@ pub enum ForesterError { #[error("Failed to register epoch {epoch}: {error}")] RegistrationFailed { epoch: u64, error: String }, - #[error("Batch processing error: {0}")] - BatchProcessing(#[from] BatchProcessError), - #[error("RPC error: {0}")] Rpc(#[from] RpcError), diff --git a/forester/src/forester_status.rs b/forester/src/forester_status.rs index 282929d65d..a1f94803ca 100644 --- a/forester/src/forester_status.rs +++ b/forester/src/forester_status.rs @@ -173,7 +173,8 @@ pub async fn fetch_forester_status(args: &StatusArgs) -> crate::Result<()> { debug!("RPC URL: {}", config.external_services.rpc_url); let mut rpc = LightClient::new(LightClientConfig { url: config.external_services.rpc_url.to_string(), - photon_url: None, + photon_url: config.external_services.indexer_url.clone(), + api_key: config.external_services.photon_api_key.clone(), commitment_config: None, fetch_active_tree: false, }) diff --git a/forester/src/indexer_type.rs b/forester/src/indexer_type.rs deleted file mode 100644 index f4d69b7850..0000000000 --- a/forester/src/indexer_type.rs +++ /dev/null @@ -1,386 +0,0 @@ -use std::{marker::PhantomData, sync::Arc}; - -use async_trait::async_trait; -use forester_utils::forester_epoch::TreeAccounts; -use light_batched_merkle_tree::{ - merkle_tree::BatchedMerkleTreeAccount, queue::BatchedQueueAccount, -}; -use light_client::{ - indexer::{photon_indexer::PhotonIndexer, Indexer, StateMerkleTreeAccounts}, - rpc::Rpc, -}; -use light_compressed_account::TreeType; -use light_hasher::{Hasher, Poseidon}; -use light_merkle_tree_reference::MerkleTree; -use light_program_test::indexer::{ - state_tree::StateMerkleTreeBundle, TestIndexer, TestIndexerExtensions, -}; -use light_sdk::constants::{STATE_MERKLE_TREE_CANOPY_DEPTH, STATE_MERKLE_TREE_HEIGHT}; -use solana_program::pubkey::Pubkey; -use solana_sdk::{signature::Keypair, signer::Signer}; -use tokio::sync::Mutex; -use tracing::info; - -use crate::{ - errors::ForesterError, - rollover::{perform_address_merkle_tree_rollover, perform_state_merkle_tree_rollover_forester}, - ForesterConfig, -}; - -mod sealed { - use super::*; - pub trait Sealed {} - impl Sealed for TestIndexer {} - impl Sealed for PhotonIndexer {} -} - -#[async_trait] -pub trait IndexerType: Indexer + sealed::Sealed { - fn rpc_phantom(&self) -> PhantomData { - PhantomData - } - fn handle_state_bundle( - &mut self, - new_merkle_tree: Pubkey, - new_queue: Pubkey, - new_cpi_context: Pubkey, - ); - - fn handle_address_bundle(&mut self, new_merkle_tree: &Keypair, new_queue: &Keypair); - - async fn finalize_batch_address_tree_update( - &mut self, - rpc: &mut R, - new_merkle_tree_pubkey: Pubkey, - ); - - async fn update_test_indexer_after_nullification( - &mut self, - rpc: &mut R, - merkle_tree_pubkey: Pubkey, - batch_index: usize, - ); - - async fn update_test_indexer_after_append( - &mut self, - rpc: &mut R, - merkle_tree_pubkey: Pubkey, - output_queue: Pubkey, - ); -} - -#[async_trait] -impl IndexerType for TestIndexer { - fn handle_state_bundle( - &mut self, - new_merkle_tree: Pubkey, - new_queue: Pubkey, - new_cpi_context: Pubkey, - ) { - let state_bundle = StateMerkleTreeBundle { - rollover_fee: 0, - accounts: StateMerkleTreeAccounts { - merkle_tree: new_merkle_tree, - nullifier_queue: new_queue, - cpi_context: new_cpi_context, - }, - tree_type: TreeType::StateV1, - output_queue_elements: vec![], - merkle_tree: Box::new(MerkleTree::::new( - STATE_MERKLE_TREE_HEIGHT, - STATE_MERKLE_TREE_CANOPY_DEPTH, - )), - input_leaf_indices: vec![], - num_inserted_batches: 0, - output_queue_batch_size: None, - }; - self.add_state_bundle(state_bundle); - } - - fn handle_address_bundle(&mut self, new_merkle_tree: &Keypair, new_queue: &Keypair) { - self.add_address_merkle_tree_accounts(new_merkle_tree, new_queue, None); - } - - async fn finalize_batch_address_tree_update( - &mut self, - rpc: &mut R, - merkle_tree_pubkey: Pubkey, - ) { - let mut account = rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap(); - self.finalize_batched_address_tree_update(merkle_tree_pubkey, account.data.as_mut_slice()) - .await; - } - - async fn update_test_indexer_after_nullification( - &mut self, - rpc: &mut R, - merkle_tree_pubkey: Pubkey, - batch_index: usize, - ) { - let state_merkle_tree_bundle = self - .state_merkle_trees - .iter_mut() - .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey) - .unwrap(); - - let mut merkle_tree_account = rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap(); - let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( - merkle_tree_account.data.as_mut_slice(), - &merkle_tree_pubkey.into(), - ) - .unwrap(); - - let batch = &merkle_tree.queue_batches.batches[batch_index]; - let batch_size = batch.zkp_batch_size; - let leaf_indices_tx_hashes = - state_merkle_tree_bundle.input_leaf_indices[..batch_size as usize].to_vec(); - for leaf_info in leaf_indices_tx_hashes.iter() { - let index = leaf_info.leaf_index as usize; - let leaf = leaf_info.leaf; - let mut index_32_bytes = [0u8; 32]; - index_32_bytes[24..].copy_from_slice(index.to_be_bytes().as_slice()); - - let nullifier = Poseidon::hashv(&[&leaf, &index_32_bytes, &leaf_info.tx_hash]).unwrap(); - - state_merkle_tree_bundle.input_leaf_indices.remove(0); - let result = state_merkle_tree_bundle - .merkle_tree - .update(&nullifier, index); - if result.is_err() { - let num_missing_leaves = - (index + 1) - state_merkle_tree_bundle.merkle_tree.rightmost_index; - state_merkle_tree_bundle - .merkle_tree - .append_batch(&vec![&[0u8; 32]; num_missing_leaves]) - .unwrap(); - state_merkle_tree_bundle - .merkle_tree - .update(&nullifier, index) - .unwrap(); - } - } - } - - async fn update_test_indexer_after_append( - &mut self, - rpc: &mut R, - merkle_tree_pubkey: Pubkey, - output_queue_pubkey: Pubkey, - ) { - let state_merkle_tree_bundle = self - .state_merkle_trees - .iter_mut() - .find(|x| x.accounts.merkle_tree == merkle_tree_pubkey) - .unwrap(); - - let (merkle_tree_next_index, root) = { - let mut merkle_tree_account = - rpc.get_account(merkle_tree_pubkey).await.unwrap().unwrap(); - let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( - merkle_tree_account.data.as_mut_slice(), - &merkle_tree_pubkey.into(), - ) - .unwrap(); - ( - merkle_tree.next_index as usize, - *merkle_tree.root_history.last().unwrap(), - ) - }; - - let zkp_batch_size = { - let mut output_queue_account = - rpc.get_account(output_queue_pubkey).await.unwrap().unwrap(); - let output_queue = - BatchedQueueAccount::output_from_bytes(output_queue_account.data.as_mut_slice()) - .unwrap(); - - output_queue.batch_metadata.zkp_batch_size - }; - - let leaves = state_merkle_tree_bundle.output_queue_elements.to_vec(); - let batch_update_leaves = leaves[0..zkp_batch_size as usize].to_vec(); - - for (i, (new_leaf, _)) in batch_update_leaves.iter().enumerate() { - let index = merkle_tree_next_index + i - zkp_batch_size as usize; - // This is dangerous it should call self.get_leaf_by_index() but it - // can t for mutable borrow - // TODO: call a get_leaf_by_index equivalent, we could move the method to the reference merkle tree - let leaf = state_merkle_tree_bundle - .merkle_tree - .get_leaf(index) - .unwrap_or_default(); - if leaf == [0u8; 32] { - let result = state_merkle_tree_bundle.merkle_tree.update(new_leaf, index); - if result.is_err() && state_merkle_tree_bundle.merkle_tree.rightmost_index == index - { - state_merkle_tree_bundle - .merkle_tree - .append(new_leaf) - .unwrap(); - } else { - result.unwrap(); - } - } - } - assert_eq!( - root, - state_merkle_tree_bundle.merkle_tree.root(), - "update indexer after append root invalid" - ); - - for _ in 0..zkp_batch_size { - state_merkle_tree_bundle.output_queue_elements.remove(0); - } - } -} - -// Implementation for PhotonIndexer - no-op -#[async_trait] -impl IndexerType for PhotonIndexer { - fn handle_state_bundle( - &mut self, - _new_merkle_tree: Pubkey, - _new_queue: Pubkey, - _new_cpi_context: Pubkey, - ) { - // No-op for production indexer - } - - fn handle_address_bundle(&mut self, _new_merkle_tree: &Keypair, _new_queue: &Keypair) { - // No-op for production indexer - } - - async fn finalize_batch_address_tree_update( - &mut self, - _rpc: &mut R, - _new_merkle_tree_pubkey: Pubkey, - ) { - // No-op for production indexer - } - - async fn update_test_indexer_after_nullification( - &mut self, - _rpc: &mut R, - _merkle_tree_pubkey: Pubkey, - _batch_index: usize, - ) { - // No-op for production indexer - } - - async fn update_test_indexer_after_append( - &mut self, - _rpc: &mut R, - _merkle_tree_pubkey: Pubkey, - _output_queue: Pubkey, - ) { - // No-op for production indexer - } -} - -pub async fn rollover_state_merkle_tree>( - config: Arc, - rpc: &mut R, - indexer: Arc>, - tree_accounts: &TreeAccounts, - epoch: u64, -) -> Result<(), ForesterError> { - let new_nullifier_queue_keypair = Keypair::new(); - let new_merkle_tree_keypair = Keypair::new(); - let new_cpi_signature_keypair = Keypair::new(); - - let rollover_signature = perform_state_merkle_tree_rollover_forester( - &config.payer_keypair, - &config.derivation_pubkey, - rpc, - &new_nullifier_queue_keypair, - &new_merkle_tree_keypair, - &new_cpi_signature_keypair, - &tree_accounts.merkle_tree, - &tree_accounts.queue, - &Pubkey::default(), - epoch, - ) - .await?; - - info!("State rollover signature: {:?}", rollover_signature); - - let mut indexer_lock = indexer.lock().await; - indexer_lock.handle_state_bundle( - new_merkle_tree_keypair.pubkey(), - new_nullifier_queue_keypair.pubkey(), - new_cpi_signature_keypair.pubkey(), - ); - - Ok(()) -} - -pub async fn rollover_address_merkle_tree>( - config: Arc, - rpc: &mut R, - indexer: Arc>, - tree_accounts: &TreeAccounts, - epoch: u64, -) -> Result<(), ForesterError> { - let new_nullifier_queue_keypair = Keypair::new(); - let new_merkle_tree_keypair = Keypair::new(); - - let rollover_signature = perform_address_merkle_tree_rollover( - &config.payer_keypair, - &config.derivation_pubkey, - rpc, - &new_nullifier_queue_keypair, - &new_merkle_tree_keypair, - &tree_accounts.merkle_tree, - &tree_accounts.queue, - epoch, - ) - .await?; - - info!("Address rollover signature: {:?}", rollover_signature); - - let mut indexer_lock = indexer.lock().await; - indexer_lock.handle_address_bundle(&new_merkle_tree_keypair, &new_nullifier_queue_keypair); - - Ok(()) -} - -pub async fn finalize_batch_address_tree_update>( - rpc: &mut R, - indexer: Arc>, - new_merkle_tree_pubkey: Pubkey, -) -> Result<(), ForesterError> { - let mut indexer_lock = indexer.lock().await; - indexer_lock - .finalize_batch_address_tree_update(rpc, new_merkle_tree_pubkey) - .await; - - Ok(()) -} - -pub async fn update_test_indexer_after_nullification>( - rpc: &mut R, - indexer: Arc>, - merkle_tree_pubkey: Pubkey, - batch_index: usize, -) -> Result<(), ForesterError> { - let mut indexer_lock = indexer.lock().await; - indexer_lock - .update_test_indexer_after_nullification(rpc, merkle_tree_pubkey, batch_index) - .await; - - Ok(()) -} - -pub async fn update_test_indexer_after_append>( - rpc: &mut R, - indexer: Arc>, - merkle_tree_pubkey: Pubkey, - output_queue: Pubkey, -) -> Result<(), ForesterError> { - let mut indexer_lock = indexer.lock().await; - indexer_lock - .update_test_indexer_after_append(rpc, merkle_tree_pubkey, output_queue) - .await; - - Ok(()) -} diff --git a/forester/src/lib.rs b/forester/src/lib.rs index c77168ee43..272cecfab2 100644 --- a/forester/src/lib.rs +++ b/forester/src/lib.rs @@ -6,7 +6,6 @@ pub mod epoch_manager; pub mod errors; pub mod forester_status; pub mod helius_priority_fee_types; -mod indexer_type; pub mod metrics; pub mod pagerduty; pub mod processor; @@ -38,7 +37,6 @@ use tracing::debug; use crate::{ epoch_manager::{run_service, WorkReport}, - indexer_type::IndexerType, metrics::QUEUE_LENGTH, processor::tx_cache::ProcessedHashCache, queue_helpers::{ @@ -55,7 +53,8 @@ pub async fn run_queue_info( ) -> Result<()> { let mut rpc = LightClient::new(LightClientConfig { url: config.external_services.rpc_url.to_string(), - photon_url: None, + photon_url: config.external_services.indexer_url.clone(), + api_key: config.external_services.photon_api_key.clone(), commitment_config: None, fetch_active_tree: false, }) @@ -111,7 +110,7 @@ pub async fn run_queue_info( Ok(()) } -pub async fn run_pipeline + 'static>( +pub async fn run_pipeline( config: Arc, rpc_rate_limiter: Option, send_tx_rate_limiter: Option, diff --git a/forester/src/processor/v1/helpers.rs b/forester/src/processor/v1/helpers.rs index 66bffdc0f6..a079649263 100644 --- a/forester/src/processor/v1/helpers.rs +++ b/forester/src/processor/v1/helpers.rs @@ -90,9 +90,9 @@ pub async fn fetch_proofs_and_create_instructions( }; let indexer_guard = indexer.lock().await; - let mut rpc = pool.get_connection().await?; + let rpc = pool.get_connection().await?; - if let Err(e) = wait_for_indexer(&mut *rpc, &*indexer_guard).await { + if let Err(e) = wait_for_indexer(&*rpc, &*indexer_guard).await { warn!("Indexer not fully caught up, but proceeding anyway: {}", e); } diff --git a/forester/src/processor/v1/send_transaction.rs b/forester/src/processor/v1/send_transaction.rs index 31173c23b3..bed289d84d 100644 --- a/forester/src/processor/v1/send_transaction.rs +++ b/forester/src/processor/v1/send_transaction.rs @@ -22,7 +22,7 @@ use solana_sdk::{ transaction::Transaction, }; use tokio::time::Instant; -use tracing::{debug, error, trace, warn}; +use tracing::{debug, error, info, trace, warn}; use crate::{ epoch_manager::WorkItem, @@ -63,6 +63,13 @@ pub async fn send_batched_transactions( ) -> Result { let function_start_time = Instant::now(); + info!( + "V1_TPS_METRIC: operation_start tree_type={} tree={} epoch={}", + tree_accounts.tree_type, + tree_accounts.merkle_tree, + transaction_builder.epoch() + ); + let num_sent_transactions = Arc::new(AtomicUsize::new(0)); let operation_cancel_signal = Arc::new(AtomicBool::new(false)); @@ -152,6 +159,16 @@ pub async fn send_batched_transactions( let total_sent_successfully = num_sent_transactions.load(Ordering::SeqCst); trace!(tree = %tree_accounts.merkle_tree, "Transaction sending loop finished. Total transactions sent successfully: {}", total_sent_successfully); + + let total_duration = function_start_time.elapsed(); + let tps = if total_duration.as_secs_f64() > 0.0 { + total_sent_successfully as f64 / total_duration.as_secs_f64() + } else { + 0.0 + }; + + info!("V1_TPS_METRIC: operation_complete tree_type={} tree={} epoch={} transactions={} duration_ms={} tps={:.2}", tree_accounts.tree_type, tree_accounts.merkle_tree, transaction_builder.epoch(), total_sent_successfully, total_duration.as_millis(), tps); + Ok(total_sent_successfully) } @@ -210,10 +227,21 @@ async fn prepare_batch_prerequisites( }; if queue_item_data.is_empty() { + info!( + "QUEUE_METRIC: queue_empty tree_type={} tree={}", + tree_accounts.tree_type, tree_accounts.merkle_tree + ); trace!(tree = %tree_id_str, "Queue is empty, no transactions to send."); return Ok(None); // Return None to indicate no work } + info!( + "QUEUE_METRIC: queue_has_elements tree_type={} tree={} count={}", + tree_accounts.tree_type, + tree_accounts.merkle_tree, + queue_item_data.len() + ); + let (recent_blockhash, last_valid_block_height) = { let mut rpc = pool.get_connection().await.map_err(|e| { error!(tree = %tree_id_str, "Failed to get RPC for blockhash: {:?}", e); diff --git a/forester/src/processor/v2/address.rs b/forester/src/processor/v2/address.rs index 6d87083221..9343913159 100644 --- a/forester/src/processor/v2/address.rs +++ b/forester/src/processor/v2/address.rs @@ -1,132 +1,73 @@ +use anyhow::Error; use borsh::BorshSerialize; -use forester_utils::instructions::address_batch_update::create_batch_update_address_tree_instruction_data; +use forester_utils::instructions::address_batch_update::{ + get_address_update_instruction_stream, AddressUpdateConfig, +}; +use futures::stream::{Stream, StreamExt}; +use light_batched_merkle_tree::merkle_tree::InstructionDataAddressAppendInputs; use light_client::{indexer::Indexer, rpc::Rpc}; -use light_merkle_tree_metadata::events::MerkleTreeEvent; use light_registry::account_compression_cpi::sdk::create_batch_update_address_tree_instruction; +use solana_program::instruction::Instruction; use solana_sdk::signer::Signer; -use tracing::{debug, info, instrument, log::error, trace}; - -use super::{ - common::BatchContext, - error::{BatchProcessError, Result}, -}; -use crate::indexer_type::{finalize_batch_address_tree_update, IndexerType}; +use tracing::{info, instrument}; + +use super::common::{process_stream, BatchContext, ParsedMerkleTreeData}; +use crate::Result; + +async fn create_stream_future( + ctx: &BatchContext, + merkle_tree_data: ParsedMerkleTreeData, +) -> Result<( + impl Stream>> + Send, + u16, +)> +where + R: Rpc, + I: Indexer + 'static, +{ + let config = AddressUpdateConfig { + rpc_pool: ctx.rpc_pool.clone(), + indexer: ctx.indexer.clone(), + merkle_tree_pubkey: ctx.merkle_tree, + prover_url: ctx.prover_url.clone(), + polling_interval: ctx.prover_polling_interval, + max_wait_time: ctx.prover_max_wait_time, + ixs_per_tx: ctx.ixs_per_tx, + }; + let (stream, size) = get_address_update_instruction_stream(config, merkle_tree_data) + .await + .map_err(Error::from)?; + let stream = stream.map(|item| item.map_err(Error::from)); + Ok((stream, size)) +} #[instrument(level = "debug", skip(context), fields(tree = %context.merkle_tree))] -pub(crate) async fn process_batch>( +pub(crate) async fn process_batch( context: &BatchContext, + merkle_tree_data: ParsedMerkleTreeData, ) -> Result { - trace!("Processing address batch operation"); - - let mut rpc = context.rpc_pool.get_connection().await?; - - let (instruction_data_vec, zkp_batch_size) = create_batch_update_address_tree_instruction_data( - &mut *rpc, - &mut *context.indexer.lock().await, - &context.merkle_tree, - context.prover_url.clone(), - context.prover_polling_interval, - context.prover_max_wait_time, - ) - .await - .map_err(|e| { - error!( - "Failed to create batch update address tree instruction data: {}", - e - ); - BatchProcessError::InstructionData(e.to_string()) - })?; - - if instruction_data_vec.is_empty() { - trace!("No ZKP batches to process for address tree"); - let mut cache = context.ops_cache.lock().await; - cache.cleanup(); - return Ok(0); - } - - info!( - "Processing {} ZKP batch updates for address tree", - instruction_data_vec.len() - ); - - let mut batches_processed = 0; - - for (chunk_idx, instruction_chunk) in - instruction_data_vec.chunks(context.ixs_per_tx).enumerate() - { - debug!( - "Sending address update transaction chunk {}/{} for tree: {}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - context.merkle_tree - ); - - let mut instructions = Vec::with_capacity(context.ixs_per_tx); - for instruction_data in instruction_chunk { - debug!( - "Instruction data size: {} bytes", - instruction_data.try_to_vec().map(|v| v.len()).unwrap_or(0) - ); - - instructions.push(create_batch_update_address_tree_instruction( - context.authority.pubkey(), - context.derivation, - context.merkle_tree, - context.epoch, - instruction_data.try_to_vec().map_err(|e| { - BatchProcessError::InstructionData(format!( - "Failed to serialize instruction data: {}", - e - )) - })?, - )); - - batches_processed += 1; - } - - let tx = match rpc - .create_and_send_transaction_with_event::( - &instructions, - &context.authority.pubkey(), - &[&context.authority], - ) - .await - { - Ok(tx) => { - info!( - "Address update transaction chunk {}/{} sent successfully: {:?}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - tx - ); - tx - } - Err(e) => { - error!( - "Failed to send address update transaction chunk {}/{} for tree {}: {:?}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - context.merkle_tree, - e - ); - return Err(e.into()); - } - }; - - debug!("Address batch transaction: {:?}", tx); - - finalize_batch_address_tree_update(&mut *rpc, context.indexer.clone(), context.merkle_tree) - .await - .map_err(|e| { - error!("Failed to finalize batch address tree update: {:?}", e); - BatchProcessError::Indexer(e.to_string()) - })?; - } - info!( - "Address batch processing completed successfully. Processed {} batches", - batches_processed + "V2_TPS_METRIC: operation_start tree_type=AddressV2 tree={} epoch={}", + context.merkle_tree, context.epoch ); - - Ok(batches_processed * zkp_batch_size as usize) + let instruction_builder = |data: &InstructionDataAddressAppendInputs| -> Instruction { + let serialized_data = data.try_to_vec().unwrap(); + create_batch_update_address_tree_instruction( + context.authority.pubkey(), + context.derivation, + context.merkle_tree, + context.epoch, + serialized_data, + ) + }; + + let stream_future = create_stream_future(context, merkle_tree_data); + process_stream( + context, + stream_future, + instruction_builder, + "AddressV2", + None, + ) + .await } diff --git a/forester/src/processor/v2/common.rs b/forester/src/processor/v2/common.rs index 0b46d80b33..57fe5f12f2 100644 --- a/forester/src/processor/v2/common.rs +++ b/forester/src/processor/v2/common.rs @@ -1,20 +1,36 @@ -use std::{sync::Arc, time::Duration}; +use std::{future::Future, sync::Arc, time::Duration}; +use borsh::BorshSerialize; use forester_utils::rpc_pool::SolanaRpcPool; +pub use forester_utils::{ParsedMerkleTreeData, ParsedQueueData}; +use futures::{pin_mut, stream::StreamExt, Stream}; use light_batched_merkle_tree::{ - batch::{Batch, BatchState}, - merkle_tree::BatchedMerkleTreeAccount, - queue::BatchedQueueAccount, + batch::BatchState, merkle_tree::BatchedMerkleTreeAccount, queue::BatchedQueueAccount, }; -use light_client::{indexer::Indexer, rpc::Rpc}; +use light_client::rpc::Rpc; use light_compressed_account::TreeType; -use solana_program::pubkey::Pubkey; -use solana_sdk::signature::Keypair; +use light_program_test::Indexer; +use solana_sdk::{instruction::Instruction, pubkey::Pubkey, signature::Keypair, signer::Signer}; use tokio::sync::Mutex; -use tracing::{debug, error, trace}; +use tracing::{debug, error, info, trace}; -use super::{address, error::Result, state, BatchProcessError}; -use crate::{indexer_type::IndexerType, processor::tx_cache::ProcessedHashCache}; +use super::{address, state}; +use crate::{errors::ForesterError, processor::tx_cache::ProcessedHashCache, Result}; + +#[derive(Debug)] +pub enum BatchReadyState { + NotReady, + AddressReadyForAppend { + merkle_tree_data: ParsedMerkleTreeData, + }, + StateReadyForAppend { + merkle_tree_data: ParsedMerkleTreeData, + output_queue_data: ParsedQueueData, + }, + StateReadyForNullify { + merkle_tree_data: ParsedMerkleTreeData, + }, +} #[derive(Debug)] pub struct BatchContext { @@ -33,19 +49,120 @@ pub struct BatchContext { } #[derive(Debug)] -pub enum BatchReadyState { - NotReady, - ReadyForAppend, - ReadyForNullify, -} - -#[derive(Debug)] -pub struct BatchProcessor> { +pub struct BatchProcessor { context: BatchContext, tree_type: TreeType, } -impl> BatchProcessor { +/// Processes a stream of batched instruction data into transactions. +pub(crate) async fn process_stream( + context: &BatchContext, + stream_creator_future: FutC, + instruction_builder: impl Fn(&D) -> Instruction, + tree_type_str: &str, + operation: Option<&str>, +) -> Result +where + R: Rpc, + I: Indexer, + S: Stream>> + Send, + D: BorshSerialize, + FutC: Future> + Send, +{ + let start_time = std::time::Instant::now(); + trace!("Executing batched stream processor (hybrid)"); + + let (batch_stream, zkp_batch_size) = stream_creator_future.await?; + + if zkp_batch_size == 0 { + trace!("ZKP batch size is 0, no work to do."); + return Ok(0); + } + + pin_mut!(batch_stream); + let mut total_instructions_processed = 0; + let mut transactions_sent = 0; + + while let Some(batch_result) = batch_stream.next().await { + let instruction_batch = batch_result?; + + if instruction_batch.is_empty() { + continue; + } + + let instructions: Vec = + instruction_batch.iter().map(&instruction_builder).collect(); + + let tx_start = std::time::Instant::now(); + let signature = send_transaction_batch(context, instructions).await?; + transactions_sent += 1; + total_instructions_processed += instruction_batch.len(); + let tx_duration = tx_start.elapsed(); + + let operation_suffix = operation + .map(|op| format!(" operation={}", op)) + .unwrap_or_default(); + info!( + "V2_TPS_METRIC: transaction_sent tree_type={}{} tree={} tx_num={} signature={} instructions={} tx_duration_ms={} (hybrid)", + tree_type_str, operation_suffix, context.merkle_tree, transactions_sent, signature, instruction_batch.len(), tx_duration.as_millis() + ); + } + + if total_instructions_processed == 0 { + trace!("No instructions were processed from the stream."); + return Ok(0); + } + + let total_duration = start_time.elapsed(); + let total_items_processed = total_instructions_processed * zkp_batch_size as usize; + let tps = if total_duration.as_secs_f64() > 0.0 { + transactions_sent as f64 / total_duration.as_secs_f64() + } else { + 0.0 + }; + let ips = if total_duration.as_secs_f64() > 0.0 { + total_instructions_processed as f64 / total_duration.as_secs_f64() + } else { + 0.0 + }; + + let operation_suffix = operation + .map(|op| format!(" operation={}", op)) + .unwrap_or_default(); + info!( + "V2_TPS_METRIC: operation_complete tree_type={}{} tree={} epoch={} zkp_batches={} transactions={} instructions={} duration_ms={} tps={:.2} ips={:.2} items_processed={} (hybrid)", + tree_type_str, operation_suffix, context.merkle_tree, context.epoch, total_instructions_processed, transactions_sent, total_instructions_processed, + total_duration.as_millis(), tps, ips, total_items_processed + ); + + info!( + "Batched stream processing complete. Processed {} total items.", + total_items_processed + ); + + Ok(total_items_processed) +} + +pub(crate) async fn send_transaction_batch( + context: &BatchContext, + instructions: Vec, +) -> Result { + info!( + "Sending transaction with {} instructions...", + instructions.len() + ); + let mut rpc = context.rpc_pool.get_connection().await?; + let signature = rpc + .create_and_send_transaction( + &instructions, + &context.authority.pubkey(), + &[&context.authority], + ) + .await?; + Ok(signature.to_string()) +} + +impl BatchProcessor { pub fn new(context: BatchContext, tree_type: TreeType) -> Self { Self { context, tree_type } } @@ -58,66 +175,65 @@ impl> BatchProcessor { let state = self.verify_batch_ready().await; match state { - BatchReadyState::ReadyForAppend => match self.tree_type { - TreeType::AddressV2 => { - trace!( - "Processing address append for tree: {}", - self.context.merkle_tree - ); + BatchReadyState::AddressReadyForAppend { merkle_tree_data } => { + trace!( + "Processing address append for tree: {}", + self.context.merkle_tree + ); - let batch_hash = format!( - "address_batch_{}_{}", - self.context.merkle_tree, self.context.epoch - ); - { - let mut cache = self.context.ops_cache.lock().await; - if cache.contains(&batch_hash) { - debug!("Skipping already processed address batch: {}", batch_hash); - return Ok(0); - } - cache.add(&batch_hash); + let batch_hash = format!( + "address_batch_{}_{}", + self.context.merkle_tree, self.context.epoch + ); + { + let mut cache = self.context.ops_cache.lock().await; + if cache.contains(&batch_hash) { + debug!("Skipping already processed address batch: {}", batch_hash); + return Ok(0); } + cache.add(&batch_hash); + } - let result = address::process_batch(&self.context).await; + let result = address::process_batch(&self.context, merkle_tree_data).await; - if let Err(ref e) = result { - error!( - "Address append failed for tree {}: {:?}", - self.context.merkle_tree, e - ); - } + if let Err(ref e) = result { + error!( + "Address append failed for tree {}: {:?}", + self.context.merkle_tree, e + ); + } - let mut cache = self.context.ops_cache.lock().await; - cache.cleanup_by_key(&batch_hash); - trace!("Cache cleaned up for batch: {}", batch_hash); + let mut cache = self.context.ops_cache.lock().await; + cache.cleanup_by_key(&batch_hash); + trace!("Cache cleaned up for batch: {}", batch_hash); - result - } - TreeType::StateV2 => { - trace!( - "Process state append for tree: {}", - self.context.merkle_tree + result + } + BatchReadyState::StateReadyForAppend { + merkle_tree_data, + output_queue_data, + } => { + trace!( + "Process state append for tree: {}", + self.context.merkle_tree + ); + let result = self + .process_state_append_hybrid(merkle_tree_data, output_queue_data) + .await; + if let Err(ref e) = result { + error!( + "State append failed for tree {}: {:?}", + self.context.merkle_tree, e ); - let result = self.process_state_append().await; - if let Err(ref e) = result { - error!( - "State append failed for tree {}: {:?}", - self.context.merkle_tree, e - ); - } - result - } - _ => { - error!("Unsupported tree type for append: {:?}", self.tree_type); - Err(BatchProcessError::UnsupportedTreeType(self.tree_type)) } - }, - BatchReadyState::ReadyForNullify => { + result + } + BatchReadyState::StateReadyForNullify { merkle_tree_data } => { trace!( "Processing batch for nullify, tree: {}", self.context.merkle_tree ); - let result = self.process_state_nullify().await; + let result = self.process_state_nullify_hybrid(merkle_tree_data).await; if let Err(ref e) = result { error!( "State nullify failed for tree {}: {:?}", @@ -137,16 +253,45 @@ impl> BatchProcessor { } async fn verify_batch_ready(&self) -> BatchReadyState { - let mut rpc = match self.context.rpc_pool.get_connection().await { + let rpc = match self.context.rpc_pool.get_connection().await { Ok(rpc) => rpc, Err(_) => return BatchReadyState::NotReady, }; - let input_ready = self.verify_input_queue_batch_ready(&mut rpc).await; - let output_ready = if self.tree_type == TreeType::StateV2 { - self.verify_output_queue_batch_ready(&mut rpc).await + let merkle_tree_account = rpc + .get_account(self.context.merkle_tree) + .await + .ok() + .flatten(); + let output_queue_account = if self.tree_type == TreeType::StateV2 { + rpc.get_account(self.context.output_queue) + .await + .ok() + .flatten() + } else { + None + }; + + let (merkle_tree_data, input_ready) = if let Some(mut account) = merkle_tree_account { + match self.parse_merkle_tree_account(&mut account) { + Ok((data, ready)) => (Some(data), ready), + Err(_) => (None, false), + } + } else { + (None, false) + }; + + let (output_queue_data, output_ready) = if self.tree_type == TreeType::StateV2 { + if let Some(mut account) = output_queue_account { + match self.parse_output_queue_account(&mut account) { + Ok((data, ready)) => (Some(data), ready), + Err(_) => (None, false), + } + } else { + (None, false) + } } else { - false + (None, false) }; trace!( @@ -156,197 +301,162 @@ impl> BatchProcessor { output_ready ); + if !input_ready && !output_ready { + info!( + "QUEUE_METRIC: queue_empty tree_type={} tree={}", + self.tree_type, self.context.merkle_tree + ); + } else { + info!("QUEUE_METRIC: queue_has_elements tree_type={} tree={} input_ready={} output_ready={}", + self.tree_type, self.context.merkle_tree, input_ready, output_ready); + } + if self.tree_type == TreeType::AddressV2 { return if input_ready { - BatchReadyState::ReadyForAppend + if let Some(mt_data) = merkle_tree_data { + BatchReadyState::AddressReadyForAppend { + merkle_tree_data: mt_data, + } + } else { + BatchReadyState::NotReady + } } else { BatchReadyState::NotReady }; } - // For State tree type, balance append and nullify operations + // For State tree type, balance appends and nullifies operations // based on the queue states match (input_ready, output_ready) { (true, true) => { - // If both queues are ready, check their fill levels - let input_fill = self.get_input_queue_completion(&mut rpc).await; - let output_fill = self.get_output_queue_completion(&mut rpc).await; + if let (Some(mt_data), Some(oq_data)) = (merkle_tree_data, output_queue_data) { + // If both queues are ready, check their fill levels + let input_fill = Self::calculate_completion_from_parsed( + mt_data.num_inserted_zkps, + mt_data.current_zkp_batch_index, + ); + let output_fill = Self::calculate_completion_from_parsed( + oq_data.num_inserted_zkps, + oq_data.current_zkp_batch_index, + ); - trace!( - "Input queue fill: {:.2}, Output queue fill: {:.2}", - input_fill, - output_fill - ); - if input_fill > output_fill { - BatchReadyState::ReadyForNullify + trace!( + "Input queue fill: {:.2}, Output queue fill: {:.2}", + input_fill, + output_fill + ); + if input_fill > output_fill { + BatchReadyState::StateReadyForNullify { + merkle_tree_data: mt_data, + } + } else { + BatchReadyState::StateReadyForAppend { + merkle_tree_data: mt_data, + output_queue_data: oq_data, + } + } } else { - BatchReadyState::ReadyForAppend + BatchReadyState::NotReady + } + } + (true, false) => { + if let Some(mt_data) = merkle_tree_data { + BatchReadyState::StateReadyForNullify { + merkle_tree_data: mt_data, + } + } else { + BatchReadyState::NotReady + } + } + (false, true) => { + if let (Some(mt_data), Some(oq_data)) = (merkle_tree_data, output_queue_data) { + BatchReadyState::StateReadyForAppend { + merkle_tree_data: mt_data, + output_queue_data: oq_data, + } + } else { + BatchReadyState::NotReady } } - (true, false) => BatchReadyState::ReadyForNullify, - (false, true) => BatchReadyState::ReadyForAppend, (false, false) => BatchReadyState::NotReady, } } - /// Get the completion percentage of the input queue - async fn get_input_queue_completion(&self, rpc: &mut R) -> f64 { - let mut account = match rpc.get_account(self.context.merkle_tree).await { - Ok(Some(account)) => account, - _ => return 0.0, - }; - - Self::calculate_completion_from_tree(account.data.as_mut_slice()) - } - - /// Get the completion percentage of the output queue - async fn get_output_queue_completion(&self, rpc: &mut R) -> f64 { - let mut account = match rpc.get_account(self.context.output_queue).await { - Ok(Some(account)) => account, - _ => return 0.0, - }; - - Self::calculate_completion_from_queue(account.data.as_mut_slice()) - } - - /// Calculate completion percentage from a merkle tree account - fn calculate_completion_from_tree(data: &mut [u8]) -> f64 { - let tree = match BatchedMerkleTreeAccount::state_from_bytes(data, &Pubkey::default().into()) - { - Ok(tree) => tree, - Err(_) => return 0.0, - }; - - let batch_index = tree.queue_batches.pending_batch_index; - match tree.queue_batches.batches.get(batch_index as usize) { - Some(batch) => Self::calculate_completion(batch), - None => 0.0, - } - } - - /// Calculate completion percentage from a queue account - fn calculate_completion_from_queue(data: &mut [u8]) -> f64 { - let queue = match BatchedQueueAccount::output_from_bytes(data) { - Ok(queue) => queue, - Err(_) => return 0.0, - }; - - let batch_index = queue.batch_metadata.pending_batch_index; - match queue.batch_metadata.batches.get(batch_index as usize) { - Some(batch) => Self::calculate_completion(batch), - None => 0.0, - } - } - - /// Calculate completion percentage for a batch - fn calculate_completion(batch: &Batch) -> f64 { - let total = batch.get_num_zkp_batches(); - if total == 0 { - return 0.0; - } - - let remaining = total - batch.get_num_inserted_zkps(); - remaining as f64 / total as f64 - } - /// Process state append operation - async fn process_state_append(&self) -> Result { - let mut rpc = self.context.rpc_pool.get_connection().await?; - let (_, zkp_batch_size) = self.get_num_inserted_zkps(&mut rpc).await?; + async fn process_state_nullify_hybrid( + &self, + merkle_tree_data: ParsedMerkleTreeData, + ) -> Result { + let zkp_batch_size = merkle_tree_data.zkp_batch_size as usize; let batch_hash = format!( - "state_append_{}_{}", + "state_nullify_hybrid_{}_{}", self.context.merkle_tree, self.context.epoch ); + { let mut cache = self.context.ops_cache.lock().await; if cache.contains(&batch_hash) { trace!( - "Skipping already processed state append batch: {}", + "Skipping already processed state nullify batch (hybrid): {}", batch_hash ); return Ok(0); } cache.add(&batch_hash); } - state::perform_append(&self.context, &mut rpc).await?; + + state::perform_nullify(&self.context, merkle_tree_data).await?; + trace!( - "State append operation completed for tree: {}", + "State nullify operation (hybrid) completed for tree: {}", self.context.merkle_tree ); - let mut cache = self.context.ops_cache.lock().await; cache.cleanup_by_key(&batch_hash); + trace!("Cache cleaned up for batch: {}", batch_hash); Ok(zkp_batch_size) } - /// Process state nullify operation - async fn process_state_nullify(&self) -> Result { - let mut rpc = self.context.rpc_pool.get_connection().await?; - let (inserted_zkps_count, zkp_batch_size) = self.get_num_inserted_zkps(&mut rpc).await?; - trace!( - "ZKP batch size: {}, inserted ZKPs count: {}", - zkp_batch_size, - inserted_zkps_count - ); + async fn process_state_append_hybrid( + &self, + merkle_tree_data: ParsedMerkleTreeData, + output_queue_data: ParsedQueueData, + ) -> Result { + let zkp_batch_size = output_queue_data.zkp_batch_size as usize; let batch_hash = format!( - "state_nullify_{}_{}", + "state_append_hybrid_{}_{}", self.context.merkle_tree, self.context.epoch ); - { let mut cache = self.context.ops_cache.lock().await; if cache.contains(&batch_hash) { trace!( - "Skipping already processed state nullify batch: {}", + "Skipping already processed state append batch (hybrid): {}", batch_hash ); return Ok(0); } cache.add(&batch_hash); } - - state::perform_nullify(&self.context, &mut rpc).await?; - + state::perform_append(&self.context, merkle_tree_data, output_queue_data).await?; trace!( - "State nullify operation completed for tree: {}", + "State append operation (hybrid) completed for tree: {}", self.context.merkle_tree ); + let mut cache = self.context.ops_cache.lock().await; cache.cleanup_by_key(&batch_hash); - trace!("Cache cleaned up for batch: {}", batch_hash); Ok(zkp_batch_size) } - /// Get the number of inserted ZKPs and the ZKP batch size - async fn get_num_inserted_zkps(&self, rpc: &mut R) -> Result<(u64, usize)> { - let (num_inserted_zkps, zkp_batch_size) = { - let mut output_queue_account = - rpc.get_account(self.context.output_queue).await?.unwrap(); - let output_queue = - BatchedQueueAccount::output_from_bytes(output_queue_account.data.as_mut_slice()) - .map_err(|e| BatchProcessError::QueueParsing(e.to_string()))?; - - let batch_index = output_queue.batch_metadata.pending_batch_index; - let zkp_batch_size = output_queue.batch_metadata.zkp_batch_size; - - ( - output_queue.batch_metadata.batches[batch_index as usize].get_num_inserted_zkps(), - zkp_batch_size as usize, - ) - }; - Ok((num_inserted_zkps, zkp_batch_size)) - } - - /// Verify if the input queue batch is ready for processing - async fn verify_input_queue_batch_ready(&self, rpc: &mut R) -> bool { - let mut account = match rpc.get_account(self.context.merkle_tree).await { - Ok(Some(account)) => account, - _ => return false, - }; - + /// Parse merkle tree account and check if batch is ready + fn parse_merkle_tree_account( + &self, + account: &mut solana_sdk::account::Account, + ) -> Result<(ParsedMerkleTreeData, bool)> { let merkle_tree = match self.tree_type { TreeType::AddressV2 => BatchedMerkleTreeAccount::address_from_bytes( account.data.as_mut_slice(), @@ -356,50 +466,89 @@ impl> BatchProcessor { account.data.as_mut_slice(), &self.context.merkle_tree.into(), ), - _ => return false, + _ => return Err(ForesterError::InvalidTreeType(self.tree_type).into()), + }?; + + let batch_index = merkle_tree.queue_batches.pending_batch_index; + let batch = merkle_tree + .queue_batches + .batches + .get(batch_index as usize) + .ok_or_else(|| anyhow::anyhow!("Batch not found"))?; + + let num_inserted_zkps = batch.get_num_inserted_zkps(); + let current_zkp_batch_index = batch.get_current_zkp_batch_index(); + + let mut leaves_hash_chains = Vec::new(); + for i in num_inserted_zkps..current_zkp_batch_index { + leaves_hash_chains + .push(merkle_tree.hash_chain_stores[batch_index as usize][i as usize]); + } + + let parsed_data = ParsedMerkleTreeData { + next_index: merkle_tree.next_index, + current_root: *merkle_tree.root_history.last().unwrap(), + root_history: merkle_tree.root_history.to_vec(), + zkp_batch_size: batch.zkp_batch_size as u16, + pending_batch_index: batch_index as u32, + num_inserted_zkps, + current_zkp_batch_index, + leaves_hash_chains, }; - if let Ok(tree) = merkle_tree { - let batch_index = tree.queue_batches.pending_batch_index; - let full_batch = tree - .queue_batches - .batches - .get(batch_index as usize) - .unwrap(); + let is_ready = batch.get_state() != BatchState::Inserted + && batch.get_current_zkp_batch_index() > batch.get_num_inserted_zkps(); - full_batch.get_state() != BatchState::Inserted - && full_batch.get_current_zkp_batch_index() > full_batch.get_num_inserted_zkps() - } else { - false - } + Ok((parsed_data, is_ready)) } - /// Verify if the output queue batch is ready for processing - async fn verify_output_queue_batch_ready(&self, rpc: &mut R) -> bool { - let mut account = match rpc.get_account(self.context.output_queue).await { - Ok(Some(account)) => account, - _ => return false, - }; + /// Parse output queue account and check if batch is ready + fn parse_output_queue_account( + &self, + account: &mut solana_sdk::account::Account, + ) -> Result<(ParsedQueueData, bool)> { + let output_queue = BatchedQueueAccount::output_from_bytes(account.data.as_mut_slice())?; + + let batch_index = output_queue.batch_metadata.pending_batch_index; + let batch = output_queue + .batch_metadata + .batches + .get(batch_index as usize) + .ok_or_else(|| anyhow::anyhow!("Batch not found"))?; + + let num_inserted_zkps = batch.get_num_inserted_zkps(); + let current_zkp_batch_index = batch.get_current_zkp_batch_index(); + + let mut leaves_hash_chains = Vec::new(); + for i in num_inserted_zkps..current_zkp_batch_index { + leaves_hash_chains + .push(output_queue.hash_chain_stores[batch_index as usize][i as usize]); + } - let output_queue = match self.tree_type { - TreeType::StateV2 => { - BatchedQueueAccount::output_from_bytes(account.data.as_mut_slice()) - } - _ => return false, + let parsed_data = ParsedQueueData { + zkp_batch_size: output_queue.batch_metadata.zkp_batch_size as u16, + pending_batch_index: batch_index as u32, + num_inserted_zkps, + current_zkp_batch_index, + leaves_hash_chains, }; - if let Ok(queue) = output_queue { - let batch_index = queue.batch_metadata.pending_batch_index; - let full_batch = queue - .batch_metadata - .batches - .get(batch_index as usize) - .unwrap(); + let is_ready = batch.get_state() != BatchState::Inserted + && batch.get_current_zkp_batch_index() > batch.get_num_inserted_zkps(); - full_batch.get_state() != BatchState::Inserted - && full_batch.get_current_zkp_batch_index() > full_batch.get_num_inserted_zkps() - } else { - false + Ok((parsed_data, is_ready)) + } + + /// Calculate completion percentage from parsed data + fn calculate_completion_from_parsed( + num_inserted_zkps: u64, + current_zkp_batch_index: u64, + ) -> f64 { + let total = current_zkp_batch_index; + if total == 0 { + return 0.0; } + let remaining = total - num_inserted_zkps; + remaining as f64 / total as f64 } } diff --git a/forester/src/processor/v2/error.rs b/forester/src/processor/v2/error.rs deleted file mode 100644 index 77a468653b..0000000000 --- a/forester/src/processor/v2/error.rs +++ /dev/null @@ -1,54 +0,0 @@ -use forester_utils::rpc_pool::PoolError; -use light_compressed_account::TreeType; -use solana_client::rpc_request::RpcError; -use thiserror::Error; - -pub type Result = std::result::Result; - -#[derive(Debug, Error)] -pub enum BatchProcessError { - #[error("Failed to parse queue account: {0}")] - QueueParsing(String), - - #[error("Failed to parse merkle tree account: {0}")] - MerkleTreeParsing(String), - - #[error("Failed to create instruction data: {0}")] - InstructionData(String), - - #[error("Transaction failed: {0}")] - Transaction(String), - - #[error("RPC error: {0}")] - Rpc(String), - - #[error("Pool error: {0}")] - Pool(String), - - #[error("Indexer error: {0}")] - Indexer(String), - - #[error("Unsupported tree type: {0:?}")] - UnsupportedTreeType(TreeType), - - #[error(transparent)] - Other(#[from] anyhow::Error), -} - -impl From for BatchProcessError { - fn from(e: light_client::rpc::RpcError) -> Self { - Self::Rpc(e.to_string()) - } -} - -impl From for BatchProcessError { - fn from(e: RpcError) -> Self { - Self::Rpc(e.to_string()) - } -} - -impl From for BatchProcessError { - fn from(e: PoolError) -> Self { - Self::Pool(e.to_string()) - } -} diff --git a/forester/src/processor/v2/mod.rs b/forester/src/processor/v2/mod.rs index f33793ce76..3f5367800b 100644 --- a/forester/src/processor/v2/mod.rs +++ b/forester/src/processor/v2/mod.rs @@ -1,13 +1,13 @@ mod address; mod common; -mod error; mod state; use common::BatchProcessor; -use error::Result; use light_client::rpc::Rpc; use tracing::{instrument, trace}; +use crate::Result; + #[instrument( level = "debug", fields( @@ -17,7 +17,7 @@ use tracing::{instrument, trace}; ), skip(context) )] -pub async fn process_batched_operations>( +pub async fn process_batched_operations( context: BatchContext, tree_type: TreeType, ) -> Result { @@ -27,8 +27,5 @@ pub async fn process_batched_operations>( } pub use common::BatchContext; -pub use error::BatchProcessError; use light_client::indexer::Indexer; use light_compressed_account::TreeType; - -use crate::indexer_type::IndexerType; diff --git a/forester/src/processor/v2/state.rs b/forester/src/processor/v2/state.rs index 74393a742d..abf21056b4 100644 --- a/forester/src/processor/v2/state.rs +++ b/forester/src/processor/v2/state.rs @@ -1,259 +1,142 @@ +use anyhow::{Error, Ok}; use borsh::BorshSerialize; use forester_utils::instructions::{ - state_batch_append::create_append_batch_ix_data, - state_batch_nullify::create_nullify_batch_ix_data, + state_batch_append::get_append_instruction_stream, + state_batch_nullify::get_nullify_instruction_stream, +}; +use futures::stream::{Stream, StreamExt}; +use light_batched_merkle_tree::merkle_tree::{ + InstructionDataBatchAppendInputs, InstructionDataBatchNullifyInputs, }; use light_client::{indexer::Indexer, rpc::Rpc}; use light_registry::account_compression_cpi::sdk::{ create_batch_append_instruction, create_batch_nullify_instruction, }; +use solana_program::instruction::Instruction; use solana_sdk::signer::Signer; -use tracing::{debug, info, instrument, log::error, trace}; - -use super::{ - common::BatchContext, - error::{BatchProcessError, Result}, -}; -use crate::indexer_type::{ - update_test_indexer_after_append, update_test_indexer_after_nullification, IndexerType, -}; - -#[instrument( - level = "debug", - fields( - forester = %context.derivation, - epoch = %context.derivation, - merkle_tree = %context.merkle_tree, - output_queue = %context.output_queue, - ), skip(context, rpc)) -] -pub(crate) async fn perform_append>( - context: &BatchContext, - rpc: &mut R, -) -> Result<()> { - let instruction_data_vec = create_append_batch_ix_data( - rpc, - &mut *context.indexer.lock().await, - context.merkle_tree, - context.output_queue, - context.prover_url.clone(), - context.prover_polling_interval, - context.prover_max_wait_time, +use tracing::{info, instrument}; + +use super::common::{process_stream, BatchContext, ParsedMerkleTreeData, ParsedQueueData}; +use crate::Result; + +async fn create_nullify_stream_future( + ctx: &BatchContext, + merkle_tree_data: ParsedMerkleTreeData, +) -> Result<( + impl Stream>> + Send, + u16, +)> +where + R: Rpc, + I: Indexer + 'static, +{ + let (stream, size) = get_nullify_instruction_stream( + ctx.rpc_pool.clone(), + ctx.indexer.clone(), + ctx.merkle_tree, + ctx.prover_url.clone(), + ctx.prover_polling_interval, + ctx.prover_max_wait_time, + merkle_tree_data, + ctx.ixs_per_tx, ) .await - .map_err(|e| { - error!("Failed to create append batch instruction data: {}", e); - BatchProcessError::InstructionData(e.to_string()) - })?; + .map_err(Error::from)?; + let stream = stream.map(|item| item.map_err(Error::from)); + Ok((stream, size)) +} - if instruction_data_vec.is_empty() { - trace!("No zkp batches to append"); - let mut cache = context.ops_cache.lock().await; - cache.cleanup(); - return Ok(()); - } +async fn create_append_stream_future( + ctx: &BatchContext, + merkle_tree_data: ParsedMerkleTreeData, + output_queue_data: ParsedQueueData, +) -> Result<( + impl Stream>> + Send, + u16, +)> +where + R: Rpc, + I: Indexer + 'static, +{ + let (stream, size) = get_append_instruction_stream( + ctx.rpc_pool.clone(), + ctx.indexer.clone(), + ctx.merkle_tree, + ctx.prover_url.clone(), + ctx.prover_polling_interval, + ctx.prover_max_wait_time, + merkle_tree_data, + output_queue_data, + ctx.ixs_per_tx, + ) + .await + .map_err(Error::from)?; + let stream = stream.map(|item| item.map_err(Error::from)); + Ok((stream, size)) +} +#[instrument(level = "debug", skip(context))] +pub(crate) async fn perform_nullify( + context: &BatchContext, + merkle_tree_data: ParsedMerkleTreeData, +) -> Result<()> { info!( - "Processing {} ZKP batch appends", - instruction_data_vec.len() + "V2_TPS_METRIC: operation_start tree_type=StateV2 operation=nullify tree={} epoch={} (hybrid)", + context.merkle_tree, context.epoch ); - for (chunk_idx, instruction_chunk) in - instruction_data_vec.chunks(context.ixs_per_tx).enumerate() - { - debug!( - "Sending append transaction chunk {}/{} for tree: {}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - context.merkle_tree - ); - - let mut instructions = Vec::with_capacity(context.ixs_per_tx); - for instruction_data in instruction_chunk { - debug!( - "Instruction data size: {} bytes", - instruction_data.try_to_vec().map(|v| v.len()).unwrap_or(0) - ); - - instructions.push(create_batch_append_instruction( - context.authority.pubkey(), - context.derivation, - context.merkle_tree, - context.output_queue, - context.epoch, - instruction_data - .try_to_vec() - .map_err(|e| BatchProcessError::InstructionData(e.to_string()))?, - )); - } - - match rpc - .create_and_send_transaction( - &instructions, - &context.authority.pubkey(), - &[&context.authority], - ) - .await - { - Ok(tx) => { - info!( - "Append transaction chunk {}/{} sent successfully: {}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - tx - ); - } - Err(e) => { - error!( - "Failed to send append transaction chunk {}/{} for tree {}: {:?}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - context.merkle_tree, - e - ); - return Err(e.into()); - } - } - - update_test_indexer_after_append( - rpc, - context.indexer.clone(), + let instruction_builder = |data: &InstructionDataBatchNullifyInputs| -> Instruction { + create_batch_nullify_instruction( + context.authority.pubkey(), + context.derivation, context.merkle_tree, - context.output_queue, + context.epoch, + data.try_to_vec().unwrap(), ) - .await - .map_err(|e| { - error!("Failed to update test indexer after append: {:?}", e); - BatchProcessError::Indexer(e.to_string()) - })?; - } + }; + + let stream_future = create_nullify_stream_future(context, merkle_tree_data); + process_stream( + context, + stream_future, + instruction_builder, + "StateV2", + Some("nullify"), + ) + .await?; Ok(()) } -/// Perform a state nullify operation for a Merkle tree -#[instrument( - level = "debug", - fields( - forester = %context.derivation, - epoch = %context.epoch, - merkle_tree = %context.merkle_tree, - ), - skip(context, rpc) -)] -pub(crate) async fn perform_nullify>( +#[instrument(level = "debug", skip(context))] +pub(crate) async fn perform_append( context: &BatchContext, - rpc: &mut R, + merkle_tree_data: ParsedMerkleTreeData, + output_queue_data: ParsedQueueData, ) -> Result<()> { - let batch_index = get_batch_index(context, rpc).await?; - let instruction_data_vec = create_nullify_batch_ix_data( - rpc, - &mut *context.indexer.lock().await, - context.merkle_tree, - context.prover_url.clone(), - context.prover_polling_interval, - context.prover_max_wait_time, - ) - .await - .map_err(|e| { - error!("Failed to create nullify batch instruction data: {}", e); - BatchProcessError::InstructionData(e.to_string()) - })?; - - if instruction_data_vec.is_empty() { - trace!("No zkp batches to nullify"); - let mut cache = context.ops_cache.lock().await; - cache.cleanup(); - return Ok(()); - } - info!( - "Processing {} ZKP batch nullifications", - instruction_data_vec.len() + "V2_TPS_METRIC: operation_start tree_type=StateV2 operation=append tree={} epoch={} (hybrid)", + context.merkle_tree, context.epoch ); - - for (chunk_idx, instruction_chunk) in - instruction_data_vec.chunks(context.ixs_per_tx).enumerate() - { - debug!( - "Processing nullify transaction chunk {}/{}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx) - ); - - let mut instructions = Vec::with_capacity(context.ixs_per_tx); - for instruction_data in instruction_chunk { - instructions.push(create_batch_nullify_instruction( - context.authority.pubkey(), - context.derivation, - context.merkle_tree, - context.epoch, - instruction_data - .try_to_vec() - .map_err(|e| BatchProcessError::InstructionData(e.to_string()))?, - )); - } - - match rpc - .create_and_send_transaction( - &instructions, - &context.authority.pubkey(), - &[&context.authority], - ) - .await - { - Ok(tx) => { - info!( - "Nullify transaction chunk {}/{} sent successfully: {}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - tx - ); - tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; - } - Err(e) => { - error!( - "Failed to send nullify transaction chunk {}/{} for tree {}: {:?}", - chunk_idx + 1, - instruction_data_vec.len().div_ceil(context.ixs_per_tx), - context.merkle_tree, - e - ); - return Err(e.into()); - } - } - - update_test_indexer_after_nullification( - rpc, - context.indexer.clone(), + let instruction_builder = |data: &InstructionDataBatchAppendInputs| -> Instruction { + create_batch_append_instruction( + context.authority.pubkey(), + context.derivation, context.merkle_tree, - batch_index, + context.output_queue, + context.epoch, + data.try_to_vec().unwrap(), ) - .await - .map_err(|e| { - error!("Failed to update test indexer after nullification: {:?}", e); - BatchProcessError::Indexer(e.to_string()) - })?; - } - + }; + + let stream_future = create_append_stream_future(context, merkle_tree_data, output_queue_data); + process_stream( + context, + stream_future, + instruction_builder, + "StateV2", + Some("append"), + ) + .await?; Ok(()) } - -/// Get the current batch index from the Merkle tree account -async fn get_batch_index( - context: &BatchContext, - rpc: &mut R, -) -> Result { - let mut account = rpc.get_account(context.merkle_tree).await?.ok_or_else(|| { - BatchProcessError::Rpc(format!("Account not found: {}", context.merkle_tree)) - })?; - - let merkle_tree = - light_batched_merkle_tree::merkle_tree::BatchedMerkleTreeAccount::state_from_bytes( - account.data.as_mut_slice(), - &context.merkle_tree.into(), - ) - .map_err(|e| BatchProcessError::MerkleTreeParsing(e.to_string()))?; - - Ok(merkle_tree.queue_batches.pending_batch_index as usize) -} diff --git a/forester/src/rollover/mod.rs b/forester/src/rollover/mod.rs index 0c4985de26..d1f6657bb4 100644 --- a/forester/src/rollover/mod.rs +++ b/forester/src/rollover/mod.rs @@ -1,8 +1,69 @@ mod operations; mod state; +use std::sync::Arc; +use forester_utils::forester_epoch::TreeAccounts; +use light_client::rpc::Rpc; pub use operations::{ get_tree_fullness, is_tree_ready_for_rollover, perform_address_merkle_tree_rollover, perform_state_merkle_tree_rollover_forester, }; +use solana_sdk::{pubkey::Pubkey, signature::Keypair}; pub use state::RolloverState; +use tracing::info; + +use crate::{errors::ForesterError, ForesterConfig}; + +pub async fn rollover_state_merkle_tree( + config: Arc, + rpc: &mut R, + tree_accounts: &TreeAccounts, + epoch: u64, +) -> Result<(), ForesterError> { + let new_nullifier_queue_keypair = Keypair::new(); + let new_merkle_tree_keypair = Keypair::new(); + let new_cpi_signature_keypair = Keypair::new(); + + let rollover_signature = perform_state_merkle_tree_rollover_forester( + &config.payer_keypair, + &config.derivation_pubkey, + rpc, + &new_nullifier_queue_keypair, + &new_merkle_tree_keypair, + &new_cpi_signature_keypair, + &tree_accounts.merkle_tree, + &tree_accounts.queue, + &Pubkey::default(), + epoch, + ) + .await?; + + info!("State rollover signature: {:?}", rollover_signature); + Ok(()) +} + +pub async fn rollover_address_merkle_tree( + config: Arc, + rpc: &mut R, + tree_accounts: &TreeAccounts, + epoch: u64, +) -> Result<(), ForesterError> { + let new_nullifier_queue_keypair = Keypair::new(); + let new_merkle_tree_keypair = Keypair::new(); + + let rollover_signature = perform_address_merkle_tree_rollover( + &config.payer_keypair, + &config.derivation_pubkey, + rpc, + &new_nullifier_queue_keypair, + &new_merkle_tree_keypair, + &tree_accounts.merkle_tree, + &tree_accounts.queue, + epoch, + ) + .await?; + + info!("Address rollover signature: {:?}", rollover_signature); + + Ok(()) +} diff --git a/forester/tests/address_v2_test.rs b/forester/tests/address_v2_test.rs index d1a4671436..65729024a2 100644 --- a/forester/tests/address_v2_test.rs +++ b/forester/tests/address_v2_test.rs @@ -103,7 +103,24 @@ async fn test_create_v2_address() { println!("num_batches: {:?}", num_batches); println!("remaining_addresses: {:?}", remaining_addresses); - for i in 0..num_batches { + let mut address_tree_account = rpc + .get_account(env.v2_address_trees[0]) + .await + .unwrap() + .unwrap(); + + let address_tree = BatchedMerkleTreeAccount::address_from_bytes( + address_tree_account.data.as_mut_slice(), + &env.v2_address_trees[0].into(), + ) + .unwrap(); + + println!("Address tree metadata: {:?}", address_tree.get_metadata()); + + let (service_handle, shutdown_sender, mut work_report_receiver) = + setup_forester_pipeline(&config).await; + + for i in 0..num_batches * 10 { println!("====== Creating v2 address {} ======", i); let result = create_v2_addresses( &mut rpc, @@ -133,23 +150,6 @@ async fn test_create_v2_address() { println!("====== result: {:?} ======", result); } - let mut address_tree_account = rpc - .get_account(env.v2_address_trees[0]) - .await - .unwrap() - .unwrap(); - - let address_tree = BatchedMerkleTreeAccount::address_from_bytes( - address_tree_account.data.as_mut_slice(), - &env.v2_address_trees[0].into(), - ) - .unwrap(); - - println!("Address tree metadata: {:?}", address_tree.get_metadata()); - - let (service_handle, shutdown_sender, mut work_report_receiver) = - setup_forester_pipeline(&config).await; - wait_for_work_report(&mut work_report_receiver, &tree_params).await; verify_root_changed(&mut rpc, &env.v2_address_trees[0], &pre_root).await; diff --git a/forester/tests/e2e_v1_test.rs b/forester/tests/e2e_v1_test.rs new file mode 100644 index 0000000000..f2149f5457 --- /dev/null +++ b/forester/tests/e2e_v1_test.rs @@ -0,0 +1,556 @@ +use std::{collections::HashSet, sync::Arc, time::Duration}; + +use account_compression::{ + utils::constants::{ADDRESS_QUEUE_VALUES, STATE_NULLIFIER_QUEUE_VALUES}, + AddressMerkleTreeAccount, +}; +use forester::{queue_helpers::fetch_queue_item_data, run_pipeline, utils::get_protocol_config}; +use forester_utils::{ + registry::register_test_forester, + rpc_pool::{SolanaRpcPool, SolanaRpcPoolBuilder}, +}; +use light_client::{ + indexer::{AddressMerkleTreeAccounts, StateMerkleTreeAccounts}, + local_test_validator::{LightValidatorConfig, ProverConfig}, + rpc::{client::RpcUrl, LightClient, LightClientConfig, Rpc, RpcError}, +}; +use light_program_test::{accounts::test_accounts::TestAccounts, indexer::TestIndexer}; +use light_registry::{utils::get_forester_epoch_pda_from_authority, ForesterEpochPda}; +use light_test_utils::{e2e_test_env::E2ETestEnv, update_test_forester}; +use serial_test::serial; +use solana_sdk::{ + commitment_config::CommitmentConfig, native_token::LAMPORTS_PER_SOL, pubkey::Pubkey, + signature::Keypair, signer::Signer, +}; +use tokio::{ + sync::{mpsc, oneshot, Mutex}, + time::sleep, +}; + +mod test_utils; +use test_utils::*; + +#[serial] +#[tokio::test(flavor = "multi_thread", worker_threads = 32)] +async fn test_e2e_v1() { + init(Some(LightValidatorConfig { + enable_indexer: false, + wait_time: 90, + prover_config: Some(ProverConfig::default()), + sbf_programs: vec![], + limit_ledger_size: None, + })) + .await; + let forester_keypair1 = Keypair::new(); + let forester_keypair2 = Keypair::new(); + + let mut test_accounts = TestAccounts::get_local_test_validator_accounts(); + test_accounts.protocol.forester = forester_keypair1.insecure_clone(); + + let mut config1 = forester_config(); + config1.payer_keypair = forester_keypair1.insecure_clone(); + + let mut config2 = forester_config(); + config2.payer_keypair = forester_keypair2.insecure_clone(); + + let pool = SolanaRpcPoolBuilder::::default() + .url(config1.external_services.rpc_url.to_string()) + .commitment(CommitmentConfig::confirmed()) + .build() + .await + .unwrap(); + + let mut rpc = LightClient::new(LightClientConfig::local_no_indexer()) + .await + .unwrap(); + rpc.payer = forester_keypair1.insecure_clone(); + + // Airdrop to both foresters and governance authority + for keypair in [ + &forester_keypair1, + &forester_keypair2, + &test_accounts.protocol.governance_authority, + ] { + rpc.airdrop_lamports(&keypair.pubkey(), LAMPORTS_PER_SOL * 100_000) + .await + .unwrap(); + } + + // Register both foresters + for forester_keypair in [&forester_keypair1, &forester_keypair2] { + register_test_forester( + &mut rpc, + &test_accounts.protocol.governance_authority, + &forester_keypair.pubkey(), + light_registry::ForesterConfig::default(), + ) + .await + .unwrap(); + } + + let new_forester_keypair1 = Keypair::new(); + let new_forester_keypair2 = Keypair::new(); + + for forester_keypair in [&new_forester_keypair1, &new_forester_keypair2] { + rpc.airdrop_lamports(&forester_keypair.pubkey(), LAMPORTS_PER_SOL * 100_000) + .await + .unwrap(); + } + + update_test_forester( + &mut rpc, + &forester_keypair1, + &forester_keypair1.pubkey(), + Some(&new_forester_keypair1), + light_registry::ForesterConfig::default(), + ) + .await + .unwrap(); + + update_test_forester( + &mut rpc, + &forester_keypair2, + &forester_keypair2.pubkey(), + Some(&new_forester_keypair2), + light_registry::ForesterConfig::default(), + ) + .await + .unwrap(); + + config1.derivation_pubkey = forester_keypair1.pubkey(); + config1.payer_keypair = new_forester_keypair1.insecure_clone(); + + config2.derivation_pubkey = forester_keypair2.pubkey(); + config2.payer_keypair = new_forester_keypair2.insecure_clone(); + + let config1 = Arc::new(config1); + let config2 = Arc::new(config2); + + let indexer: TestIndexer = + TestIndexer::init_from_acounts(&config1.payer_keypair, &test_accounts, 0).await; + + let mut env = E2ETestEnv::::new( + rpc, + indexer, + &test_accounts, + keypair_action_config(), + general_action_config(), + 0, + Some(0), + ) + .await; + // removing batched Merkle tree + env.indexer.state_merkle_trees.remove(1); + // removing batched address tree + env.indexer.address_merkle_trees.remove(1); + let user_index = 0; + let balance = env + .rpc + .get_balance(&env.users[user_index].keypair.pubkey()) + .await + .unwrap(); + env.compress_sol(user_index, balance).await; + // Create state and address trees which can be rolled over + env.create_address_tree(Some(0)).await; + env.create_state_tree(Some(0)).await; + let state_tree_with_rollover_threshold_0 = + env.indexer.state_merkle_trees[1].accounts.merkle_tree; + let address_tree_with_rollover_threshold_0 = + env.indexer.address_merkle_trees[1].accounts.merkle_tree; + + println!( + "State tree with rollover threshold 0: {:?}", + state_tree_with_rollover_threshold_0 + ); + println!( + "Address tree with rollover threshold 0: {:?}", + address_tree_with_rollover_threshold_0 + ); + + let state_trees: Vec = env + .indexer + .state_merkle_trees + .iter() + .map(|x| x.accounts) + .collect(); + let address_trees: Vec = env + .indexer + .address_merkle_trees + .iter() + .map(|x| x.accounts) + .collect(); + + println!("Address trees: {:?}", address_trees); + + // Two rollovers plus other work + let mut total_expected_work = 2; + { + let iterations = 5; + for i in 0..iterations { + println!("Round {} of {}", i, iterations); + let user_keypair = env.users[0].keypair.insecure_clone(); + env.transfer_sol_deterministic(&user_keypair, &user_keypair.pubkey(), Some(1)) + .await + .unwrap(); + env.transfer_sol_deterministic(&user_keypair, &user_keypair.pubkey().clone(), Some(0)) + .await + .unwrap(); + sleep(Duration::from_millis(100)).await; + env.create_address(None, Some(1)).await; + env.create_address(None, Some(0)).await; + } + assert_queue_len( + &pool, + &state_trees, + &address_trees, + &mut total_expected_work, + 0, + true, + ) + .await; + } + + let (shutdown_sender1, shutdown_receiver1) = oneshot::channel(); + let (shutdown_sender2, shutdown_receiver2) = oneshot::channel(); + let (work_report_sender1, mut work_report_receiver1) = mpsc::channel(100); + let (work_report_sender2, mut work_report_receiver2) = mpsc::channel(100); + + let indexer = Arc::new(Mutex::new(env.indexer)); + + let service_handle1 = tokio::spawn(run_pipeline::( + config1.clone(), + None, + None, + indexer.clone(), + shutdown_receiver1, + work_report_sender1, + )); + let service_handle2 = tokio::spawn(run_pipeline::( + config2.clone(), + None, + None, + indexer, + shutdown_receiver2, + work_report_sender2, + )); + + const EXPECTED_EPOCHS: u64 = 3; // We expect to process 2 epochs (0 and 1) + + let mut processed_epochs = HashSet::new(); + let mut total_processed = 0; + while processed_epochs.len() < EXPECTED_EPOCHS as usize { + tokio::select! { + Some(report) = work_report_receiver1.recv() => { + println!("Received work report from forester 1: {:?}", report); + total_processed += report.processed_items; + processed_epochs.insert(report.epoch); + } + Some(report) = work_report_receiver2.recv() => { + println!("Received work report from forester 2: {:?}", report); + total_processed += report.processed_items; + processed_epochs.insert(report.epoch); + } + else => break, + } + } + + println!("Processed {} items", total_processed); + + // Verify that we've processed the expected number of epochs + assert_eq!( + processed_epochs.len(), + EXPECTED_EPOCHS as usize, + "Processed {} epochs, expected {}", + processed_epochs.len(), + EXPECTED_EPOCHS + ); + + // Verify that we've processed epochs 0 and 1 + // assert!(processed_epochs.contains(&0), "Epoch 0 was not processed"); + assert!(processed_epochs.contains(&1), "Epoch 1 was not processed"); + + assert_trees_are_rolledover( + &pool, + &state_tree_with_rollover_threshold_0, + &address_tree_with_rollover_threshold_0, + ) + .await; + // assert queues have been emptied + assert_queue_len(&pool, &state_trees, &address_trees, &mut 0, 0, false).await; + let mut rpc = pool.get_connection().await.unwrap(); + let forester_pubkeys = [config1.derivation_pubkey, config2.derivation_pubkey]; + + // assert that foresters registered for epoch 1 and 2 (no new work is emitted after epoch 0) + // Assert that foresters have registered all processed epochs and the next epoch (+1) + for epoch in 1..=EXPECTED_EPOCHS { + let total_processed_work = + assert_foresters_registered(&forester_pubkeys[..], &mut rpc, epoch) + .await + .unwrap(); + if epoch == 1 { + assert_eq!( + total_processed_work, total_expected_work, + "Not all items were processed." + ); + } else { + assert_eq!( + total_processed_work, 0, + "Not all items were processed in prior epoch." + ); + } + } + + shutdown_sender1 + .send(()) + .expect("Failed to send shutdown signal to forester 1"); + shutdown_sender2 + .send(()) + .expect("Failed to send shutdown signal to forester 2"); + service_handle1.await.unwrap().unwrap(); + service_handle2.await.unwrap().unwrap(); +} +pub async fn assert_trees_are_rolledover( + pool: &SolanaRpcPool, + state_tree_with_rollover_threshold_0: &Pubkey, + address_tree_with_rollover_threshold_0: &Pubkey, +) { + let rpc = pool.get_connection().await.unwrap(); + let address_merkle_tree = rpc + .get_anchor_account::(address_tree_with_rollover_threshold_0) + .await + .unwrap() + .unwrap(); + assert_ne!( + address_merkle_tree + .metadata + .rollover_metadata + .rolledover_slot, + u64::MAX, + "address_merkle_tree: {:?}", + address_merkle_tree + ); + let state_merkle_tree = rpc + .get_anchor_account::(state_tree_with_rollover_threshold_0) + .await + .unwrap() + .unwrap(); + assert_ne!( + state_merkle_tree.metadata.rollover_metadata.rolledover_slot, + u64::MAX, + "state_merkle_tree: {:?}", + state_merkle_tree + ); +} + +async fn assert_foresters_registered( + foresters: &[Pubkey], + rpc: &mut LightClient, + epoch: u64, +) -> Result { + let mut performed_work = 0; + for (i, forester) in foresters.iter().enumerate() { + let forester_epoch_pda = get_forester_epoch_pda_from_authority(forester, epoch).0; + let forester_epoch_pda = rpc + .get_anchor_account::(&forester_epoch_pda) + .await?; + println!("forester_epoch_pda {}: {:?}", i, forester_epoch_pda); + + if let Some(forester_epoch_pda) = forester_epoch_pda { + // If one forester is first for both queues there will be no work left + // - this assert is flaky + // assert!( + // forester_epoch_pda.work_counter > 0, + // "forester {} did not perform any work", + // i + // ); + performed_work += forester_epoch_pda.work_counter; + } else { + return Err(RpcError::CustomError(format!( + "Forester {} not registered", + i, + ))); + } + } + Ok(performed_work) +} + +#[serial] +#[tokio::test(flavor = "multi_thread", worker_threads = 32)] +async fn test_epoch_double_registration() { + println!("*****************************************************************"); + + init(Some(LightValidatorConfig { + enable_indexer: false, + wait_time: 90, + prover_config: Some(ProverConfig::default()), + sbf_programs: vec![], + limit_ledger_size: None, + })) + .await; + + let forester_keypair = Keypair::new(); + + let mut test_accounts = TestAccounts::get_local_test_validator_accounts(); + test_accounts.protocol.forester = forester_keypair.insecure_clone(); + + let mut config = forester_config(); + config.payer_keypair = forester_keypair.insecure_clone(); + + let pool = SolanaRpcPoolBuilder::::default() + .url(config.external_services.rpc_url.to_string()) + .commitment(CommitmentConfig::confirmed()) + .build() + .await + .unwrap(); + + let mut rpc = LightClient::new(LightClientConfig { + url: RpcUrl::Localnet.to_string(), + photon_url: None, + api_key: None, + commitment_config: Some(CommitmentConfig::confirmed()), + fetch_active_tree: false, + }) + .await + .unwrap(); + rpc.payer = forester_keypair.insecure_clone(); + + rpc.airdrop_lamports(&forester_keypair.pubkey(), LAMPORTS_PER_SOL * 100_000) + .await + .unwrap(); + + rpc.airdrop_lamports( + &test_accounts.protocol.governance_authority.pubkey(), + LAMPORTS_PER_SOL * 100_000, + ) + .await + .unwrap(); + + register_test_forester( + &mut rpc, + &test_accounts.protocol.governance_authority, + &forester_keypair.pubkey(), + light_registry::ForesterConfig::default(), + ) + .await + .unwrap(); + + let new_forester_keypair = Keypair::new(); + + rpc.airdrop_lamports(&new_forester_keypair.pubkey(), LAMPORTS_PER_SOL * 100_000) + .await + .unwrap(); + update_test_forester( + &mut rpc, + &forester_keypair, + &forester_keypair.pubkey(), + Some(&new_forester_keypair), + light_registry::ForesterConfig::default(), + ) + .await + .unwrap(); + + config.derivation_pubkey = forester_keypair.pubkey(); + config.payer_keypair = new_forester_keypair.insecure_clone(); + + let config = Arc::new(config); + + let mut indexer: TestIndexer = + TestIndexer::init_from_acounts(&config.payer_keypair, &test_accounts, 0).await; + indexer.state_merkle_trees.remove(1); + let indexer = Arc::new(Mutex::new(indexer)); + + for _ in 0..10 { + let (shutdown_sender, shutdown_receiver) = oneshot::channel(); + let (work_report_sender, _work_report_receiver) = mpsc::channel(100); + + // Run the forester pipeline + let service_handle = tokio::spawn(run_pipeline::( + config.clone(), + None, + None, + indexer.clone(), + shutdown_receiver, + work_report_sender.clone(), + )); + + sleep(Duration::from_secs(2)).await; + + shutdown_sender + .send(()) + .expect("Failed to send shutdown signal"); + let result = service_handle.await.unwrap(); + assert!(result.is_ok(), "Registration should succeed"); + } + + let mut rpc = pool.get_connection().await.unwrap(); + let protocol_config = get_protocol_config(&mut *rpc).await; + let solana_slot = rpc.get_slot().await.unwrap(); + let current_epoch = protocol_config.get_current_epoch(solana_slot); + + let forester_epoch_pda_address = + get_forester_epoch_pda_from_authority(&config.derivation_pubkey, current_epoch).0; + + let forester_epoch_pda = rpc + .get_anchor_account::(&forester_epoch_pda_address) + .await + .unwrap(); + + assert!( + forester_epoch_pda.is_some(), + "Forester should be registered" + ); + let forester_epoch_pda = forester_epoch_pda.unwrap(); + assert_eq!( + forester_epoch_pda.epoch, current_epoch, + "Registered epoch should match current epoch" + ); +} + +pub async fn assert_queue_len( + pool: &SolanaRpcPool, + state_trees: &[StateMerkleTreeAccounts], + address_trees: &[AddressMerkleTreeAccounts], + total_expected_work: &mut u64, + expected_len: usize, + not_empty: bool, +) { + for tree in state_trees.iter() { + let mut rpc = pool.get_connection().await.unwrap(); + let queue_length = fetch_queue_item_data( + &mut *rpc, + &tree.nullifier_queue, + 0, + STATE_NULLIFIER_QUEUE_VALUES, + STATE_NULLIFIER_QUEUE_VALUES, + ) + .await + .unwrap() + .len(); + if not_empty { + assert_ne!(queue_length, 0); + } else { + assert_eq!(queue_length, expected_len); + } + *total_expected_work += queue_length as u64; + } + + for tree in address_trees.iter() { + let mut rpc = pool.get_connection().await.unwrap(); + let queue_length = fetch_queue_item_data( + &mut *rpc, + &tree.queue, + 0, + ADDRESS_QUEUE_VALUES, + ADDRESS_QUEUE_VALUES, + ) + .await + .unwrap() + .len(); + if not_empty { + assert_ne!(queue_length, 0); + } else { + assert_eq!(queue_length, expected_len); + } + *total_expected_work += queue_length as u64; + } +} diff --git a/forester/tests/e2e_v2_test.rs b/forester/tests/e2e_v2_test.rs new file mode 100644 index 0000000000..f60205f0ca --- /dev/null +++ b/forester/tests/e2e_v2_test.rs @@ -0,0 +1,1384 @@ +use std::{collections::HashMap, env, sync::Arc, time::Duration}; + +use account_compression::{state::StateMerkleTreeAccount, AddressMerkleTreeAccount}; +use anchor_lang::Discriminator; +use borsh::BorshSerialize; +use create_address_test_program::create_invoke_cpi_instruction; +use forester::{ + config::{ExternalServicesConfig, GeneralConfig, RpcPoolConfig, TransactionConfig}, + epoch_manager::WorkReport, + run_pipeline, + utils::get_protocol_config, + ForesterConfig, +}; +use forester_utils::utils::wait_for_indexer; +use light_batched_merkle_tree::{ + initialize_state_tree::InitStateTreeAccountsInstructionData, + merkle_tree::BatchedMerkleTreeAccount, +}; +use light_client::{ + indexer::{ + photon_indexer::PhotonIndexer, AddressWithTree, + GetCompressedTokenAccountsByOwnerOrDelegateOptions, Indexer, + }, + local_test_validator::{LightValidatorConfig, ProverConfig}, + rpc::{merkle_tree::MerkleTreeExt, LightClient, LightClientConfig, Rpc}, +}; +use light_compressed_account::{ + address::{derive_address, derive_address_legacy}, + compressed_account::CompressedAccount, + instruction_data::{ + compressed_proof::CompressedProof, + data::{NewAddressParams, NewAddressParamsAssigned}, + with_readonly::InstructionDataInvokeCpiWithReadOnly, + }, + TreeType, +}; +use light_compressed_token::process_transfer::{ + transfer_sdk::{create_transfer_instruction, to_account_metas}, + TokenTransferOutputData, +}; +use light_hasher::Poseidon; +use light_program_test::accounts::test_accounts::TestAccounts; +use light_prover_client::prover::spawn_prover; +use light_sdk::token::TokenDataWithMerkleContext; +use light_test_utils::{ + conversions::sdk_to_program_token_data, get_concurrent_merkle_tree, get_indexed_merkle_tree, + pack::pack_new_address_params_assigned, spl::create_mint_helper_with_keypair, + system_program::create_invoke_instruction, +}; +use rand::{prelude::SliceRandom, rngs::StdRng, Rng, SeedableRng}; +use serial_test::serial; +use solana_program::{native_token::LAMPORTS_PER_SOL, pubkey::Pubkey}; +use solana_sdk::{ + signature::{Keypair, Signature}, + signer::Signer, +}; +use tokio::{ + sync::{mpsc, oneshot, Mutex}, + time::{sleep, timeout}, +}; + +use crate::test_utils::{get_registration_phase_start_slot, init, wait_for_slot}; + +mod test_utils; + +const MINT_TO_NUM: u64 = 5; +const DEFAULT_TIMEOUT_SECONDS: u64 = 60 * 15; +const COMPUTE_BUDGET_LIMIT: u32 = 1_000_000; + +#[derive(Debug, Clone, Copy, PartialEq)] +enum TestMode { + Local, + Devnet, +} + +impl TestMode { + fn from_env() -> Self { + match env::var("TEST_MODE").as_deref() { + Ok("local") => TestMode::Local, + Ok("devnet") => TestMode::Devnet, + _ => TestMode::Devnet, // Default to devnet + } + } +} + +fn get_env_var(key: &str) -> String { + env::var(key).unwrap_or_else(|_| panic!("{} environment variable is not set", key)) +} + +fn get_rpc_url() -> String { + match TestMode::from_env() { + TestMode::Local => "http://localhost:8899".to_string(), + TestMode::Devnet => get_env_var("PHOTON_RPC_URL"), + } +} + +fn get_ws_rpc_url() -> String { + match TestMode::from_env() { + TestMode::Local => "ws://localhost:8900".to_string(), + TestMode::Devnet => get_env_var("PHOTON_WSS_RPC_URL"), + } +} + +fn get_indexer_url() -> String { + match TestMode::from_env() { + TestMode::Local => "http://localhost:8784".to_string(), + TestMode::Devnet => get_env_var("PHOTON_INDEXER_URL"), + } +} + +fn get_prover_url() -> String { + match TestMode::from_env() { + TestMode::Local => "http://localhost:3001".to_string(), + TestMode::Devnet => get_env_var("PHOTON_PROVER_URL"), + } +} + +fn get_api_key() -> Option { + match TestMode::from_env() { + TestMode::Local => None, + TestMode::Devnet => Some(get_env_var("PHOTON_API_KEY")), + } +} + +fn get_forester_keypair() -> Keypair { + match TestMode::from_env() { + TestMode::Local => Keypair::new(), + TestMode::Devnet => { + let keypair_string = get_env_var("FORESTER_KEYPAIR"); + + if keypair_string.starts_with('[') && keypair_string.ends_with(']') { + let bytes_str = &keypair_string[1..keypair_string.len() - 1]; // Remove [ ] + let bytes: Result, _> = bytes_str + .split(',') + .map(|s| s.trim().parse::()) + .collect(); + + match bytes { + Ok(byte_vec) => { + if byte_vec.len() == 64 { + return Keypair::from_bytes(&byte_vec) + .expect("Failed to create keypair from byte array"); + } else { + panic!( + "Keypair byte array must be exactly 64 bytes, got {}", + byte_vec.len() + ); + } + } + Err(e) => panic!("Failed to parse keypair byte array: {}", e), + } + } + + match bs58::decode(&keypair_string).into_vec() { + Ok(bytes) => { + Keypair::from_bytes(&bytes).expect("Failed to create keypair from base58 bytes") + } + Err(_) => panic!( + "FORESTER_KEYPAIR must be either base58 encoded or byte array format [1,2,3,...]" + ), + } + } + } +} + +fn is_v1_state_test_enabled() -> bool { + env::var("TEST_V1_STATE").unwrap_or_else(|_| "true".to_string()) == "true" +} + +fn is_v2_state_test_enabled() -> bool { + env::var("TEST_V2_STATE").unwrap_or_else(|_| "true".to_string()) == "true" +} + +fn is_v1_address_test_enabled() -> bool { + env::var("TEST_V1_ADDRESS").unwrap_or_else(|_| "true".to_string()) == "true" +} + +fn is_v2_address_test_enabled() -> bool { + env::var("TEST_V2_ADDRESS").unwrap_or_else(|_| "true".to_string()) == "true" +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 16)] +#[serial] +async fn test_e2e_v2() { + let state_tree_params = InitStateTreeAccountsInstructionData::test_default(); + let env = TestAccounts::get_local_test_validator_accounts(); + let config = ForesterConfig { + external_services: ExternalServicesConfig { + rpc_url: get_rpc_url(), + ws_rpc_url: Some(get_ws_rpc_url()), + indexer_url: Some(get_indexer_url()), + prover_url: Some(get_prover_url()), + photon_api_key: get_api_key(), + pushgateway_url: None, + pagerduty_routing_key: None, + rpc_rate_limit: None, + photon_rate_limit: None, + send_tx_rate_limit: None, + }, + retry_config: Default::default(), + queue_config: Default::default(), + indexer_config: Default::default(), + transaction_config: TransactionConfig { + batch_ixs_per_tx: 4, + ..Default::default() + }, + general_config: GeneralConfig { + slot_update_interval_seconds: 10, + tree_discovery_interval_seconds: 5, + enable_metrics: false, + skip_v1_state_trees: false, + skip_v2_state_trees: false, + skip_v1_address_trees: false, + skip_v2_address_trees: false, + }, + rpc_pool_config: RpcPoolConfig { + max_size: 50, + connection_timeout_secs: 15, + idle_timeout_secs: 300, + max_retries: 10, + initial_retry_delay_ms: 1000, + max_retry_delay_ms: 16000, + }, + registry_pubkey: light_registry::ID, + payer_keypair: env.protocol.forester.insecure_clone(), + derivation_pubkey: env.protocol.forester.pubkey(), + address_tree_data: vec![], + state_tree_data: vec![], + }; + let test_mode = TestMode::from_env(); + + if test_mode == TestMode::Local { + init(Some(LightValidatorConfig { + enable_indexer: true, + wait_time: 60, + prover_config: None, + sbf_programs: vec![( + "FNt7byTHev1k5x2cXZLBr8TdWiC3zoP5vcnZR4P682Uy".to_string(), + "../target/deploy/create_address_test_program.so".to_string(), + )], + limit_ledger_size: None, + })) + .await; + spawn_prover(ProverConfig::default()).await; + } + + let mut rpc = setup_rpc_connection(&env.protocol.forester).await; + if test_mode == TestMode::Local { + ensure_sufficient_balance( + &mut rpc, + &env.protocol.forester.pubkey(), + LAMPORTS_PER_SOL * 100, + ) + .await; + ensure_sufficient_balance( + &mut rpc, + &env.protocol.governance_authority.pubkey(), + LAMPORTS_PER_SOL * 100, + ) + .await; + } + + // Get initial state for V1 state tree if enabled + let pre_state_v1_root = if is_v1_state_test_enabled() { + let (_, _, root) = get_initial_merkle_tree_state( + &mut rpc, + &env.v1_state_trees[0].merkle_tree, + TreeType::StateV1, + ) + .await; + Some(root) + } else { + None + }; + + // Get initial state for V1 address tree if enabled + let pre_address_v1_root = if is_v1_address_test_enabled() { + let (_, _, root) = get_initial_merkle_tree_state( + &mut rpc, + &env.v1_address_trees[0].merkle_tree, + TreeType::AddressV1, + ) + .await; + Some(root) + } else { + None + }; + + // Get initial state for V2 state tree if enabled + let pre_state_v2_root = if is_v2_state_test_enabled() { + let (_, _, root) = get_initial_merkle_tree_state( + &mut rpc, + &env.v2_state_trees[0].merkle_tree, + TreeType::StateV2, + ) + .await; + Some(root) + } else { + None + }; + + // Get initial state for V2 address tree if enabled + let pre_address_v2_root = if is_v2_address_test_enabled() { + let (_, _, root) = + get_initial_merkle_tree_state(&mut rpc, &env.v2_address_trees[0], TreeType::AddressV2) + .await; + Some(root) + } else { + None + }; + + let payer = get_forester_keypair(); + println!("payer pubkey: {:?}", payer.pubkey()); + + if test_mode == TestMode::Local { + ensure_sufficient_balance(&mut rpc, &payer.pubkey(), LAMPORTS_PER_SOL * 100).await; + } else { + ensure_sufficient_balance(&mut rpc, &payer.pubkey(), LAMPORTS_PER_SOL).await; + } + + // V1 mint if V1 test enabled + let legacy_mint_pubkey = if is_v1_state_test_enabled() { + let legacy_mint_keypair = Keypair::new(); + let pubkey = create_mint_helper_with_keypair(&mut rpc, &payer, &legacy_mint_keypair).await; + + let sig = mint_to( + &mut rpc, + &env.v1_state_trees[0].merkle_tree, + &payer, + &pubkey, + ) + .await; + println!("v1 mint_to: {:?}", sig); + Some(pubkey) + } else { + println!("Skipping V1 mint - V1 state test disabled"); + None + }; + + // V2 mint if V2 test enabled + let batch_mint_pubkey = if is_v2_state_test_enabled() { + let batch_mint_keypair = Keypair::new(); + let pubkey = create_mint_helper_with_keypair(&mut rpc, &payer, &batch_mint_keypair).await; + + let sig = mint_to( + &mut rpc, + &env.v2_state_trees[0].output_queue, + &payer, + &pubkey, + ) + .await; + println!("v2 mint_to: {:?}", sig); + Some(pubkey) + } else { + println!("Skipping V2 mint - V2 state test disabled"); + None + }; + + let mut sender_batched_accs_counter = 0; + let mut sender_legacy_accs_counter = 0; + let mut sender_batched_token_counter: u64 = MINT_TO_NUM * 2; + let mut address_v1_counter = 0; + let mut address_v2_counter = 0; + + let rng_seed = rand::thread_rng().gen::(); + println!("seed {}", rng_seed); + let rng = &mut StdRng::seed_from_u64(rng_seed); + + let mut photon_indexer = create_photon_indexer(); + let protocol_config = get_protocol_config(&mut rpc).await; + + // spawn foresters on registration phase start slot + let registration_phase_slot = + get_registration_phase_start_slot(&mut rpc, &protocol_config).await; + wait_for_slot(&mut rpc, registration_phase_slot).await; + + let (service_handle, shutdown_sender, mut work_report_receiver) = + setup_forester_pipeline(&config).await; + + // execute transactions after forester is ready + let active_phase_slot = get_registration_phase_start_slot(&mut rpc, &protocol_config).await; + wait_for_slot(&mut rpc, active_phase_slot).await; + + execute_test_transactions( + &mut rpc, + &mut photon_indexer, + rng, + &env, + &payer, + legacy_mint_pubkey.as_ref(), + batch_mint_pubkey.as_ref(), + &mut sender_batched_accs_counter, + &mut sender_legacy_accs_counter, + &mut sender_batched_token_counter, + &mut address_v1_counter, + &mut address_v2_counter, + ) + .await; + + wait_for_work_report(&mut work_report_receiver, &state_tree_params).await; + + // Verify root changes based on enabled tests + if is_v1_state_test_enabled() { + if let Some(pre_root) = pre_state_v1_root { + verify_root_changed( + &mut rpc, + &env.v1_state_trees[0].merkle_tree, + &pre_root, + TreeType::StateV1, + ) + .await; + } + } + + if is_v2_state_test_enabled() { + if let Some(pre_root) = pre_state_v2_root { + verify_root_changed( + &mut rpc, + &env.v2_state_trees[0].merkle_tree, + &pre_root, + TreeType::StateV2, + ) + .await; + } + } + + if is_v1_address_test_enabled() { + if let Some(pre_root) = pre_address_v1_root { + verify_root_changed( + &mut rpc, + &env.v1_address_trees[0].merkle_tree, + &pre_root, + TreeType::AddressV1, + ) + .await; + } + } + + if is_v2_address_test_enabled() { + if let Some(pre_root) = pre_address_v2_root { + verify_root_changed( + &mut rpc, + &env.v2_address_trees[0], + &pre_root, + TreeType::AddressV2, + ) + .await; + } + } + + shutdown_sender + .send(()) + .expect("Failed to send shutdown signal"); + service_handle.await.unwrap().unwrap(); +} + +async fn setup_rpc_connection(forester: &Keypair) -> LightClient { + let mut rpc = LightClient::new(if TestMode::from_env() == TestMode::Local { + LightClientConfig::local() + } else { + LightClientConfig::new(get_rpc_url(), Some(get_indexer_url()), get_api_key()) + }) + .await + .unwrap(); + rpc.payer = forester.insecure_clone(); + rpc +} + +async fn ensure_sufficient_balance(rpc: &mut LightClient, pubkey: &Pubkey, target_balance: u64) { + if rpc.get_balance(pubkey).await.unwrap() < target_balance { + rpc.airdrop_lamports(pubkey, target_balance).await.unwrap(); + } +} + +fn create_photon_indexer() -> PhotonIndexer { + PhotonIndexer::new(get_indexer_url(), get_api_key()) +} + +async fn get_initial_merkle_tree_state( + rpc: &mut LightClient, + merkle_tree_pubkey: &Pubkey, + kind: TreeType, +) -> (u64, u64, [u8; 32]) { + match kind { + TreeType::StateV1 => { + let account = rpc + .get_anchor_account::(merkle_tree_pubkey) + .await + .unwrap() + .unwrap(); + + let merkle_tree = + get_concurrent_merkle_tree::( + rpc, + *merkle_tree_pubkey, + ) + .await; + + let next_index = merkle_tree.next_index() as u64; + let sequence_number = account.metadata.rollover_metadata.rolledover_slot; + let root = merkle_tree.root(); + + (next_index, sequence_number, root) + } + TreeType::AddressV1 => { + let account = rpc + .get_anchor_account::(merkle_tree_pubkey) + .await + .unwrap() + .unwrap(); + + let merkle_tree = get_indexed_merkle_tree::< + AddressMerkleTreeAccount, + LightClient, + Poseidon, + usize, + 26, + 16, + >(rpc, *merkle_tree_pubkey) + .await; + + let next_index = merkle_tree.next_index() as u64; + let sequence_number = account.metadata.rollover_metadata.rolledover_slot; + let root = merkle_tree.root(); + + (next_index, sequence_number, root) + } + TreeType::StateV2 => { + let mut merkle_tree_account = + rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap(); + let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( + merkle_tree_account.data.as_mut_slice(), + &merkle_tree_pubkey.into(), + ) + .unwrap(); + + let initial_next_index = merkle_tree.get_metadata().next_index; + let initial_sequence_number = merkle_tree.get_metadata().sequence_number; + ( + initial_next_index, + initial_sequence_number, + merkle_tree.get_root().unwrap(), + ) + } + TreeType::AddressV2 => { + let mut merkle_tree_account = + rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap(); + let merkle_tree = BatchedMerkleTreeAccount::address_from_bytes( + merkle_tree_account.data.as_mut_slice(), + &merkle_tree_pubkey.into(), + ) + .unwrap(); + + let initial_next_index = merkle_tree.get_metadata().next_index; + let initial_sequence_number = merkle_tree.get_metadata().sequence_number; + ( + initial_next_index, + initial_sequence_number, + merkle_tree.get_root().unwrap(), + ) + } + } +} + +async fn verify_root_changed( + rpc: &mut LightClient, + merkle_tree_pubkey: &Pubkey, + pre_root: &[u8; 32], + kind: TreeType, +) { + let current_root = match kind { + TreeType::StateV1 => { + let merkle_tree = + get_concurrent_merkle_tree::( + rpc, + *merkle_tree_pubkey, + ) + .await; + + println!( + "Final V1 state tree next_index: {}", + merkle_tree.next_index() + ); + merkle_tree.root() + } + TreeType::AddressV1 => { + let merkle_tree = get_indexed_merkle_tree::< + AddressMerkleTreeAccount, + LightClient, + Poseidon, + usize, + 26, + 16, + >(rpc, *merkle_tree_pubkey) + .await; + + println!( + "Final V1 address tree next_index: {}", + merkle_tree.next_index() + ); + merkle_tree.root() + } + TreeType::StateV2 => { + let mut merkle_tree_account = + rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap(); + let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( + merkle_tree_account.data.as_mut_slice(), + &merkle_tree_pubkey.into(), + ) + .unwrap(); + + println!( + "Final V2 state tree metadata: {:?}", + merkle_tree.get_metadata() + ); + merkle_tree.get_root().unwrap() + } + TreeType::AddressV2 => { + let mut merkle_tree_account = + rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap(); + let merkle_tree = BatchedMerkleTreeAccount::address_from_bytes( + merkle_tree_account.data.as_mut_slice(), + &merkle_tree_pubkey.into(), + ) + .unwrap(); + + println!( + "Final V2 address tree metadata: {:?}", + merkle_tree.get_metadata() + ); + merkle_tree.get_root().unwrap() + } + }; + + assert_ne!( + *pre_root, current_root, + "Root should have changed for {:?}", + kind + ); +} + +async fn get_state_v2_batch_size(rpc: &mut R, merkle_tree_pubkey: &Pubkey) -> u64 { + let mut merkle_tree_account = rpc.get_account(*merkle_tree_pubkey).await.unwrap().unwrap(); + let merkle_tree = BatchedMerkleTreeAccount::state_from_bytes( + merkle_tree_account.data.as_mut_slice(), + &merkle_tree_pubkey.into(), + ) + .unwrap(); + + merkle_tree.get_metadata().queue_batches.batch_size +} + +async fn setup_forester_pipeline( + config: &ForesterConfig, +) -> ( + tokio::task::JoinHandle>, + oneshot::Sender<()>, + mpsc::Receiver, +) { + let (shutdown_sender, shutdown_receiver) = oneshot::channel(); + let (work_report_sender, work_report_receiver) = mpsc::channel(100); + + let forester_photon_indexer = create_photon_indexer(); + let service_handle = tokio::spawn(run_pipeline::( + Arc::from(config.clone()), + None, + None, + Arc::new(Mutex::new(forester_photon_indexer)), + shutdown_receiver, + work_report_sender, + )); + + (service_handle, shutdown_sender, work_report_receiver) +} + +async fn wait_for_work_report( + work_report_receiver: &mut mpsc::Receiver, + tree_params: &InitStateTreeAccountsInstructionData, +) { + let batch_size = tree_params.output_queue_zkp_batch_size as usize; + // With increased test size, expect more processed items + let minimum_processed_items: usize = if is_v2_state_test_enabled() { + (tree_params.output_queue_batch_size as usize) * 4 // Expect at least 4 batches worth + } else { + tree_params.output_queue_batch_size as usize + }; + let mut total_processed_items: usize = 0; + let timeout_duration = Duration::from_secs(DEFAULT_TIMEOUT_SECONDS); + + println!("Waiting for work reports..."); + println!("Batch size: {}", batch_size); + println!( + "Minimum required processed items: {}", + minimum_processed_items + ); + + let start_time = tokio::time::Instant::now(); + while total_processed_items < minimum_processed_items { + match timeout( + timeout_duration.saturating_sub(start_time.elapsed()), + work_report_receiver.recv(), + ) + .await + { + Ok(Some(report)) => { + println!("Received work report: {:?}", report); + total_processed_items += report.processed_items; + } + Ok(None) => { + println!("Work report channel closed unexpectedly"); + break; + } + Err(_) => { + println!("Timed out after waiting for {:?}", timeout_duration); + break; + } + } + } + + println!("Total processed items: {}", total_processed_items); + assert!( + total_processed_items >= minimum_processed_items, + "Processed fewer items ({}) than required ({})", + total_processed_items, + minimum_processed_items + ); +} + +#[allow(clippy::too_many_arguments)] +async fn execute_test_transactions( + rpc: &mut R, + indexer: &mut I, + rng: &mut StdRng, + env: &TestAccounts, + payer: &Keypair, + v1_mint_pubkey: Option<&Pubkey>, + v2_mint_pubkey: Option<&Pubkey>, + sender_batched_accs_counter: &mut u64, + sender_legacy_accs_counter: &mut u64, + sender_batched_token_counter: &mut u64, + address_v1_counter: &mut u64, + address_v2_counter: &mut u64, +) { + let mut iterations = 10; + if is_v2_state_test_enabled() { + let batch_size = + get_state_v2_batch_size(rpc, &env.v2_state_trees[0].merkle_tree).await as usize; + iterations = batch_size * 2; + } + + println!("Executing {} test transactions", iterations); + println!("==========================================="); + for i in 0..iterations { + if is_v2_state_test_enabled() { + let batch_compress_sig = compress( + rpc, + &env.v2_state_trees[0].output_queue, + payer, + if i == 0 { 5_000_000 } else { 2_000_000 }, // Ensure sufficient for rent exemption + sender_batched_accs_counter, + ) + .await; + println!("{} v2 compress: {:?}", i, batch_compress_sig); + + let batch_transfer_sig = transfer::( + rpc, + indexer, + &env.v2_state_trees[0].output_queue, + payer, + sender_batched_accs_counter, + env, + ) + .await; + println!("{} v2 transfer: {:?}", i, batch_transfer_sig); + + if let Some(mint_pubkey) = v2_mint_pubkey { + let batch_transfer_token_sig = compressed_token_transfer( + rpc, + indexer, + &env.v2_state_trees[0].output_queue, + payer, + mint_pubkey, + sender_batched_token_counter, + ) + .await; + println!("{} v2 token transfer: {:?}", i, batch_transfer_token_sig); + } + } + + if is_v1_state_test_enabled() { + let compress_sig = compress( + rpc, + &env.v1_state_trees[0].merkle_tree, + payer, + 2_000_000, // Ensure sufficient for rent exemption + sender_legacy_accs_counter, + ) + .await; + println!("{} v1 compress: {:?}", i, compress_sig); + + let legacy_transfer_sig = transfer::( + rpc, + indexer, + &env.v1_state_trees[0].merkle_tree, + payer, + sender_legacy_accs_counter, + env, + ) + .await; + println!("{} v1 transfer: {:?}", i, legacy_transfer_sig); + + if let Some(mint_pubkey) = v1_mint_pubkey { + let legacy_transfer_token_sig = compressed_token_transfer( + rpc, + indexer, + &env.v1_state_trees[0].merkle_tree, + payer, + mint_pubkey, + sender_batched_token_counter, + ) + .await; + println!("{} v1 token transfer: {:?}", i, legacy_transfer_token_sig); + } + } + + // V1 Address operations + if is_v1_address_test_enabled() { + let sig_v1_addr = create_v1_address( + rpc, + indexer, + rng, + &env.v1_address_trees[0].merkle_tree, + &env.v1_address_trees[0].queue, + payer, + address_v1_counter, + ) + .await; + println!("{} v1 address: {:?}", i, sig_v1_addr); + } + + // V2 Address operations + if is_v2_address_test_enabled() { + let sig_v2_addr = create_v2_addresses( + rpc, + &env.v2_address_trees[0], + &env.protocol.registered_program_pda, + payer, + env, + rng, + 2, + ) + .await; + + if sig_v2_addr.is_ok() { + *address_v2_counter += 2; + } + println!("{} v2 address create: {:?}", i, sig_v2_addr); + } + } +} + +async fn mint_to( + rpc: &mut R, + merkle_tree_pubkey: &Pubkey, + payer: &Keypair, + mint_pubkey: &Pubkey, +) -> Signature { + let mint_to_ix = light_compressed_token::process_mint::mint_sdk::create_mint_to_instruction( + &payer.pubkey(), + &payer.pubkey(), + mint_pubkey, + merkle_tree_pubkey, + vec![100_000; MINT_TO_NUM as usize], + vec![payer.pubkey(); MINT_TO_NUM as usize], + None, + false, + 0, + ); + let instructions = vec![ + solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( + COMPUTE_BUDGET_LIMIT, + ), + mint_to_ix, + ]; + rpc.create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await + .unwrap() +} + +async fn compressed_token_transfer( + rpc: &mut R, + indexer: &I, + merkle_tree_pubkey: &Pubkey, + payer: &Keypair, + mint: &Pubkey, + counter: &mut u64, +) -> Signature { + wait_for_indexer(rpc, indexer).await.unwrap(); + let mut input_compressed_accounts: Vec = indexer + .get_compressed_token_accounts_by_owner( + &payer.pubkey(), + Some(GetCompressedTokenAccountsByOwnerOrDelegateOptions { + mint: Some(*mint), + cursor: None, + limit: None, + }), + None, + ) + .await + .unwrap() + .into(); + if input_compressed_accounts.is_empty() { + return Signature::default(); + } + + let rng = &mut rand::thread_rng(); + input_compressed_accounts.shuffle(rng); + input_compressed_accounts.truncate(1); + + let tokens = input_compressed_accounts[0].token_data.amount; + let compressed_account_hashes = vec![input_compressed_accounts[0] + .compressed_account + .hash() + .unwrap()]; + + let proof_for_compressed_accounts = indexer + .get_validity_proof(compressed_account_hashes, vec![], None) + .await + .unwrap(); + + let root_indices = proof_for_compressed_accounts.value.get_root_indices(); + let merkle_contexts = vec![ + input_compressed_accounts[0] + .compressed_account + .merkle_context, + ]; + + let compressed_accounts = vec![TokenTransferOutputData { + amount: tokens, + owner: payer.pubkey(), + lamports: None, + merkle_tree: *merkle_tree_pubkey, + }]; + + let proof = proof_for_compressed_accounts + .value + .proof + .0 + .map(|p| CompressedProof { + a: p.a, + b: p.b, + c: p.c, + }); + let input_token_data = vec![sdk_to_program_token_data( + input_compressed_accounts[0].token_data.clone(), + )]; + let input_compressed_accounts_data = vec![input_compressed_accounts[0] + .compressed_account + .compressed_account + .clone()]; + + let instruction = create_transfer_instruction( + &payer.pubkey(), + &payer.pubkey(), + &merkle_contexts, + &compressed_accounts, + &root_indices, + &proof, + &input_token_data, + &input_compressed_accounts_data, + *mint, + None, + false, + None, + None, + None, + true, + None, + None, + false, + &[], + false, + ) + .unwrap(); + + let instructions = vec![ + solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( + COMPUTE_BUDGET_LIMIT, + ), + instruction, + ]; + let sig = rpc + .create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await + .unwrap(); + *counter += compressed_accounts.len() as u64; + *counter -= input_compressed_accounts.len() as u64; + sig +} + +async fn transfer( + rpc: &mut R, + indexer: &I, + merkle_tree_pubkey: &Pubkey, + payer: &Keypair, + counter: &mut u64, + test_accounts: &TestAccounts, +) -> Signature { + println!("transfer V2: {} merkle_tree: {}", V2, merkle_tree_pubkey); + wait_for_indexer(rpc, indexer).await.unwrap(); + let mut input_compressed_accounts = indexer + .get_compressed_accounts_by_owner(&payer.pubkey(), None, None) + .await + .map(|response| response.value.items) + .unwrap_or(vec![]); + + input_compressed_accounts = if V2 { + input_compressed_accounts + .into_iter() + .filter(|x| { + test_accounts + .v2_state_trees + .iter() + .any(|y| y.merkle_tree == x.tree_info.tree) + }) + .collect() + } else { + input_compressed_accounts + .into_iter() + .filter(|x| { + test_accounts + .v1_state_trees + .iter() + .any(|y| y.merkle_tree == x.tree_info.tree) + }) + .collect() + }; + + if input_compressed_accounts.is_empty() { + return Signature::default(); + } + + let rng = &mut rand::thread_rng(); + input_compressed_accounts.shuffle(rng); + input_compressed_accounts.truncate(1); + + let lamports = input_compressed_accounts[0].lamports; + let compressed_account_hashes = vec![input_compressed_accounts[0].hash]; + + wait_for_indexer(rpc, indexer).await.unwrap(); + let proof_for_compressed_accounts = indexer + .get_validity_proof(compressed_account_hashes, vec![], None) + .await + .unwrap(); + let root_indices = proof_for_compressed_accounts.value.get_root_indices(); + let merkle_contexts = vec![ + light_compressed_account::compressed_account::MerkleContext { + merkle_tree_pubkey: input_compressed_accounts[0].tree_info.tree.into(), + queue_pubkey: input_compressed_accounts[0].tree_info.queue.into(), + leaf_index: input_compressed_accounts[0].leaf_index, + prove_by_index: false, + tree_type: if V2 { + TreeType::StateV2 + } else { + TreeType::StateV1 + }, + }, + ]; + + let compressed_accounts = vec![CompressedAccount { + lamports, + owner: payer.pubkey().into(), + address: None, + data: None, + }]; + let proof = proof_for_compressed_accounts + .value + .proof + .0 + .map(|p| CompressedProof { + a: p.a, + b: p.b, + c: p.c, + }); + let input_compressed_accounts_data = vec![CompressedAccount { + lamports: input_compressed_accounts[0].lamports, + owner: input_compressed_accounts[0].owner.into(), + address: input_compressed_accounts[0].address, + data: input_compressed_accounts[0].data.clone(), + }]; + + let instruction = create_invoke_instruction( + &payer.pubkey(), + &payer.pubkey(), + &input_compressed_accounts_data, + &compressed_accounts, + &merkle_contexts, + &[*merkle_tree_pubkey], + &root_indices, + &[], + proof, + None, + false, + None, + true, + ); + + let instructions = vec![ + solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( + COMPUTE_BUDGET_LIMIT, + ), + instruction, + ]; + let sig = rpc + .create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await + .unwrap(); + *counter += compressed_accounts.len() as u64; + *counter -= input_compressed_accounts_data.len() as u64; + sig +} + +async fn compress( + rpc: &mut R, + merkle_tree_pubkey: &Pubkey, + payer: &Keypair, + lamports: u64, + counter: &mut u64, +) -> Signature { + let payer_balance = rpc.get_balance(&payer.pubkey()).await.unwrap(); + println!("payer balance: {}", payer_balance); + + // Ensure payer has enough balance for compress amount + transaction fees + rent exemption buffer + let rent_exemption_buffer = 50_000_000; // 0.05 SOL buffer for rent exemption (compression creates multiple accounts) + // Ensure the compress amount itself is sufficient for rent exemption + let min_rent_exempt = 2_000_000; // Minimum 0.002 SOL for rent exemption + let actual_lamports = std::cmp::max(lamports, min_rent_exempt); + + let required_balance = actual_lamports + rent_exemption_buffer; + + if payer_balance < required_balance { + // Try to airdrop more funds + let airdrop_amount = required_balance * 2; // Airdrop 2x what we need + println!( + "Insufficient balance. Requesting airdrop of {} lamports", + airdrop_amount + ); + if let Err(e) = rpc.airdrop_lamports(&payer.pubkey(), airdrop_amount).await { + println!("Airdrop failed: {:?}. Proceeding anyway...", e); + } else { + // Wait a bit for airdrop to process + sleep(Duration::from_millis(1000)).await; + let new_balance = rpc.get_balance(&payer.pubkey()).await.unwrap(); + println!("New payer balance after airdrop: {}", new_balance); + } + } + + let compress_account = CompressedAccount { + lamports: actual_lamports, + owner: payer.pubkey().into(), + address: None, + data: None, + }; + let instruction = create_invoke_instruction( + &payer.pubkey(), + &payer.pubkey(), + &[], + &[compress_account], + &[], + &[*merkle_tree_pubkey], + &[], + &[], + None, + Some(actual_lamports), + true, + None, + true, + ); + let instructions = vec![ + solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( + COMPUTE_BUDGET_LIMIT, + ), + instruction, + ]; + match rpc + .create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await + { + Ok(sig) => { + *counter += 1; + sig + } + Err(e) => { + panic!("compress error: {:?}", e); + } + } +} + +async fn create_v1_address( + rpc: &mut R, + indexer: &mut I, + rng: &mut StdRng, + merkle_tree_pubkey: &Pubkey, + queue: &Pubkey, + payer: &Keypair, + counter: &mut u64, +) -> Signature { + let seed = rng.gen::<[u8; 32]>(); + let address = derive_address_legacy( + &light_compressed_account::Pubkey::from(*merkle_tree_pubkey), + &seed, + ) + .unwrap(); + let address_proof_inputs = vec![AddressWithTree { + address, + tree: *merkle_tree_pubkey, + }]; + + wait_for_indexer(rpc, indexer).await.unwrap(); + let proof_for_addresses = indexer + .get_validity_proof(vec![], address_proof_inputs, None) + .await + .unwrap(); + + let new_address_params = vec![NewAddressParams { + seed, + address_queue_pubkey: (*queue).into(), + address_merkle_tree_pubkey: (*merkle_tree_pubkey).into(), + address_merkle_tree_root_index: proof_for_addresses.value.get_address_root_indices()[0], + }]; + + let proof = proof_for_addresses.value.proof.0.map(|p| CompressedProof { + a: p.a, + b: p.b, + c: p.c, + }); + let root = proof_for_addresses.value.addresses[0].root; + let index = proof_for_addresses.value.addresses[0].root_index; + + println!("indexer root: {:?}, index: {}", root, index); + + { + let account = rpc + .get_anchor_account::(merkle_tree_pubkey) + .await + .unwrap(); + println!("address merkle tree account: {:?}", account); + let merkle_tree = + get_indexed_merkle_tree::( + rpc, + *merkle_tree_pubkey, + ) + .await; + + for (idx, root) in merkle_tree.roots.iter().enumerate() { + println!("root[{}]: {:?}", idx, root); + } + println!("root index: {}", merkle_tree.root_index()); + } + + let instruction = create_invoke_instruction( + &payer.pubkey(), + &payer.pubkey(), + &[], + &[], + &[], + &[], + &proof_for_addresses.value.get_root_indices(), + &new_address_params, + proof, + None, + false, + None, + false, + ); + + let instructions = vec![ + solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( + COMPUTE_BUDGET_LIMIT, + ), + instruction, + ]; + let sig = rpc + .create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await + .unwrap(); + *counter += 1; + sig +} + +async fn create_v2_addresses( + rpc: &mut R, + batch_address_merkle_tree: &Pubkey, + _registered_program_pda: &Pubkey, + payer: &Keypair, + _env: &TestAccounts, + rng: &mut StdRng, + num_addresses: usize, +) -> Result<(), light_client::rpc::RpcError> { + let mut address_seeds = Vec::with_capacity(num_addresses); + let mut addresses = Vec::with_capacity(num_addresses); + + for _ in 0..num_addresses { + let seed = rng.gen(); + let address = derive_address( + &seed, + &batch_address_merkle_tree.to_bytes(), + &create_address_test_program::ID.to_bytes(), + ); + + address_seeds.push(seed); + addresses.push(address); + } + + let address_with_trees = addresses + .into_iter() + .map(|address| AddressWithTree { + address, + tree: *batch_address_merkle_tree, + }) + .collect::>(); + + let proof_result = rpc + .get_validity_proof(Vec::new(), address_with_trees, None) + .await + .unwrap(); + + let new_address_params = address_seeds + .iter() + .enumerate() + .map(|(i, seed)| NewAddressParamsAssigned { + seed: *seed, + address_queue_pubkey: (*batch_address_merkle_tree).into(), + address_merkle_tree_pubkey: (*batch_address_merkle_tree).into(), + address_merkle_tree_root_index: proof_result.value.get_address_root_indices()[i], + assigned_account_index: None, + }) + .collect::>(); + + let mut remaining_accounts = HashMap::::new(); + let packed_new_address_params = + pack_new_address_params_assigned(&new_address_params, &mut remaining_accounts); + + let ix_data = InstructionDataInvokeCpiWithReadOnly { + mode: 0, + bump: 255, + with_cpi_context: false, + invoking_program_id: create_address_test_program::ID.into(), + proof: proof_result.value.proof.0, + new_address_params: packed_new_address_params, + is_compress: false, + compress_or_decompress_lamports: 0, + output_compressed_accounts: Default::default(), + input_compressed_accounts: Default::default(), + with_transaction_hash: true, + read_only_accounts: Vec::new(), + read_only_addresses: Vec::new(), + cpi_context: Default::default(), + }; + + let remaining_accounts_metas = to_account_metas(remaining_accounts); + + let instruction = create_invoke_cpi_instruction( + payer.pubkey(), + [ + light_system_program::instruction::InvokeCpiWithReadOnly::DISCRIMINATOR.to_vec(), + ix_data.try_to_vec()?, + ] + .concat(), + remaining_accounts_metas, + None, + ); + + let instructions = vec![ + solana_sdk::compute_budget::ComputeBudgetInstruction::set_compute_unit_limit( + COMPUTE_BUDGET_LIMIT, + ), + instruction, + ]; + + rpc.create_and_send_transaction(&instructions, &payer.pubkey(), &[payer]) + .await + .map(|_| ()) +} diff --git a/forester/tests/batched_address_test.rs b/forester/tests/legacy/batched_address_test.rs similarity index 100% rename from forester/tests/batched_address_test.rs rename to forester/tests/legacy/batched_address_test.rs diff --git a/forester/tests/batched_state_async_indexer_test.rs b/forester/tests/legacy/batched_state_async_indexer_test.rs similarity index 99% rename from forester/tests/batched_state_async_indexer_test.rs rename to forester/tests/legacy/batched_state_async_indexer_test.rs index a661af39c1..af65022065 100644 --- a/forester/tests/batched_state_async_indexer_test.rs +++ b/forester/tests/legacy/batched_state_async_indexer_test.rs @@ -72,9 +72,7 @@ const COMPUTE_BUDGET_LIMIT: u32 = 1_000_000; // ``` // 2025-05-13T22:43:27.825147Z ERROR process_queue{forester=En9a97stB3Ek2n6Ey3NJwCUJnmTzLMMEA5C69upGDuQP epoch=0 tree=HLKs5NJ8FXkJg8BrzJt56adFYYuwg5etzDtBbQYTsixu}:process_light_slot{forester=En9a97stB3Ek2n6Ey3NJwCUJnmTzLMMEA5C69upGDuQP epoch=0 tree=HLKs5NJ8FXkJg8BrzJt56adFYYuwg5etzDtBbQYTsixu}:process_batched_operations{epoch=0 tree=HLKs5NJ8FXkJg8BrzJt56adFYYuwg5etzDtBbQYTsixu tree_type=StateV2}: forester::processor::v2::common: State append failed for tree HLKs5NJ8FXkJg8BrzJt56adFYYuwg5etzDtBbQYTsixu: InstructionData("prover error: \"Failed to send request: error sending request for url (http://localhost:3001/prove): error trying to connect: dns error: task 145 was cancelled\"") // ``` -#[ignore = "multiple flaky errors post light-client refactor"] #[tokio::test(flavor = "multi_thread", worker_threads = 16)] -#[serial] async fn test_state_indexer_async_batched() { let tree_params = InitStateTreeAccountsInstructionData::test_default(); @@ -179,7 +177,7 @@ async fn test_state_indexer_async_batched() { &env.v2_state_trees[0].output_queue, ) .await; - wait_for_indexer(&mut rpc, &photon_indexer).await.unwrap(); + wait_for_indexer(&rpc, &photon_indexer).await.unwrap(); let input_compressed_accounts = get_token_accounts::(&photon_indexer, &batch_payer.pubkey(), &mint_pubkey) diff --git a/forester/tests/batched_state_indexer_test.rs b/forester/tests/legacy/batched_state_indexer_test.rs similarity index 100% rename from forester/tests/batched_state_indexer_test.rs rename to forester/tests/legacy/batched_state_indexer_test.rs diff --git a/forester/tests/batched_state_test.rs b/forester/tests/legacy/batched_state_test.rs similarity index 100% rename from forester/tests/batched_state_test.rs rename to forester/tests/legacy/batched_state_test.rs diff --git a/forester/tests/e2e_test.rs b/forester/tests/legacy/e2e_test.rs similarity index 100% rename from forester/tests/e2e_test.rs rename to forester/tests/legacy/e2e_test.rs diff --git a/forester/tests/legacy/priority_fee_test.rs b/forester/tests/legacy/priority_fee_test.rs new file mode 100644 index 0000000000..5b17b90660 --- /dev/null +++ b/forester/tests/legacy/priority_fee_test.rs @@ -0,0 +1,226 @@ +use forester::{ + cli::StartArgs, + processor::v1::{ + config::CapConfig, + helpers::{get_capped_priority_fee, request_priority_fee_estimate}, + }, + ForesterConfig, +}; +use light_client::rpc::{LightClient, LightClientConfig, Rpc}; +use reqwest::Url; +use solana_sdk::signature::Signer; + +use crate::test_utils::init; +mod test_utils; + +#[tokio::test] +#[ignore] +async fn test_priority_fee_request() { + dotenvy::dotenv().ok(); + + init(None).await; + + let args = StartArgs { + rpc_url: Some( + std::env::var("FORESTER_RPC_URL").expect("FORESTER_RPC_URL must be set in environment"), + ), + push_gateway_url: None, + pagerduty_routing_key: None, + ws_rpc_url: Some( + std::env::var("FORESTER_WS_RPC_URL") + .expect("FORESTER_WS_RPC_URL must be set in environment"), + ), + indexer_url: Some( + std::env::var("FORESTER_INDEXER_URL") + .expect("FORESTER_INDEXER_URL must be set in environment"), + ), + prover_url: Some( + std::env::var("FORESTER_PROVER_URL") + .expect("FORESTER_PROVER_URL must be set in environment"), + ), + payer: Some( + std::env::var("FORESTER_PAYER").expect("FORESTER_PAYER must be set in environment"), + ), + derivation: Some( + std::env::var("FORESTER_DERIVATION_PUBKEY") + .expect("FORESTER_DERIVATION_PUBKEY must be set in environment"), + ), + photon_api_key: Some( + std::env::var("PHOTON_API_KEY").expect("PHOTON_API_KEY must be set in environment"), + ), + indexer_batch_size: 50, + indexer_max_concurrent_batches: 10, + legacy_ixs_per_tx: 1, + batch_ixs_per_tx: 4, + transaction_max_concurrent_batches: 20, + tx_cache_ttl_seconds: 15, + ops_cache_ttl_seconds: 180, + cu_limit: 1_000_000, + enable_priority_fees: true, + rpc_pool_size: 20, + rpc_pool_connection_timeout_secs: 1, + rpc_pool_idle_timeout_secs: 1, + rpc_pool_max_retries: 10, + rpc_pool_initial_retry_delay_ms: 1000, + rpc_pool_max_retry_delay_ms: 16000, + slot_update_interval_seconds: 10, + tree_discovery_interval_seconds: 5, + max_retries: 3, + retry_delay: 1000, + retry_timeout: 30000, + state_queue_start_index: 0, + state_queue_processing_length: 28807, + address_queue_start_index: 0, + address_queue_processing_length: 28807, + rpc_rate_limit: None, + photon_rate_limit: None, + send_tx_rate_limit: None, + }; + + let config = ForesterConfig::new_for_start(&args).expect("Failed to create config"); + + // Setup RPC connection using config + let mut rpc = LightClient::new(LightClientConfig::local()).await.unwrap(); + rpc.payer = config.payer_keypair.insecure_clone(); + + let account_keys = vec![config.payer_keypair.pubkey()]; + + let url = Url::parse(&rpc.get_url()).expect("Failed to parse URL"); + println!("URL: {}", url); + let priority_fee = request_priority_fee_estimate(&url, account_keys) + .await + .unwrap(); + + println!("Priority fee: {:?}", priority_fee); + assert!(priority_fee > 0, "Priority fee should be greater than 0"); +} +#[test] + +fn test_capped_priority_fee() { + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 50_000, + min_fee_lamports: 10_000, + max_fee_lamports: 100_000, + // 1_000_000 cu x 50_000 microlamports per cu = 50_000 lamports total + compute_unit_limit: 1_000_000, + }; + let expected = 50_000; + + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 10_000, + min_fee_lamports: 10_000, + max_fee_lamports: 100_000, + compute_unit_limit: 1_000_000, + }; + let expected = 10_000; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 100_000, + min_fee_lamports: 10_000, + max_fee_lamports: 100_000, + compute_unit_limit: 1_000_000, + }; + let expected = 100_000; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 10_000, + min_fee_lamports: 20_000, + max_fee_lamports: 100_000, + compute_unit_limit: 1_000_000, + }; + let expected = 20_000; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 200_000, + min_fee_lamports: 10_000, + max_fee_lamports: 100_000, + compute_unit_limit: 1_000_000, + }; + let expected = 100_000; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 10_000, + min_fee_lamports: 0, + max_fee_lamports: 0, + compute_unit_limit: 1_000_000, + }; + let expected = 0; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 10_000, + min_fee_lamports: 10_000, + max_fee_lamports: 0, + compute_unit_limit: 1_000_000, + }; + println!("expecting panic"); + let result = std::panic::catch_unwind(|| get_capped_priority_fee(cap_config)); + assert!( + result.is_err(), + "Expected panic for max fee less than min fee" + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 10_000, + min_fee_lamports: 50_000, + max_fee_lamports: 50_000, + compute_unit_limit: 1_000_000, + }; + let expected = 50_000; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); + + let cap_config = CapConfig { + rec_fee_microlamports_per_cu: 100_000, + min_fee_lamports: 50_000, + max_fee_lamports: 50_000, + compute_unit_limit: 1_000_000, + }; + let expected = 50_000; + let result = get_capped_priority_fee(cap_config); + assert_eq!( + result, expected, + "Priority fee capping failed for input {}", + cap_config.rec_fee_microlamports_per_cu + ); +} diff --git a/forester/tests/legacy/test_utils.rs b/forester/tests/legacy/test_utils.rs new file mode 100644 index 0000000000..a93a410b64 --- /dev/null +++ b/forester/tests/legacy/test_utils.rs @@ -0,0 +1,287 @@ +use forester::{ + config::{ExternalServicesConfig, GeneralConfig, RpcPoolConfig}, + metrics::register_metrics, + telemetry::setup_telemetry, + ForesterConfig, +}; +use light_client::{ + indexer::{photon_indexer::PhotonIndexer, Indexer, NewAddressProofWithContext}, + local_test_validator::{spawn_validator, LightValidatorConfig}, +}; +use light_program_test::{accounts::test_accounts::TestAccounts, indexer::TestIndexerExtensions}; +use light_test_utils::e2e_test_env::{GeneralActionConfig, KeypairActionConfig, User}; +use solana_sdk::{ + pubkey::Pubkey, + signature::{Keypair, Signer}, +}; +use tracing::debug; + +#[allow(dead_code)] +pub async fn init(config: Option) { + setup_telemetry(); + register_metrics(); + spawn_test_validator(config).await; +} + +#[allow(dead_code)] +pub async fn spawn_test_validator(config: Option) { + let config = config.unwrap_or_default(); + spawn_validator(config).await; +} + +#[allow(dead_code)] +pub fn keypair_action_config() -> KeypairActionConfig { + KeypairActionConfig { + compress_sol: Some(1.0), + decompress_sol: Some(1.0), + transfer_sol: Some(1.0), + create_address: Some(1.0), + compress_spl: Some(1.0), + decompress_spl: Some(1.0), + mint_spl: Some(1.0), + transfer_spl: Some(1.0), + max_output_accounts: Some(3), + fee_assert: false, + approve_spl: None, + revoke_spl: None, + freeze_spl: None, + thaw_spl: None, + burn_spl: None, + } +} + +#[allow(dead_code)] +pub fn general_action_config() -> GeneralActionConfig { + GeneralActionConfig { + add_keypair: Some(1.0), + create_state_mt: Some(1.0), + create_address_mt: Some(1.0), + nullify_compressed_accounts: Some(1.0), + empty_address_queue: Some(1.0), + rollover: None, + add_forester: None, + disable_epochs: true, + } +} + +#[allow(dead_code)] +pub fn forester_config() -> ForesterConfig { + let mut test_accounts = TestAccounts::get_program_test_test_accounts(); + test_accounts.protocol.forester = Keypair::new(); + + ForesterConfig { + external_services: ExternalServicesConfig { + rpc_url: "http://localhost:8899".to_string(), + ws_rpc_url: Some("ws://localhost:8900".to_string()), + indexer_url: Some("http://localhost:8784".to_string()), + prover_url: Some("http://localhost:3001".to_string()), + photon_api_key: None, + pushgateway_url: None, + pagerduty_routing_key: None, + rpc_rate_limit: None, + photon_rate_limit: None, + send_tx_rate_limit: None, + }, + retry_config: Default::default(), + queue_config: Default::default(), + indexer_config: Default::default(), + transaction_config: Default::default(), + general_config: GeneralConfig { + slot_update_interval_seconds: 10, + tree_discovery_interval_seconds: 5, + enable_metrics: false, + skip_v1_state_trees: false, + skip_v2_state_trees: false, + skip_v1_address_trees: false, + skip_v2_address_trees: false, + }, + rpc_pool_config: RpcPoolConfig { + max_size: 50, + connection_timeout_secs: 15, + idle_timeout_secs: 300, + max_retries: 10, + initial_retry_delay_ms: 1000, + max_retry_delay_ms: 16000, + }, + registry_pubkey: light_registry::ID, + payer_keypair: test_accounts.protocol.forester.insecure_clone(), + derivation_pubkey: test_accounts.protocol.forester.pubkey(), + address_tree_data: vec![], + state_tree_data: vec![], + } +} + +// truncate to <254 bit +#[allow(dead_code)] +pub fn generate_pubkey_254() -> Pubkey { + let mock_address: Pubkey = Pubkey::new_unique(); + let mut mock_address_less_than_254_bit: [u8; 32] = mock_address.to_bytes(); + mock_address_less_than_254_bit[0] = 0; + Pubkey::from(mock_address_less_than_254_bit) +} + +#[allow(dead_code)] +pub async fn assert_new_address_proofs_for_photon_and_test_indexer< + I: Indexer + TestIndexerExtensions, +>( + indexer: &mut I, + trees: &[Pubkey], + addresses: &[Pubkey], + photon_indexer: &PhotonIndexer, +) { + for (tree, address) in trees.iter().zip(addresses.iter()) { + let address_proof_test_indexer = indexer + .get_multiple_new_address_proofs(tree.to_bytes(), vec![address.to_bytes()], None) + .await; + + let address_proof_photon = photon_indexer + .get_multiple_new_address_proofs(tree.to_bytes(), vec![address.to_bytes()], None) + .await; + + if address_proof_photon.is_err() { + panic!("Photon error: {:?}", address_proof_photon); + } + + if address_proof_test_indexer.is_err() { + panic!("Test indexer error: {:?}", address_proof_test_indexer); + } + + let photon_result: NewAddressProofWithContext = address_proof_photon + .unwrap() + .value + .items + .first() + .unwrap() + .clone(); + let test_indexer_result: NewAddressProofWithContext = address_proof_test_indexer + .unwrap() + .value + .items + .first() + .unwrap() + .clone(); + debug!( + "assert proofs for address: {} photon result: {:?} test indexer result: {:?}", + address, photon_result, test_indexer_result + ); + + assert_eq!(photon_result.merkle_tree, test_indexer_result.merkle_tree); + assert_eq!( + photon_result.low_address_index, + test_indexer_result.low_address_index + ); + assert_eq!( + photon_result.low_address_value, + test_indexer_result.low_address_value + ); + assert_eq!( + photon_result.low_address_next_index, + test_indexer_result.low_address_next_index + ); + assert_eq!( + photon_result.low_address_next_value, + test_indexer_result.low_address_next_value + ); + assert_eq!( + photon_result.low_address_proof.len(), + test_indexer_result.low_address_proof.len() + ); + + assert_eq!(photon_result.root, test_indexer_result.root); + assert_eq!(photon_result.root_seq, test_indexer_result.root_seq); + + for (photon_proof_hash, test_indexer_proof_hash) in photon_result + .low_address_proof + .iter() + .zip(test_indexer_result.low_address_proof.iter()) + { + assert_eq!(photon_proof_hash, test_indexer_proof_hash); + } + } +} + +#[allow(dead_code)] +pub async fn assert_accounts_by_owner( + indexer: &mut I, + user: &User, + photon_indexer: &PhotonIndexer, +) { + let mut photon_accs = photon_indexer + .get_compressed_accounts_by_owner(&user.keypair.pubkey(), None, None) + .await + .unwrap() + .value + .items; + photon_accs.sort_by_key(|a| a.hash); + + let mut test_accs = indexer + .get_compressed_accounts_by_owner(&user.keypair.pubkey(), None, None) + .await + .unwrap(); + test_accs.value.items.sort_by_key(|a| a.hash); + + debug!( + "asserting accounts for user: {} Test accs: {:?} Photon accs: {:?}", + user.keypair.pubkey().to_string(), + test_accs.value.items.len(), + photon_accs.len() + ); + assert_eq!(test_accs.value.items.len(), photon_accs.len()); + + debug!("test_accs: {:?}", test_accs); + debug!("photon_accs: {:?}", photon_accs); + + for (test_acc, indexer_acc) in test_accs.value.items.iter().zip(photon_accs.iter()) { + assert_eq!(test_acc, indexer_acc); + } +} + +#[allow(dead_code)] +pub async fn assert_account_proofs_for_photon_and_test_indexer< + I: Indexer + TestIndexerExtensions, +>( + indexer: &mut I, + user_pubkey: &Pubkey, + photon_indexer: &PhotonIndexer, +) { + let accs = indexer + .get_compressed_accounts_by_owner(user_pubkey, None, None) + .await; + for account in accs.unwrap().value.items { + let photon_result = photon_indexer + .get_multiple_compressed_account_proofs(vec![account.hash], None) + .await; + let test_indexer_result = indexer + .get_multiple_compressed_account_proofs(vec![account.hash], None) + .await; + + if photon_result.is_err() { + panic!("Photon error: {:?}", photon_result); + } + + if test_indexer_result.is_err() { + panic!("Test indexer error: {:?}", test_indexer_result); + } + + let photon_result = photon_result.unwrap().value.items; + let test_indexer_result = test_indexer_result.unwrap().value.items; + + assert_eq!(photon_result.len(), test_indexer_result.len()); + for (photon_proof, test_indexer_proof) in + photon_result.iter().zip(test_indexer_result.iter()) + { + assert_eq!(photon_proof.hash, test_indexer_proof.hash); + assert_eq!(photon_proof.leaf_index, test_indexer_proof.leaf_index); + assert_eq!(photon_proof.merkle_tree, test_indexer_proof.merkle_tree); + assert_eq!(photon_proof.root_seq, test_indexer_proof.root_seq); + assert_eq!(photon_proof.proof.len(), test_indexer_proof.proof.len()); + for (photon_proof_hash, test_indexer_proof_hash) in photon_proof + .proof + .iter() + .zip(test_indexer_proof.proof.iter()) + { + assert_eq!(photon_proof_hash, test_indexer_proof_hash); + } + } + } +} diff --git a/forester/tests/test_utils.rs b/forester/tests/test_utils.rs index a93a410b64..8bf48c3a40 100644 --- a/forester/tests/test_utils.rs +++ b/forester/tests/test_utils.rs @@ -1,19 +1,28 @@ +use std::time::Duration; + use forester::{ config::{ExternalServicesConfig, GeneralConfig, RpcPoolConfig}, metrics::register_metrics, telemetry::setup_telemetry, ForesterConfig, }; +use forester_utils::forester_epoch::get_epoch_phases; use light_client::{ indexer::{photon_indexer::PhotonIndexer, Indexer, NewAddressProofWithContext}, local_test_validator::{spawn_validator, LightValidatorConfig}, + rpc::{LightClient, Rpc}, }; use light_program_test::{accounts::test_accounts::TestAccounts, indexer::TestIndexerExtensions}; +use light_registry::{ + protocol_config::state::{ProtocolConfig, ProtocolConfigPda}, + utils::get_protocol_config_pda_address, +}; use light_test_utils::e2e_test_env::{GeneralActionConfig, KeypairActionConfig, User}; use solana_sdk::{ pubkey::Pubkey, signature::{Keypair, Signer}, }; +use tokio::time::sleep; use tracing::debug; #[allow(dead_code)] @@ -285,3 +294,47 @@ pub async fn assert_account_proofs_for_photon_and_test_indexer< } } } + +#[allow(dead_code)] +pub async fn get_registration_phase_start_slot( + rpc: &mut R, + protocol_config: &ProtocolConfig, +) -> u64 { + let current_slot = rpc.get_slot().await.unwrap(); + let current_epoch = protocol_config.get_current_epoch(current_slot); + let phases = get_epoch_phases(protocol_config, current_epoch); + phases.registration.start +} + +#[allow(dead_code)] +pub async fn get_active_phase_start_slot( + rpc: &mut R, + protocol_config: &ProtocolConfig, +) -> u64 { + let current_slot = rpc.get_slot().await.unwrap(); + let current_epoch = protocol_config.get_current_epoch(current_slot); + let phases = get_epoch_phases(protocol_config, current_epoch); + phases.active.start +} + +#[allow(dead_code)] +pub async fn wait_for_slot(rpc: &mut LightClient, target_slot: u64) { + while rpc.get_slot().await.unwrap() < target_slot { + println!( + "waiting for active phase slot: {}, current slot: {}", + target_slot, + rpc.get_slot().await.unwrap() + ); + sleep(Duration::from_millis(400)).await; + } +} + +#[allow(dead_code)] +async fn get_protocol_config(rpc: &mut LightClient) -> ProtocolConfig { + let protocol_config_pda_address = get_protocol_config_pda_address().0; + rpc.get_anchor_account::(&protocol_config_pda_address) + .await + .unwrap() + .unwrap() + .config +} diff --git a/program-tests/utils/src/setup_accounts.rs b/program-tests/utils/src/setup_accounts.rs index c52fd86002..f63244fe7b 100644 --- a/program-tests/utils/src/setup_accounts.rs +++ b/program-tests/utils/src/setup_accounts.rs @@ -14,6 +14,7 @@ pub async fn setup_accounts(keypairs: TestKeypairs, url: RpcUrl) -> Result, pub photon_url: Option, + pub api_key: Option, pub fetch_active_tree: bool, } impl LightClientConfig { - pub fn new(url: String, photon_url: Option) -> Self { + pub fn new(url: String, photon_url: Option, api_key: Option) -> Self { Self { url, photon_url, + api_key, commitment_config: Some(CommitmentConfig::confirmed()), fetch_active_tree: true, } @@ -45,6 +47,7 @@ impl LightClientConfig { url: RpcUrl::Localnet.to_string(), commitment_config: Some(CommitmentConfig::confirmed()), photon_url: None, + api_key: None, fetch_active_tree: false, } } @@ -54,14 +57,16 @@ impl LightClientConfig { url: RpcUrl::Localnet.to_string(), commitment_config: Some(CommitmentConfig::confirmed()), photon_url: Some("http://127.0.0.1:8784".to_string()), + api_key: None, fetch_active_tree: false, } } - pub fn devnet(photon_url: Option) -> Self { + pub fn devnet(photon_url: Option, api_key: Option) -> Self { Self { url: RpcUrl::Devnet.to_string(), photon_url, + api_key, commitment_config: Some(CommitmentConfig::confirmed()), fetch_active_tree: true, } diff --git a/xtask/src/create_batch_address_tree.rs b/xtask/src/create_batch_address_tree.rs index ffb39ef368..1cde15c0c0 100644 --- a/xtask/src/create_batch_address_tree.rs +++ b/xtask/src/create_batch_address_tree.rs @@ -44,6 +44,7 @@ pub async fn create_batch_address_tree(options: Options) -> anyhow::Result<()> { photon_url: None, commitment_config: None, fetch_active_tree: false, + api_key: None, }) .await .unwrap(); diff --git a/xtask/src/create_batch_state_tree.rs b/xtask/src/create_batch_state_tree.rs index 0cf7810b71..7afb0b411b 100644 --- a/xtask/src/create_batch_state_tree.rs +++ b/xtask/src/create_batch_state_tree.rs @@ -49,6 +49,7 @@ pub async fn create_batch_state_tree(options: Options) -> anyhow::Result<()> { photon_url: None, commitment_config: None, fetch_active_tree: false, + api_key: None, }) .await .unwrap(); diff --git a/xtask/src/create_state_tree.rs b/xtask/src/create_state_tree.rs index e1b8a7b733..326c7f4d1c 100644 --- a/xtask/src/create_state_tree.rs +++ b/xtask/src/create_state_tree.rs @@ -49,6 +49,7 @@ pub async fn create_state_tree(options: Options) -> anyhow::Result<()> { photon_url: None, commitment_config: None, fetch_active_tree: false, + api_key: None, }) .await .unwrap(); diff --git a/xtask/src/create_update_protocol_config_ix.rs b/xtask/src/create_update_protocol_config_ix.rs index a2a136d541..e91ac4fbb4 100644 --- a/xtask/src/create_update_protocol_config_ix.rs +++ b/xtask/src/create_update_protocol_config_ix.rs @@ -37,6 +37,7 @@ pub async fn create_update_protocol_config_ix(options: Options) -> anyhow::Resul photon_url: None, commitment_config: None, fetch_active_tree: false, + api_key: None, }) .await .unwrap(); diff --git a/xtask/src/new_deployment.rs b/xtask/src/new_deployment.rs index 4cb7f6f5ff..ff9a61b7a4 100644 --- a/xtask/src/new_deployment.rs +++ b/xtask/src/new_deployment.rs @@ -58,6 +58,7 @@ pub async fn init_new_deployment(options: Options) -> anyhow::Result<()> { photon_url: None, commitment_config: None, fetch_active_tree: false, + api_key: None, }) .await .unwrap();