feat: optimize forester by parallelizing zkp proof generation for batch append and nullify (#1633)

* wip

* parallel batch proof gen

* Remove debug logs

* remove unnecessary init_logger calls from tests

* cleanup

* cleanup

* Add delays to stabilize async indexer batch tests

* test_indexer: get_queue_elements with 0 offset

* fix

* Group transactions into chunks for batch processing.

* Remove sleep delay in transaction processing.

* Refactor chunking logic with a constant for chunk size.

* add batch_ixs_per_tx configuration for transaction batching

Replaces the fixed CHUNK_SIZE with a configurable batch_ixs_per_tx parameter, allowing dynamic control over instructions per transaction.

* Change batch_ixs_per_tx to 1 in test configurations

* format

* switch to warp-ubuntu-latest-x64-4x for forester ci workflow

Author: sergeytimoshin

Cargo.lock

@@ -43,12 +43,14 @@ members = [
 
 resolver = "2"
 
+[profile.dev]
+opt-level = 0
+debug = true
+strip = "none"
+
 [profile.release]
 overflow-checks = true
 
-[profile.test]
-opt-level = 2
-
 [workspace.dependencies]
 # Solana
 solana-banks-client = "=1.18.22"

Cargo.toml

@@ -36,7 +36,7 @@ solana-sdk = { workspace = true }
 thiserror = { workspace = true }
 
 # Logging
-log = { workspace = true }
+tracing = { workspace = true }
 
 # Big numbers
 num-traits = { workspace = true }

forester-utils/Cargo.toml

@@ -0,0 +1,93 @@
+use std::{fmt, marker::PhantomData, mem, pin::Pin};
+
+use account_compression::processor::initialize_address_merkle_tree::Pubkey;
+use light_client::rpc::RpcConnection;
+use light_concurrent_merkle_tree::copy::ConcurrentMerkleTreeCopy;
+use light_hash_set::HashSet;
+use light_hasher::Hasher;
+use light_indexed_merkle_tree::copy::IndexedMerkleTreeCopy;
+use num_traits::{CheckedAdd, CheckedSub, ToBytes, Unsigned};
+use solana_sdk::account::Account;
+
+#[derive(Debug, Clone)]
+pub struct AccountZeroCopy<'a, T> {
+    pub account: Pin<Box<Account>>,
+    deserialized: *const T,
+    _phantom_data: PhantomData<&'a T>,
+}
+
+impl<'a, T> AccountZeroCopy<'a, T> {
+    pub async fn new<R: RpcConnection>(rpc: &mut R, address: Pubkey) -> AccountZeroCopy<'a, T> {
+        let account = Box::pin(rpc.get_account(address).await.unwrap().unwrap());
+        let deserialized = account.data[8..].as_ptr() as *const T;
+
+        Self {
+            account,
+            deserialized,
+            _phantom_data: PhantomData,
+        }
+    }
+
+    // Safe method to access `deserialized` ensuring the lifetime is respected
+    pub fn deserialized(&self) -> &'a T {
+        unsafe { &*self.deserialized }
+    }
+}
+
+/// Fetches the given account, then copies and serializes it as a `HashSet`.
+///
+/// # Safety
+///
+/// This is highly unsafe. Ensuring that:
+///
+/// * The correct account is used.
+/// * The account has enough space to be treated as a HashSet with specified
+///   parameters.
+/// * The account data is aligned.
+///
+/// Is the caller's responsibility.
+pub async unsafe fn get_hash_set<T, R: RpcConnection>(rpc: &mut R, pubkey: Pubkey) -> HashSet {
+    let mut data = rpc.get_account(pubkey).await.unwrap().unwrap().data.clone();
+
+    HashSet::from_bytes_copy(&mut data[8 + mem::size_of::<T>()..]).unwrap()
+}
+
+/// Fetches the given account, then copies and serializes it as a
+/// `ConcurrentMerkleTree`.
+pub async fn get_concurrent_merkle_tree<T, R, H, const HEIGHT: usize>(
+    rpc: &mut R,
+    pubkey: Pubkey,
+) -> ConcurrentMerkleTreeCopy<H, HEIGHT>
+where
+    R: RpcConnection,
+    H: Hasher,
+{
+    let account = rpc.get_account(pubkey).await.unwrap().unwrap();
+
+    ConcurrentMerkleTreeCopy::from_bytes_copy(&account.data[8 + mem::size_of::<T>()..]).unwrap()
+}
+// TODO: do discriminator check
+/// Fetches the fiven account, then copies and serializes it as an
+/// `IndexedMerkleTree`.
+pub async fn get_indexed_merkle_tree<T, R, H, I, const HEIGHT: usize, const NET_HEIGHT: usize>(
+    rpc: &mut R,
+    pubkey: Pubkey,
+) -> IndexedMerkleTreeCopy<H, I, HEIGHT, NET_HEIGHT>
+where
+    R: RpcConnection,
+    H: Hasher,
+    I: CheckedAdd
+        + CheckedSub
+        + Copy
+        + Clone
+        + fmt::Debug
+        + PartialOrd
+        + ToBytes
+        + TryFrom<usize>
+        + Unsigned,
+    usize: From<I>,
+{
+    let account = rpc.get_account(pubkey).await.unwrap().unwrap();
+
+    IndexedMerkleTreeCopy::from_bytes_copy(&account.data[8 + mem::size_of::<T>()..]).unwrap()
+}

forester-utils/src/account_zero_copy.rs

@@ -13,17 +13,18 @@ use light_hasher::Poseidon;
 use num_traits::Zero;
 use solana_sdk::pubkey::Pubkey;
 
-use crate::{get_concurrent_merkle_tree, get_hash_set, get_indexed_merkle_tree, AccountZeroCopy};
+use crate::account_zero_copy::{
+    get_concurrent_merkle_tree, get_hash_set, get_indexed_merkle_tree, AccountZeroCopy,
+};
 
 pub async fn get_address_bundle_config<R: RpcConnection>(
     rpc: &mut R,
     address_bundle: AddressMerkleTreeAccounts,
 ) -> (AddressMerkleTreeConfig, AddressQueueConfig) {
-    let address_queue_meta_data =
-        AccountZeroCopy::<account_compression::QueueAccount>::new(rpc, address_bundle.queue)
-            .await
-            .deserialized()
-            .metadata;
+    let address_queue_meta_data = AccountZeroCopy::<QueueAccount>::new(rpc, address_bundle.queue)
+        .await
+        .deserialized()
+        .metadata;
     let address_queue = unsafe { get_hash_set::<QueueAccount, R>(rpc, address_bundle.queue).await };
     let queue_config = AddressQueueConfig {
         network_fee: Some(address_queue_meta_data.rollover_metadata.network_fee),
@@ -32,13 +33,10 @@ pub async fn get_address_bundle_config<R: RpcConnection>(
         sequence_threshold: address_queue.sequence_threshold as u64,
     };
     let address_tree_meta_data =
-        AccountZeroCopy::<account_compression::AddressMerkleTreeAccount>::new(
-            rpc,
-            address_bundle.merkle_tree,
-        )
-        .await
-        .deserialized()
-        .metadata;
+        AccountZeroCopy::<AddressMerkleTreeAccount>::new(rpc, address_bundle.merkle_tree)
+            .await
+            .deserialized()
+            .metadata;
     let address_tree =
         get_indexed_merkle_tree::<AddressMerkleTreeAccount, R, Poseidon, usize, 26, 16>(
             rpc,
@@ -70,13 +68,11 @@ pub async fn get_state_bundle_config<R: RpcConnection>(
     rpc: &mut R,
     state_tree_bundle: StateMerkleTreeAccounts,
 ) -> (StateMerkleTreeConfig, NullifierQueueConfig) {
-    let address_queue_meta_data = AccountZeroCopy::<account_compression::QueueAccount>::new(
-        rpc,
-        state_tree_bundle.nullifier_queue,
-    )
-    .await
-    .deserialized()
-    .metadata;
+    let address_queue_meta_data =
+        AccountZeroCopy::<QueueAccount>::new(rpc, state_tree_bundle.nullifier_queue)
+            .await
+            .deserialized()
+            .metadata;
     let address_queue =
         unsafe { get_hash_set::<QueueAccount, R>(rpc, state_tree_bundle.nullifier_queue).await };
     let queue_config = NullifierQueueConfig {
@@ -85,13 +81,10 @@ pub async fn get_state_bundle_config<R: RpcConnection>(
         sequence_threshold: address_queue.sequence_threshold as u64,
     };
     let address_tree_meta_data =
-        AccountZeroCopy::<account_compression::StateMerkleTreeAccount>::new(
-            rpc,
-            state_tree_bundle.merkle_tree,
-        )
-        .await
-        .deserialized()
-        .metadata;
+        AccountZeroCopy::<StateMerkleTreeAccount>::new(rpc, state_tree_bundle.merkle_tree)
+            .await
+            .deserialized()
+            .metadata;
     let address_tree = get_concurrent_merkle_tree::<StateMerkleTreeAccount, R, Poseidon, 26>(
         rpc,
         state_tree_bundle.merkle_tree,
@@ -121,9 +114,7 @@ pub async fn address_tree_ready_for_rollover<R: RpcConnection>(
     rpc: &mut R,
     merkle_tree: Pubkey,
 ) -> bool {
-    let account =
-        AccountZeroCopy::<account_compression::AddressMerkleTreeAccount>::new(rpc, merkle_tree)
-            .await;
+    let account = AccountZeroCopy::<AddressMerkleTreeAccount>::new(rpc, merkle_tree).await;
     let rent_exemption = rpc
         .get_minimum_balance_for_rent_exemption(account.account.data.len())
         .await

forester-utils/src/address_merkle_tree_config.rs

@@ -0,0 +1,13 @@
+use thiserror::Error;
+
+#[derive(Error, Debug)]
+pub enum ForesterUtilsError {
+    #[error("parse error: {0:?}")]
+    Parse(String),
+    #[error("prover error: {0:?}")]
+    Prover(String),
+    #[error("rpc error: {0:?}")]
+    Rpc(String),
+    #[error("indexer error: {0:?}")]
+    Indexer(String),
+}