feat: bb8-based SolanaRpcPool (#1075)

* Replace the old RpcPool implementation with bb8-based connection pooling for improved resource management.

* Fix missing newline at end of rpc_pool.rs

* Remove unused code in forester/src/lib.rs

* chore: export environment variables in install.sh

Export the environment variables RUSTUP_HOME, CARGO_HOME, and PATH in the install.sh script.

Author: sergeytimoshin

Cargo.lock

@@ -43,6 +43,8 @@ dotenvy = "0.15.7"
 crossbeam-channel = "0.5.12"
 tokio-stream = "0.1.14"
 base64 = "0.12.3"
+async-trait = "0.1.81"
+bb8 = "0.8.5"
 
 [dev-dependencies]
 function_name = "0.3.0"

forester/Cargo.toml

@@ -1,3 +1,4 @@
+use crate::rpc_pool::PoolError;
 use account_compression::initialize_address_merkle_tree::Error as AccountCompressionError;
 use light_hash_set::HashSetError;
 use light_test_utils::indexer::IndexerError;
@@ -157,3 +158,9 @@ impl From<JoinError> for ForesterError {
         ForesterError::JoinError(err.to_string())
     }
 }
+
+impl From<PoolError> for ForesterError {
+    fn from(err: PoolError) -> Self {
+        ForesterError::Custom(err.to_string())
+    }
+}

forester/src/epoch_manager.rs

@@ -11,6 +11,7 @@ pub mod utils;
 
 use crate::epoch_manager::{fetch_queue_item_data, run_service, WorkReport};
 use crate::errors::ForesterError;
+use crate::rpc_pool::SolanaRpcPool;
 use crate::utils::get_protocol_config;
 pub use config::{ForesterConfig, ForesterEpochInfo};
 use env_logger::Env;
@@ -19,8 +20,8 @@ use light_test_utils::indexer::Indexer;
 use light_test_utils::rpc::rpc_connection::RpcConnection;
 use light_test_utils::rpc::SolanaRpcConnection;
 use log::info;
-pub use rpc_pool::RpcPool;
 pub use settings::init_config;
+use solana_sdk::commitment_config::CommitmentConfig;
 use solana_sdk::native_token::LAMPORTS_PER_SOL;
 use solana_sdk::signature::Signer;
 use std::sync::Arc;
@@ -36,16 +37,15 @@ pub async fn run_queue_info(
     trees: Vec<TreeAccounts>,
     queue_type: TreeType,
 ) {
-    let rpc = SolanaRpcConnection::new(config.external_services.rpc_url.to_string(), None);
-    let rpc = Arc::new(Mutex::new(rpc));
+    let mut rpc = SolanaRpcConnection::new(config.external_services.rpc_url.to_string(), None);
     let state_trees: Vec<_> = trees
         .iter()
         .filter(|t| t.tree_type == queue_type)
         .cloned()
         .collect();
 
     for tree_data in state_trees {
-        let queue_length = fetch_queue_item_data(rpc.clone(), &tree_data.queue)
+        let queue_length = fetch_queue_item_data(&mut rpc, &tree_data.queue)
             .await
             .unwrap()
             .len();
@@ -62,26 +62,31 @@ pub async fn run_pipeline<R: RpcConnection, I: Indexer<R>>(
     shutdown: oneshot::Receiver<()>,
     work_report_sender: mpsc::Sender<WorkReport>,
 ) -> Result<(), ForesterError> {
-    let rpc_pool = Arc::new(RpcPool::<R>::new(config.clone()).await);
+    let rpc_pool = SolanaRpcPool::<R>::new(
+        config.external_services.rpc_url.to_string(),
+        CommitmentConfig::confirmed(),
+        config.rpc_pool_size as u32,
+    )
+    .await
+    .map_err(|e| ForesterError::Custom(e.to_string()))?;
 
     {
-        let rpc = rpc_pool.get_connection().await;
-        let mut rpc = rpc.lock().await;
+        let mut rpc = rpc_pool.get_connection().await?;
         rpc.airdrop_lamports(&config.payer_keypair.pubkey(), LAMPORTS_PER_SOL * 100_000)
             .await
             .unwrap();
     }
 
     let protocol_config = {
-        let rpc = rpc_pool.get_connection().await;
-        Arc::new(get_protocol_config(rpc.clone()).await)
+        let mut rpc = rpc_pool.get_connection().await?;
+        get_protocol_config(&mut *rpc).await
     };
 
     info!("Starting Forester pipeline");
     run_service(
         config,
-        protocol_config,
-        rpc_pool,
+        Arc::new(protocol_config),
+        Arc::new(rpc_pool),
         indexer,
         shutdown,
         work_report_sender,

forester/src/errors.rs

@@ -1,4 +1,3 @@
-use std::ops::DerefMut;
 use std::sync::Arc;
 
 use light_registry::account_compression_cpi::sdk::{
@@ -15,7 +14,7 @@ use solana_sdk::transaction::Transaction;
 use tokio::sync::Mutex;
 
 use crate::errors::ForesterError;
-use crate::{ForesterConfig, RpcPool};
+use crate::ForesterConfig;
 use account_compression::utils::constants::{
     STATE_MERKLE_TREE_CANOPY_DEPTH, STATE_MERKLE_TREE_HEIGHT,
 };
@@ -40,15 +39,14 @@ use light_test_utils::{
 };
 
 pub async fn is_tree_ready_for_rollover<R: RpcConnection>(
-    rpc: Arc<Mutex<R>>,
+    rpc: &mut R,
     tree_pubkey: Pubkey,
     tree_type: TreeType,
 ) -> Result<bool, ForesterError> {
     info!(
         "Checking if tree is ready for rollover: {:?}",
         tree_pubkey.to_string()
     );
-    let mut rpc = rpc.lock().await;
     match tree_type {
         TreeType::State => {
             let account = rpc
@@ -63,7 +61,7 @@ pub async fn is_tree_ready_for_rollover<R: RpcConnection>(
             }
             let merkle_tree =
                 get_concurrent_merkle_tree::<StateMerkleTreeAccount, R, Poseidon, 26>(
-                    &mut rpc,
+                    rpc,
                     tree_pubkey,
                 )
                 .await;
@@ -87,7 +85,7 @@ pub async fn is_tree_ready_for_rollover<R: RpcConnection>(
 
             let merkle_tree =
                 get_indexed_merkle_tree::<AddressMerkleTreeAccount, R, Poseidon, usize, 26, 16>(
-                    &mut rpc,
+                    rpc,
                     tree_pubkey,
                 )
                 .await;
@@ -104,18 +102,17 @@ pub async fn is_tree_ready_for_rollover<R: RpcConnection>(
 #[allow(dead_code)]
 pub async fn rollover_state_merkle_tree<R: RpcConnection, I: Indexer<R>>(
     config: Arc<ForesterConfig>,
-    rpc_pool: Arc<RpcPool<R>>,
+    rpc: &mut R,
     indexer: Arc<Mutex<I>>,
     tree_accounts: &TreeAccounts,
 ) -> Result<(), ForesterError> {
     let new_nullifier_queue_keypair = Keypair::new();
     let new_merkle_tree_keypair = Keypair::new();
     let new_cpi_signature_keypair = Keypair::new();
 
-    let rpc = rpc_pool.get_connection().await;
     let rollover_signature = perform_state_merkle_tree_roll_over_forester(
         &config.payer_keypair,
-        rpc.clone(),
+        rpc,
         &new_nullifier_queue_keypair,
         &new_merkle_tree_keypair,
         &new_cpi_signature_keypair,
@@ -146,7 +143,7 @@ pub async fn rollover_state_merkle_tree<R: RpcConnection, I: Indexer<R>>(
 #[allow(clippy::too_many_arguments)]
 pub async fn perform_state_merkle_tree_roll_over_forester<R: RpcConnection>(
     payer: &Keypair,
-    context: Arc<Mutex<R>>,
+    context: &mut R,
     new_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
     new_cpi_context_keypair: &Keypair,
@@ -155,7 +152,7 @@ pub async fn perform_state_merkle_tree_roll_over_forester<R: RpcConnection>(
     old_cpi_context_pubkey: &Pubkey,
 ) -> Result<solana_sdk::signature::Signature, RpcError> {
     let instructions = create_rollover_state_merkle_tree_instructions(
-        context.clone(),
+        context,
         &payer.pubkey(),
         new_queue_keypair,
         new_address_merkle_tree_keypair,
@@ -165,7 +162,6 @@ pub async fn perform_state_merkle_tree_roll_over_forester<R: RpcConnection>(
         old_cpi_context_pubkey,
     )
     .await;
-    let mut context = context.lock().await;
     let blockhash = context.get_latest_blockhash().await.unwrap();
     let transaction = Transaction::new_signed_with_payer(
         &instructions,
@@ -178,16 +174,15 @@ pub async fn perform_state_merkle_tree_roll_over_forester<R: RpcConnection>(
 
 pub async fn rollover_address_merkle_tree<R: RpcConnection, I: Indexer<R>>(
     config: Arc<ForesterConfig>,
-    rpc_pool: Arc<RpcPool<R>>,
+    rpc: &mut R,
     indexer: Arc<Mutex<I>>,
     tree_data: &TreeAccounts,
 ) -> Result<(), ForesterError> {
     let new_nullifier_queue_keypair = Keypair::new();
     let new_merkle_tree_keypair = Keypair::new();
-    let rpc = rpc_pool.get_connection().await;
     perform_address_merkle_tree_roll_over(
         &config.payer_keypair,
-        rpc.clone(),
+        rpc,
         &new_nullifier_queue_keypair,
         &new_merkle_tree_keypair,
         &tree_data.merkle_tree,
@@ -205,22 +200,21 @@ pub async fn rollover_address_merkle_tree<R: RpcConnection, I: Indexer<R>>(
 
 pub async fn perform_address_merkle_tree_roll_over<R: RpcConnection>(
     payer: &Keypair,
-    context: Arc<Mutex<R>>,
+    context: &mut R,
     new_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
     old_merkle_tree_pubkey: &Pubkey,
     old_queue_pubkey: &Pubkey,
 ) -> Result<solana_sdk::signature::Signature, RpcError> {
     let instructions = create_rollover_address_merkle_tree_instructions(
-        context.clone(),
+        context,
         &payer.pubkey(),
         new_queue_keypair,
         new_address_merkle_tree_keypair,
         old_merkle_tree_pubkey,
         old_queue_pubkey,
     )
     .await;
-    let mut context = context.lock().await;
     let blockhash = context.get_latest_blockhash().await.unwrap();
     let transaction = Transaction::new_signed_with_payer(
         &instructions,
@@ -232,16 +226,15 @@ pub async fn perform_address_merkle_tree_roll_over<R: RpcConnection>(
 }
 
 pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
-    rpc: Arc<Mutex<R>>,
+    rpc: &mut R,
     authority: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_address_merkle_tree_keypair: &Keypair,
     merkle_tree_pubkey: &Pubkey,
     nullifier_queue_pubkey: &Pubkey,
 ) -> Vec<Instruction> {
-    let mut rpc = rpc.lock().await;
     let (merkle_tree_config, queue_config) = get_address_bundle_config(
-        rpc.deref_mut(),
+        rpc,
         AddressMerkleTreeAccounts {
             merkle_tree: *merkle_tree_pubkey,
             queue: *nullifier_queue_pubkey,
@@ -250,7 +243,7 @@ pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
     .await;
     let (merkle_tree_rent_exemption, queue_rent_exemption) =
         get_rent_exemption_for_address_merkle_tree_and_queue(
-            rpc.deref_mut(),
+            rpc,
             &merkle_tree_config,
             &queue_config,
         )
@@ -291,7 +284,7 @@ pub async fn create_rollover_address_merkle_tree_instructions<R: RpcConnection>(
 
 #[allow(clippy::too_many_arguments)]
 pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
-    rpc: Arc<Mutex<R>>,
+    rpc: &mut R,
     authority: &Pubkey,
     new_nullifier_queue_keypair: &Keypair,
     new_state_merkle_tree_keypair: &Keypair,
@@ -300,9 +293,8 @@ pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
     nullifier_queue_pubkey: &Pubkey,
     old_cpi_context_pubkey: &Pubkey,
 ) -> Vec<Instruction> {
-    let mut rpc = rpc.lock().await;
     let (merkle_tree_config, queue_config) = get_state_bundle_config(
-        rpc.deref_mut(),
+        rpc,
         StateMerkleTreeAccounts {
             merkle_tree: *merkle_tree_pubkey,
             nullifier_queue: *nullifier_queue_pubkey,
@@ -311,12 +303,8 @@ pub async fn create_rollover_state_merkle_tree_instructions<R: RpcConnection>(
     )
     .await;
     let (state_merkle_tree_rent_exemption, queue_rent_exemption) =
-        get_rent_exemption_for_state_merkle_tree_and_queue(
-            rpc.deref_mut(),
-            &merkle_tree_config,
-            &queue_config,
-        )
-        .await;
+        get_rent_exemption_for_state_merkle_tree_and_queue(rpc, &merkle_tree_config, &queue_config)
+            .await;
     let create_nullifier_queue_instruction = create_account_instruction(
         authority,
         queue_rent_exemption.size,