feat: allow for multiple ws endpoints in remote config (#372)

<!-- greptile_comment -->

## Greptile Summary

This PR adds support for multiple WebSocket endpoints in the remote
configuration, enabling more flexible validator setup and improved
connection management.

- Added new `CustomWithMultipleWs` variant in
`magicblock-config/src/accounts.rs` to handle multiple WebSocket URLs
- Modified `RemoteAccountUpdatesShard` and `RemoteAccountUpdatesWorker`
to use direct URL and commitment fields instead of `RpcProviderConfig`
- Updated `magicblock-mutator/src/cluster.rs` needs improvement in
WebSocket URL handling and error cases
- Added test coverage in `magicblock-config/tests/parse_config.rs` for
multiple WebSocket configurations
- Simplified configuration parsing by removing custom deserialization in
favor of serde defaults



<!-- /greptile_comment -->

---------

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>

Author: bmuddha

magicblock-account-updates/src/remote_account_updates_shard.rs

@@ -6,7 +6,6 @@ use std::{
     sync::{Arc, RwLock},
 };
 
-use conjunto_transwise::RpcProviderConfig;
 use futures_util::StreamExt;
 use log::*;
 use magicblock_metrics::metrics;
@@ -15,7 +14,7 @@ use solana_pubsub_client::nonblocking::pubsub_client::PubsubClient;
 use solana_rpc_client_api::config::RpcAccountInfoConfig;
 use solana_sdk::{
     clock::{Clock, Slot},
-    commitment_config::CommitmentConfig,
+    commitment_config::{CommitmentConfig, CommitmentLevel},
     pubkey::Pubkey,
     sysvar::clock,
 };
@@ -35,7 +34,8 @@ pub enum RemoteAccountUpdatesShardError {
 
 pub struct RemoteAccountUpdatesShard {
     shard_id: String,
-    rpc_provider_config: RpcProviderConfig,
+    url: String,
+    commitment: Option<CommitmentLevel>,
     monitoring_request_receiver: Receiver<(Pubkey, bool)>,
     first_subscribed_slots: Arc<RwLock<HashMap<Pubkey, Slot>>>,
     last_known_update_slots: Arc<RwLock<HashMap<Pubkey, Slot>>>,
@@ -44,14 +44,16 @@ pub struct RemoteAccountUpdatesShard {
 impl RemoteAccountUpdatesShard {
     pub fn new(
         shard_id: String,
-        rpc_provider_config: RpcProviderConfig,
+        url: String,
+        commitment: Option<CommitmentLevel>,
         monitoring_request_receiver: Receiver<(Pubkey, bool)>,
         first_subscribed_slots: Arc<RwLock<HashMap<Pubkey, Slot>>>,
         last_known_update_slots: Arc<RwLock<HashMap<Pubkey, Slot>>>,
     ) -> Self {
         Self {
             shard_id,
-            rpc_provider_config,
+            url,
+            commitment,
             monitoring_request_receiver,
             first_subscribed_slots,
             last_known_update_slots,
@@ -64,15 +66,13 @@ impl RemoteAccountUpdatesShard {
     ) -> Result<(), RemoteAccountUpdatesShardError> {
         // Create a pubsub client
         info!("Shard {}: Starting", self.shard_id);
-        let ws_url = self.rpc_provider_config.ws_url();
-        let pubsub_client = PubsubClient::new(ws_url)
+        let pubsub_client = PubsubClient::new(&self.url)
             .await
             .map_err(RemoteAccountUpdatesShardError::PubsubClientError)?;
         // For every account, we only want the updates, not the actual content of the accounts
         let rpc_account_info_config = Some(RpcAccountInfoConfig {
             commitment: self
-                .rpc_provider_config
-                .commitment()
+                .commitment
                 .map(|commitment| CommitmentConfig { commitment }),
             encoding: Some(UiAccountEncoding::Base64),
             data_slice: Some(UiDataSliceConfig {

magicblock-account-updates/src/remote_account_updates_worker.rs

@@ -7,9 +7,10 @@ use std::{
     time::Duration,
 };
 
-use conjunto_transwise::RpcProviderConfig;
 use log::*;
-use solana_sdk::{clock::Slot, pubkey::Pubkey};
+use solana_sdk::{
+    clock::Slot, commitment_config::CommitmentLevel, pubkey::Pubkey,
+};
 use thiserror::Error;
 use tokio::{
     sync::mpsc::{channel, Receiver, Sender},
@@ -42,7 +43,8 @@ struct RemoteAccountUpdatesWorkerRunner {
 }
 
 pub struct RemoteAccountUpdatesWorker {
-    rpc_provider_configs: Vec<RpcProviderConfig>,
+    ws_urls: Vec<String>,
+    commitment: Option<CommitmentLevel>,
     refresh_interval: Duration,
     monitoring_request_receiver: Receiver<(Pubkey, bool)>,
     monitoring_request_sender: Sender<(Pubkey, bool)>,
@@ -52,13 +54,15 @@ pub struct RemoteAccountUpdatesWorker {
 
 impl RemoteAccountUpdatesWorker {
     pub fn new(
-        rpc_provider_configs: Vec<RpcProviderConfig>,
+        ws_urls: Vec<String>,
+        commitment: Option<CommitmentLevel>,
         refresh_interval: Duration,
     ) -> Self {
         let (monitoring_request_sender, monitoring_request_receiver) =
             channel(INFLIGHT_ACCOUNT_FETCHES_LIMIT);
         Self {
-            rpc_provider_configs,
+            ws_urls,
+            commitment,
             refresh_interval,
             monitoring_request_receiver,
             monitoring_request_sender,
@@ -91,13 +95,12 @@ impl RemoteAccountUpdatesWorker {
         let mut runners = vec![];
         let mut monitored_accounts = HashSet::new();
         // Initialize all the runners for all configs
-        for (index, rpc_provider_config) in
-            self.rpc_provider_configs.iter().enumerate()
-        {
+        for (index, url) in self.ws_urls.iter().enumerate() {
             runners.push(
                 self.create_runner_from_config(
                     index,
-                    rpc_provider_config.clone(),
+                    url.clone(),
+                    self.commitment,
                     &monitored_accounts,
                 )
                 .await,
@@ -127,14 +130,15 @@ impl RemoteAccountUpdatesWorker {
                 }
                 // Periodically we refresh runners to keep them fresh
                 _ = refresh_interval.tick() => {
-                    current_refresh_index = (current_refresh_index + 1) % self.rpc_provider_configs.len();
-                    let rpc_provider_config = self.rpc_provider_configs
+                    current_refresh_index = (current_refresh_index + 1) % self.ws_urls.len();
+                    let url = self.ws_urls
                         .get(current_refresh_index)
                         .unwrap()
                         .clone();
                     let new_runner = self.create_runner_from_config(
                         current_refresh_index,
-                        rpc_provider_config,
+                        url,
+                        self.commitment,
                         &monitored_accounts
                     ).await;
                     let old_runner = std::mem::replace(&mut runners[current_refresh_index], new_runner);
@@ -157,7 +161,8 @@ impl RemoteAccountUpdatesWorker {
     async fn create_runner_from_config(
         &self,
         index: usize,
-        rpc_provider_config: RpcProviderConfig,
+        url: String,
+        commitment: Option<CommitmentLevel>,
         monitored_accounts: &HashSet<Pubkey>,
     ) -> RemoteAccountUpdatesWorkerRunner {
         let (monitoring_request_sender, monitoring_request_receiver) =
@@ -171,7 +176,8 @@ impl RemoteAccountUpdatesWorker {
         let join_handle = tokio::spawn(async move {
             let mut shard = RemoteAccountUpdatesShard::new(
                 shard_id.clone(),
-                rpc_provider_config,
+                url,
+                commitment,
                 monitoring_request_receiver,
                 first_subscribed_slots,
                 last_known_update_slots,

magicblock-account-updates/tests/remote_account_updates.rs

@@ -21,7 +21,8 @@ fn setup() -> (
 ) {
     // Create account updates worker and client
     let mut worker = RemoteAccountUpdatesWorker::new(
-        vec![RpcProviderConfig::devnet(), RpcProviderConfig::devnet()],
+        vec![RpcProviderConfig::devnet().ws_url().into(); 2],
+        Some(solana_sdk::commitment_config::CommitmentLevel::Confirmed),
         Duration::from_secs(1), // We constantly refresh stuff to make it struggle
     );
     let client = RemoteAccountUpdatesClient::new(&worker);

magicblock-accounts/src/utils/mod.rs

@@ -33,6 +33,9 @@ pub fn try_rpc_cluster_from_cluster(
         Cluster::CustomWithWs(http, ws) => {
             Ok(RpcCluster::Custom(http.to_string(), ws.to_string()))
         }
+        Cluster::CustomWithMultipleWs { http, ws } => {
+            Ok(RpcCluster::Custom(http.to_string(), ws[0].to_string()))
+        }
     }
 }
 

magicblock-api/src/external_config.rs

@@ -30,6 +30,10 @@ pub(crate) fn cluster_from_remote(
         CustomWithWs(http, ws) => {
             Cluster::CustomWithWs(http.clone(), ws.clone())
         }
+        CustomWithMultipleWs { http, ws } => Cluster::CustomWithMultipleWs {
+            http: http.clone(),
+            ws: ws.clone(),
+        },
     }
 }
 

magicblock-api/src/magic_validator.rs

@@ -271,12 +271,8 @@ impl MagicValidator {
             RemoteAccountFetcherWorker::new(remote_rpc_config.clone());
 
         let remote_account_updates_worker = RemoteAccountUpdatesWorker::new(
-            // We'll maintain 3 connections constantly (those could be on different nodes if we wanted to)
-            vec![
-                remote_rpc_config.clone(),
-                remote_rpc_config.clone(),
-                remote_rpc_config.clone(),
-            ],
+            accounts_config.remote_cluster.ws_urls(),
+            remote_rpc_config.commitment(),
             // We'll kill/refresh one connection every 50 minutes
             Duration::from_secs(60 * 50),
         );

magicblock-config/src/accounts.rs

@@ -1,10 +1,7 @@
-use std::{fmt, str::FromStr};
+use std::str::FromStr;
 
 use magicblock_accounts_db::config::AccountsDbConfig;
-use serde::{
-    de::{self, Deserializer, SeqAccess, Visitor},
-    Deserialize, Serialize,
-};
+use serde::{Deserialize, Serialize};
 use solana_sdk::{native_token::LAMPORTS_PER_SOL, pubkey::Pubkey};
 use strum_macros::EnumString;
 use url::Url;
@@ -62,90 +59,17 @@ pub enum RemoteConfig {
     #[serde(alias = "local")]
     #[serde(alias = "localhost")]
     Development,
-    #[serde(untagged, deserialize_with = "deserialize_url")]
+    #[serde(untagged)]
     Custom(Url),
-    #[serde(untagged, deserialize_with = "deserialize_tuple_url")]
+    #[serde(untagged)]
     CustomWithWs(Url, Url),
+    #[serde(untagged)]
+    CustomWithMultipleWs {
+        http: Url,
+        ws: Vec<Url>,
+    },
 }
 
-pub fn deserialize_url<'de, D>(deserializer: D) -> Result<Url, D::Error>
-where
-    D: Deserializer<'de>,
-{
-    struct UrlVisitor;
-
-    impl Visitor<'_> for UrlVisitor {
-        type Value = Url;
-
-        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
-            formatter.write_str("a valid URL string")
-        }
-
-        fn visit_str<E>(self, value: &str) -> Result<Url, E>
-        where
-            E: de::Error,
-        {
-            Url::parse(value).map_err(|e| {
-                    // The error returned here by serde is a bit unhelpful so we help out
-                    // by logging a bit more information.
-                    eprintln!(
-                        "RemoteConfig encountered invalid URL '{value}', err: ({e}).",
-                    );
-                    de::Error::custom(e)
-                })
-        }
-    }
-
-    deserializer.deserialize_str(UrlVisitor)
-}
-
-pub fn deserialize_tuple_url<'de, D>(
-    deserializer: D,
-) -> Result<(Url, Url), D::Error>
-where
-    D: Deserializer<'de>,
-{
-    struct UrlTupleVisitor;
-
-    impl<'de> Visitor<'de> for UrlTupleVisitor {
-        type Value = (Url, Url);
-
-        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
-            formatter.write_str("a sequence of two URL strings")
-        }
-
-        fn visit_seq<A>(self, mut seq: A) -> Result<(Url, Url), A::Error>
-        where
-            A: SeqAccess<'de>,
-        {
-            let first: String = seq.next_element()?.ok_or_else(|| {
-                eprintln!("expected a sequence of two URLs: http and ws");
-                de::Error::invalid_length(0, &self)
-            })?;
-            let second: String = seq.next_element()?.ok_or_else(|| {
-                eprintln!("expected a sequence of two URLs: http and ws");
-                de::Error::invalid_length(1, &self)
-            })?;
-
-            let http = Url::parse(&first).map_err(|e| {
-                eprintln!(
-                    "Invalid HTTP URL in RemoteConfig '{first}', err: ({e}).",
-                );
-                de::Error::custom(e)
-            })?;
-            let ws = Url::parse(&second).map_err(|e| {
-                eprintln!(
-                    "Invalid WS URL in RemoteConfig '{second}', err: ({e}).",
-                );
-                de::Error::custom(e)
-            })?;
-
-            Ok((http, ws))
-        }
-    }
-
-    deserializer.deserialize_seq(UrlTupleVisitor)
-}
 // -----------------
 // LifecycleMode
 // -----------------

magicblock-config/tests/parse_config.rs

@@ -244,3 +244,15 @@ payer = { init_sol = 2000, init_lamports = 300_000 }
     let config = toml::from_str::<EphemeralConfig>(toml).unwrap();
     assert!(config.accounts.payer.try_init_lamports().is_err());
 }
+
+#[test]
+fn test_custom_remote_with_multiple_ws() {
+    let toml = r#"
+[accounts]
+remote = { http = "http://localhost:8899", ws = ["ws://awesomews1.com:933", "wss://awesomews2.com:944"] }
+"#;
+
+    let res = toml::from_str::<EphemeralConfig>(toml);
+    println!("{res:?}");
+    assert!(res.is_ok());
+}

magicblock-mutator/src/cluster.rs

@@ -6,6 +6,11 @@ pub const MAINNET_URL: &str = "https://api.mainnet-beta.solana.com";
 pub const DEVNET_URL: &str = "https://api.devnet.solana.com";
 pub const DEVELOPMENT_URL: &str = "http://127.0.0.1:8899";
 
+const WS_MAINNET: &str = "wss://api.mainnet-beta.solana.com/";
+const WS_TESTNET: &str = "wss://api.testnet.solana.com/";
+pub const WS_DEVNET: &str = "wss://api.devnet.solana.com/";
+const WS_DEVELOPMENT: &str = "ws://localhost:8900";
+
 /// TODO(vbrunet)
 ///  - this probably belong in a different crate, "mutator" is specific to the data dump mechanisms
 ///  - conjunto_addresses::cluster::RpcCluster already achieve this and is a full duplicate
@@ -15,6 +20,7 @@ pub enum Cluster {
     Known(ClusterType),
     Custom(Url),
     CustomWithWs(Url, Url),
+    CustomWithMultipleWs { http: Url, ws: Vec<Url> },
 }
 
 impl From<ClusterType> for Cluster {
@@ -35,6 +41,45 @@ impl Cluster {
             },
             Cluster::Custom(url) => url.as_str(),
             Cluster::CustomWithWs(url, _) => url.as_str(),
+            Cluster::CustomWithMultipleWs { http, .. } => http.as_str(),
+        }
+    }
+
+    pub fn ws_urls(&self) -> Vec<String> {
+        use ClusterType::*;
+        const WS_SHARD_COUNT: usize = 3;
+        match self {
+            Cluster::Known(cluster) => vec![
+                match cluster {
+                    Testnet => WS_TESTNET.into(),
+                    MainnetBeta => WS_MAINNET.into(),
+                    Devnet => WS_DEVNET.into(),
+                    Development => WS_DEVELOPMENT.into(),
+                };
+                WS_SHARD_COUNT
+            ],
+            Cluster::Custom(url) => {
+                let mut ws_url = url.clone();
+                ws_url
+                    .set_scheme(if url.scheme() == "https" {
+                        "wss"
+                    } else {
+                        "ws"
+                    })
+                    .expect("valid scheme");
+                if let Some(port) = ws_url.port() {
+                    ws_url
+                        .set_port(Some(port + 1))
+                        .expect("valid url with port");
+                }
+                vec![ws_url.to_string(); WS_SHARD_COUNT]
+            }
+            Cluster::CustomWithWs(_, ws) => {
+                vec![ws.to_string(); WS_SHARD_COUNT]
+            }
+            Cluster::CustomWithMultipleWs { ws, .. } => {
+                ws.iter().map(Url::to_string).collect()
+            }
         }
     }
 }