feat: handle multiple ws addesses (#7)

Author: lucacillario

README.md

@@ -18,7 +18,7 @@ The project is structured as follows:
 ## Running the chain pusher
 
 ```bash
-cargo run -- --auth_header "Bearer <your_auth_token>" --ws_url "ws_url" --cluster "https://devnet.magicblock.app"
+cargo run -- --auth-header "Bearer <your_auth_token>" --ws-urls "ws_url1,ws_url2" --cluster "https://devnet.magicblock.app"
 ```
 
 ## Consuming Price Data in a Solana Program

src/args.rs

@@ -1,5 +1,6 @@
 use clap::{Parser, ValueEnum};
 use solana_sdk::signature::Keypair;
+use tracing::warn;
 
 #[derive(Debug, Clone, ValueEnum)]
 pub enum ChannelType {
@@ -31,8 +32,18 @@ pub struct Args {
     pub private_key: Option<String>,
     #[arg(long, help = "Authorization header for the WebSocket connection")]
     pub auth_header: Option<String>,
-    #[arg(long, help = "WebSocket URL for the price feed")]
+
+    #[arg(long,
+        help = "DEPRECATED: Single WebSocket URL for the price feed",
+        conflicts_with = "ws_urls",
+        hide = true)]
     pub ws_url: Option<String>,
+
+    #[arg(long,
+        help = "Comma-separated list of WebSocket URLs for the price feed",
+        value_delimiter = ',')]
+    pub ws_urls: Vec<String>,
+
     #[arg(long, help = "Solana cluster URL")]
     pub cluster: Option<String>,
     #[arg(long, help = "Comma-separated list of price feeds")]
@@ -44,11 +55,28 @@ pub struct Args {
     pub channel: Option<ChannelType>,
 }
 
-pub fn get_ws_url(cli_url: Option<String>) -> String {
-    std::env::var("ORACLE_WS_URL")
-        .ok()
-        .or(cli_url)
-        .unwrap_or_else(|| "ws://localhost:8765".to_string())
+pub fn get_ws_urls(cli_url: Option<String>, cli_urls: Vec<String>) -> Vec<String> {
+    if cli_url.is_some() {
+        warn!("'--ws-url' is deprecated, use '--ws-urls' with comma-separated list instead");
+    }
+
+    let env_url = std::env::var("ORACLE_WS_URL").ok();
+    let env_urls = std::env::var("ORACLE_WS_URLS").ok();
+
+    if !cli_urls.is_empty() {
+        cli_urls
+    } else if let Some(urls_str) = env_urls {
+        urls_str
+            .split(',')
+            .map(|url| url.trim().to_string())
+            .filter(|url| !url.is_empty())
+            .collect()
+    } else {
+        let single_url = cli_url
+            .or(env_url)
+            .unwrap_or_else(|| "ws://localhost:8765".to_string());
+        vec![single_url]
+    }
 }
 
 pub fn get_auth_header(cli_auth: Option<String>) -> String {

src/main.rs

@@ -27,7 +27,7 @@ use tokio_native_tls::TlsConnector;
 use tracing::{debug, error, info, warn};
 use url::Url;
 
-use crate::args::{get_auth_header, get_channel, get_price_feeds, get_private_key, get_solana_cluster, get_ws_url, Args};
+use crate::args::{get_auth_header, get_channel, get_price_feeds, get_private_key, get_solana_cluster, get_ws_urls, Args};
 use crate::pyth_lazer::chain_pusher::PythChainPusher;
 use crate::stork::chain_pusher::StorkChainPusher;
 use crate::types::ChainPusher;
@@ -39,27 +39,38 @@ async fn main() {
     let args = Args::parse();
     let private_key = get_private_key(args.private_key);
     let auth_header = get_auth_header(args.auth_header);
-    let ws_url = get_ws_url(args.ws_url);
+    let ws_urls = get_ws_urls(args.ws_url, args.ws_urls);
     let cluster_url = get_solana_cluster(args.cluster);
     let price_feeds = get_price_feeds(args.price_feeds);
     let channel = get_channel(args.channel);
 
     let payer = Keypair::from_base58_string(&private_key);
     info!(wallet_pubkey = ?payer.pubkey(), "Identity initialized");
 
-    let chain_pusher: Arc<dyn ChainPusher> = if ws_url.contains("stork") {
+    let chain_pusher: Arc<dyn ChainPusher> = if ws_urls.iter().any(|url| url.contains("stork")) {
         Arc::new(StorkChainPusher::new(&cluster_url, payer).await)
     } else {
         Arc::new(PythChainPusher::new(&cluster_url, payer).await)
     };
 
     loop {
-        if let Err(e) =
-            run_websocket_client(&chain_pusher, &ws_url, &auth_header, &price_feeds, &channel).await
-        {
-            error!(error = ?e, "WebSocket connection error, attempting reconnection");
+        let mut last_error = None;
+
+        for ws_url in &ws_urls {
+            match run_websocket_client(&chain_pusher, ws_url, &auth_header, &price_feeds, &channel).await {
+                Ok(_) => break,
+                Err(e) => {
+                    error!(error = ?e, url = ws_url, "WebSocket connection failed, trying next URL");
+                    last_error = Some(e);
+                }
+            }
+        }
+
+        // if all URLs fail, wait before trying again
+        if let Some(e) = last_error {
+            error!(error = ?e, "All WebSocket URLs failed, retrying in 5 seconds");
+            time::sleep(Duration::from_secs(5)).await;
         }
-        time::sleep(Duration::from_secs(5)).await;
     }
 }