Refactor into rpc_multi_client module

Author: merolish

src/agent.rs

@@ -83,6 +83,7 @@ pub mod pyth;
 pub mod services;
 pub mod solana;
 pub mod state;
+pub mod utils;
 
 lazy_static! {
     /// A static exit flag to indicate to running threads that we're shutting down. This is used to

src/agent/services/exporter.rs

@@ -8,14 +8,14 @@ use {
             exporter::Exporter,
             transactions::Transactions,
         },
+        utils::rpc_multi_client::RpcMultiClient,
     },
     anyhow::Result,
     futures_util::future,
     serde::{
         Deserialize,
         Serialize,
     },
-    solana_client::nonblocking::rpc_client::RpcClient,
     solana_sdk::commitment_config::CommitmentConfig,
     std::{
         sync::Arc,
@@ -27,6 +27,7 @@ use {
         time::Interval,
     },
     tracing::instrument,
+    url::Url,
 };
 
 #[derive(Clone, Serialize, Deserialize, Debug)]
@@ -111,7 +112,7 @@ pub struct NetworkState {
 /// fetching the blockhash and slot number.
 struct NetworkStateQuerier {
     /// The RPC client
-    rpc_clients: Vec<RpcClient>,
+    rpc_multi_client: RpcMultiClient,
 
     /// The interval with which to query the network state
     query_interval: Interval,
@@ -129,17 +130,14 @@ impl NetworkStateQuerier {
         )
     )]
     pub fn new(
-        rpc_urls: &Vec<String>,
+        rpc_urls: &Vec<Url>,
         rpc_timeout: Duration,
         query_interval: Interval,
         network_state_tx: watch::Sender<NetworkState>,
     ) -> Self {
-        let rpc_clients = rpc_urls
-            .iter()
-            .map(|rpc_url| RpcClient::new_with_timeout(rpc_url.clone(), rpc_timeout))
-            .collect();
+        let rpc_multi_client = RpcMultiClient::new_with_timeout(rpc_urls.clone(), rpc_timeout);
         NetworkStateQuerier {
-            rpc_clients,
+            rpc_multi_client,
             query_interval,
             network_state_tx,
         }
@@ -156,12 +154,11 @@ impl NetworkStateQuerier {
 
     #[instrument(skip(self))]
     async fn query_network_state(&mut self) -> Result<()> {
-        // TODO: These are polled every 200ms and errors are simply logged.
-        // TODO: Should we retry/fallback on failure?
         // Fetch the blockhash and current slot in parallel
-        let current_slot_future =
-            self.rpc_clients[0].get_slot_with_commitment(CommitmentConfig::confirmed());
-        let latest_blockhash_future = self.rpc_clients[0].get_latest_blockhash();
+        let current_slot_future = self
+            .rpc_multi_client
+            .get_slot_with_commitment(CommitmentConfig::confirmed());
+        let latest_blockhash_future = self.rpc_multi_client.get_latest_blockhash();
 
         let (current_slot_result, latest_blockhash_result) =
             future::join(current_slot_future, latest_blockhash_future).await;
@@ -229,8 +226,8 @@ mod exporter {
                 publish_batches,
                 Exporter,
             },
+            utils::rpc_multi_client::RpcMultiClient,
         },
-        solana_client::nonblocking::rpc_client::RpcClient,
         solana_sdk::commitment_config::CommitmentConfig,
         std::sync::Arc,
         tokio::sync::watch,
@@ -249,21 +246,14 @@ mod exporter {
         let mut dynamic_compute_unit_price_update_interval =
             tokio::time::interval(config.exporter.publish_interval_duration);
 
-        let clients: Arc<Vec<RpcClient>> = Arc::new(
-            config
-                .rpc_urls
-                .iter()
-                .map(|rpc_url| {
-                    RpcClient::new_with_timeout_and_commitment(
-                        rpc_url.clone(),
-                        config.rpc_timeout,
-                        CommitmentConfig {
-                            commitment: config.oracle.commitment,
-                        },
-                    )
-                })
-                .collect(),
-        );
+        let rpc_multi_client: Arc<RpcMultiClient> =
+            Arc::new(RpcMultiClient::new_with_timeout_and_commitment(
+                config.rpc_urls.clone(),
+                config.rpc_timeout,
+                CommitmentConfig {
+                    commitment: config.oracle.commitment,
+                },
+            ));
         let Ok(key_store) = KeyStore::new(config.key_store.clone()) else {
             tracing::warn!("Key store not available, Exporter won't start.");
             return;
@@ -282,7 +272,7 @@ mod exporter {
                             let publisher_buffer_key = Exporter::get_publisher_buffer_key(&*state).await;
                             if let Err(err) = publish_batches(
                                 state.clone(),
-                                clients.clone(),
+                                rpc_multi_client.clone(),
                                 network,
                                 &network_state_rx,
                                 key_store.accumulator_key,
@@ -310,7 +300,7 @@ mod exporter {
                             if let Err(err) = Exporter::update_recent_compute_unit_price(
                                 &*state,
                                 &publish_keypair,
-                                &clients,
+                                &rpc_multi_client,
                                 config.exporter.staleness_threshold,
                                 config.exporter.unchanged_publish_threshold,
                             ).await {
@@ -329,12 +319,12 @@ mod transaction_monitor {
         crate::agent::{
             solana::network,
             state::transactions::Transactions,
+            utils::rpc_multi_client::RpcMultiClient,
         },
         serde::{
             Deserialize,
             Serialize,
         },
-        solana_client::nonblocking::rpc_client::RpcClient,
         std::{
             sync::Arc,
             time::Duration,
@@ -369,17 +359,16 @@ mod transaction_monitor {
     where
         S: Transactions,
     {
-        let rpc_clients = config
-            .rpc_urls
-            .iter()
-            .map(|rpc_url| RpcClient::new_with_timeout(rpc_url.clone(), config.rpc_timeout))
-            .collect();
+        let rpc_multi_client =
+            RpcMultiClient::new_with_timeout(config.rpc_urls.clone(), config.rpc_timeout);
         let mut poll_interval =
             tokio::time::interval(config.exporter.transaction_monitor.poll_interval_duration);
 
         loop {
             poll_interval.tick().await;
-            if let Err(err) = Transactions::poll_transactions_status(&*state, &rpc_clients).await {
+            if let Err(err) =
+                Transactions::poll_transactions_status(&*state, &rpc_multi_client).await
+            {
                 tracing::error!(err = ?err, "Transaction monitor failed.");
             }
         }

src/agent/services/keypairs.rs

@@ -6,13 +6,13 @@ use {
     crate::agent::{
         solana::network::Network,
         state::keypairs::Keypairs,
+        utils::rpc_multi_client::RpcMultiClient,
     },
     anyhow::{
         bail,
         Result,
     },
     serde::Deserialize,
-    solana_client::nonblocking::rpc_client::RpcClient,
     solana_sdk::{
         commitment_config::CommitmentConfig,
         signature::Keypair,
@@ -23,6 +23,7 @@ use {
         sync::Arc,
     },
     tokio::task::JoinHandle,
+    url::Url,
     warp::{
         hyper::StatusCode,
         reject::Rejection,
@@ -61,8 +62,8 @@ impl Default for Config {
 }
 
 pub async fn keypairs<S>(
-    primary_rpc_urls: Vec<String>,
-    secondary_rpc_urls: Option<Vec<String>>,
+    primary_rpc_urls: Vec<Url>,
+    secondary_rpc_urls: Option<Vec<Url>>,
     config: Config,
     state: Arc<S>,
 ) -> Vec<JoinHandle<()>>
@@ -160,7 +161,7 @@ async fn handle_new_keypair<'a, 'b: 'a, S>(
     network: Network,
     new_keypair_bytes: Vec<u8>,
     min_keypair_balance_sol: u64,
-    rpc_urls: Vec<String>,
+    rpc_urls: Vec<Url>,
     network_name: &'b str,
 ) -> WithStatus<&'static str>
 where
@@ -205,9 +206,11 @@ where
 pub async fn validate_keypair(
     kp: &Keypair,
     min_keypair_balance_sol: u64,
-    rpc_urls: Vec<String>,
+    rpc_urls: Vec<Url>,
 ) -> Result<()> {
-    let balance_lamports = match get_balance(kp, rpc_urls).await {
+    let rpc_multi_client =
+        RpcMultiClient::new_with_commitment(rpc_urls, CommitmentConfig::confirmed());
+    let balance_lamports = match rpc_multi_client.get_balance(kp).await {
         Ok(balance_lamports) => balance_lamports,
         Err(_) => bail!("Could not check keypair's balance"),
     };
@@ -225,14 +228,3 @@ pub async fn validate_keypair(
         )))
     }
 }
-
-async fn get_balance(kp: &Keypair, rpc_urls: Vec<String>) -> Result<u64> {
-    for rpc_url in rpc_urls {
-        let c = RpcClient::new_with_commitment(rpc_url.clone(), CommitmentConfig::confirmed());
-        match c.get_balance(&kp.pubkey()).await {
-            Ok(balance) => return Ok(balance),
-            Err(e) => tracing::warn!("getBalance error for rpc endpoint {}: {}", rpc_url, e),
-        }
-    }
-    bail!("getBalance failed for all RPC endpoints")
-}

src/agent/services/oracle.rs

@@ -12,14 +12,12 @@ use {
             },
         },
         state::oracle::Oracle,
+        utils::rpc_multi_client::RpcMultiClient,
     },
     anyhow::Result,
     solana_account_decoder::UiAccountEncoding,
     solana_client::{
-        nonblocking::{
-            pubsub_client::PubsubClient,
-            rpc_client::RpcClient,
-        },
+        nonblocking::pubsub_client::PubsubClient,
         rpc_config::{
             RpcAccountInfoConfig,
             RpcProgramAccountsConfig,
@@ -168,21 +166,13 @@ async fn poller<S>(
 {
     // Setup an RpcClient for manual polling.
     let mut tick = tokio::time::interval(config.oracle.poll_interval_duration);
-    let clients: Arc<Vec<RpcClient>> = Arc::new(
-        config
-            .rpc_urls
-            .iter()
-            .map(|rpc_url| {
-                RpcClient::new_with_timeout_and_commitment(
-                    rpc_url.clone(),
-                    config.rpc_timeout,
-                    CommitmentConfig {
-                        commitment: config.oracle.commitment,
-                    },
-                )
-            })
-            .collect(),
-    );
+    let rpc_multi_client = Arc::new(RpcMultiClient::new_with_timeout_and_commitment(
+        config.rpc_urls.clone(),
+        config.rpc_timeout,
+        CommitmentConfig {
+            commitment: config.oracle.commitment,
+        },
+    ));
 
     loop {
         if let Err(err) = async {
@@ -194,7 +184,7 @@ async fn poller<S>(
                 oracle_program_key,
                 publish_keypair.as_ref(),
                 pyth_price_store_program_key,
-                &clients,
+                &rpc_multi_client,
                 max_lookup_batch_size,
             )
             .await?;

src/agent/solana.rs

@@ -17,6 +17,7 @@ pub mod network {
             Serialize,
         },
         std::time::Duration,
+        url::Url,
     };
 
     #[derive(Clone, Copy, Serialize, Deserialize, Debug)]
@@ -25,8 +26,8 @@ pub mod network {
         Secondary,
     }
 
-    pub fn default_rpc_url() -> Vec<String> {
-        vec!["http://localhost:8899".to_string()]
+    pub fn default_rpc_urls() -> Vec<Url> {
+        vec![Url::parse("http://localhost:8899").unwrap()]
     }
 
     pub fn default_wss_url() -> String {
@@ -41,8 +42,8 @@ pub mod network {
     #[derive(Clone, Serialize, Deserialize, Debug)]
     pub struct Config {
         /// HTTP RPC endpoint list
-        #[serde(default = "default_rpc_url")]
-        pub rpc_urls:    Vec<String>,
+        #[serde(default = "default_rpc_urls")]
+        pub rpc_urls:    Vec<Url>,
         /// WSS RPC endpoint
         #[serde(default = "default_wss_url")]
         pub wss_url:     String,