feat(lazer): add resilient client in rust

Author: keyvankhademi

Cargo.lock

@@ -1,13 +1,14 @@
 [package]
 name = "pyth-lazer-client"
-version = "0.1.3"
+version = "1.0.0"
 edition = "2021"
 description = "A Rust client for Pyth Lazer"
 license = "Apache-2.0"
 
 [dependencies]
 pyth-lazer-protocol = { path = "../protocol", version = "0.9.0" }
 tokio = { version = "1", features = ["full"] }
+tokio-stream = "0.1.17"
 tokio-tungstenite = { version = "0.20", features = ["native-tls"] }
 futures-util = "0.3"
 serde = { version = "1.0", features = ["derive"] }
@@ -17,6 +18,7 @@ anyhow = "1.0"
 tracing = "0.1"
 url = "2.4"
 derive_more = { version = "1.0.0", features = ["from"] }
+backoff = { version = "0.4.0", features = ["futures", "tokio"] }
 
 [dev-dependencies]
 bincode = "1.3.3"
@@ -25,3 +27,4 @@ hex = "0.4.3"
 libsecp256k1 = "0.7.1"
 bs58 = "0.5.1"
 alloy-primitives = "0.8.19"
+tracing-subscriber = { version = "0.3.19", features = ["env-filter", "json"] }

lazer/sdk/rust/client/Cargo.toml

@@ -1,6 +1,6 @@
 use base64::Engine;
-use futures_util::StreamExt;
-use pyth_lazer_client::{AnyResponse, LazerClient};
+use pyth_lazer_client::client::PythLazerClient;
+use pyth_lazer_client::ws_connection::AnyResponse;
 use pyth_lazer_protocol::message::{
     EvmMessage, LeEcdsaMessage, LeUnsignedMessage, Message, SolanaMessage,
 };
@@ -9,8 +9,10 @@ use pyth_lazer_protocol::router::{
     Channel, DeliveryFormat, FixedRate, Format, JsonBinaryEncoding, PriceFeedId, PriceFeedProperty,
     SubscriptionParams, SubscriptionParamsRepr,
 };
-use pyth_lazer_protocol::subscription::{Request, Response, SubscribeRequest, SubscriptionId};
+use pyth_lazer_protocol::subscription::{Response, SubscribeRequest, SubscriptionId};
 use tokio::pin;
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::EnvFilter;
 
 fn get_lazer_access_token() -> String {
     // Place your access token in your env at LAZER_ACCESS_TOKEN or set it here
@@ -20,11 +22,22 @@ fn get_lazer_access_token() -> String {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    tracing_subscriber::fmt()
+        .with_env_filter(
+            EnvFilter::builder()
+                .with_default_directive(LevelFilter::INFO.into())
+                .from_env()?,
+        )
+        .json()
+        .init();
+
     // Create and start the client
-    let mut client = LazerClient::new(
-        "wss://pyth-lazer.dourolabs.app/v1/stream",
-        &get_lazer_access_token(),
-    )?;
+    let mut client = PythLazerClient::new(
+        vec!["wss://pyth-lazer.dourolabs.app/v1/stream".to_string()],
+        get_lazer_access_token(),
+        1,
+    );
+
     let stream = client.start().await?;
     pin!(stream);
 
@@ -72,16 +85,16 @@ async fn main() -> anyhow::Result<()> {
     ];
 
     for req in subscription_requests {
-        client.subscribe(Request::Subscribe(req)).await?;
+        client.subscribe(req).await?;
     }
 
     println!("Subscribed to price feeds. Waiting for updates...");
 
     // Process the first few updates
     let mut count = 0;
-    while let Some(msg) = stream.next().await {
+    while let Some(msg) = stream.recv().await {
         // The stream gives us base64-encoded binary messages. We need to decode, parse, and verify them.
-        match msg? {
+        match msg {
             AnyResponse::Json(msg) => match msg {
                 Response::StreamUpdated(update) => {
                     println!("Received a JSON update for {:?}", update.subscription_id);
@@ -189,8 +202,6 @@ async fn main() -> anyhow::Result<()> {
         println!("Unsubscribed from {sub_id:?}");
     }
 
-    tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
-    client.close().await?;
     Ok(())
 }
 

lazer/sdk/rust/client/examples/subscribe_price_feeds.rs

@@ -0,0 +1,88 @@
+use crate::{
+    resilient_ws_connection::PythLazerResilientWSConnection, ws_connection::AnyResponse,
+    CHANNEL_CAPACITY,
+};
+use anyhow::Result;
+use futures_util::stream;
+use pyth_lazer_protocol::subscription::{SubscribeRequest, SubscriptionId};
+use tokio::sync::mpsc::{self, error::TrySendError};
+use tokio_stream::{wrappers::ReceiverStream, StreamExt};
+use tracing::{error, warn};
+
+pub struct PythLazerClient {
+    endpoints: Vec<String>,
+    access_token: String,
+    num_connections: usize,
+    ws_connections: Vec<PythLazerResilientWSConnection>,
+    receivers: Vec<mpsc::Receiver<AnyResponse>>,
+}
+
+impl PythLazerClient {
+    /// Creates a new client instance
+    ///
+    /// # Arguments
+    /// * `endpoints` - A vector of endpoint URLs
+    /// * `access_token` - The access token for authentication
+    /// * `num_connections` - The number of WebSocket connections to maintain
+    pub fn new(endpoints: Vec<String>, access_token: String, num_connections: usize) -> Self {
+        Self {
+            endpoints,
+            access_token,
+            num_connections,
+            ws_connections: Vec::with_capacity(num_connections),
+            receivers: Vec::with_capacity(num_connections),
+        }
+    }
+
+    pub async fn start(&mut self) -> Result<mpsc::Receiver<AnyResponse>> {
+        let (sender, receiver) = mpsc::channel::<AnyResponse>(CHANNEL_CAPACITY);
+
+        for i in 0..self.num_connections {
+            let endpoint = self.endpoints[i % self.endpoints.len()].clone();
+            let (sender, receiver) = mpsc::channel::<AnyResponse>(CHANNEL_CAPACITY);
+            let connection = PythLazerResilientWSConnection::new(
+                endpoint,
+                self.access_token.clone(),
+                sender.clone(),
+            );
+            self.ws_connections.push(connection);
+            self.receivers.push(receiver);
+        }
+
+        let streams: Vec<_> = self.receivers.drain(..).map(ReceiverStream::new).collect();
+        let mut merged_stream = stream::select_all(streams);
+
+        tokio::spawn(async move {
+            while let Some(response) = merged_stream.next().await {
+                match sender.try_send(response) {
+                    Ok(_) => (),
+                    Err(TrySendError::Full(r)) => {
+                        warn!("Sender channel is full, responses will be delayed");
+                        if sender.send(r).await.is_err() {
+                            error!("Sender channel is closed, stopping client");
+                        }
+                    }
+                    Err(TrySendError::Closed(_)) => {
+                        error!("Sender channel is closed, stopping client");
+                    }
+                }
+            }
+        });
+
+        Ok(receiver)
+    }
+
+    pub async fn subscribe(&mut self, subscribe_request: SubscribeRequest) -> Result<()> {
+        for connection in &mut self.ws_connections {
+            connection.subscribe(subscribe_request.clone()).await?;
+        }
+        Ok(())
+    }
+
+    pub async fn unsubscribe(&mut self, subscription_id: SubscriptionId) -> Result<()> {
+        for connection in &mut self.ws_connections {
+            connection.unsubscribe(subscription_id).await?;
+        }
+        Ok(())
+    }
+}

lazer/sdk/rust/client/src/client.rs

@@ -1,138 +1,5 @@
-use anyhow::Result;
-use derive_more::From;
-use futures_util::{SinkExt, StreamExt, TryStreamExt};
-use pyth_lazer_protocol::{
-    binary_update::BinaryWsUpdate,
-    subscription::{ErrorResponse, Request, Response, SubscriptionId, UnsubscribeRequest},
-};
-use tokio_tungstenite::{connect_async, tungstenite::Message};
-use url::Url;
+const CHANNEL_CAPACITY: usize = 1000;
 
-/// A WebSocket client for consuming Pyth Lazer price feed updates
-///
-/// This client provides a simple interface to:
-/// - Connect to a Lazer WebSocket endpoint
-/// - Subscribe to price feed updates
-/// - Receive updates as a stream of messages
-///
-pub struct LazerClient {
-    endpoint: Url,
-    access_token: String,
-    ws_sender: Option<
-        futures_util::stream::SplitSink<
-            tokio_tungstenite::WebSocketStream<
-                tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
-            >,
-            Message,
-        >,
-    >,
-}
-
-#[derive(Debug, Clone, PartialEq, Eq, Hash, From)]
-pub enum AnyResponse {
-    Json(Response),
-    Binary(BinaryWsUpdate),
-}
-
-impl LazerClient {
-    /// Creates a new Lazer client instance
-    ///
-    /// # Arguments
-    /// * `endpoint` - The WebSocket URL of the Lazer service
-    /// * `access_token` - Access token for authentication
-    ///
-    /// # Returns
-    /// Returns a new client instance (not yet connected)
-    pub fn new(endpoint: &str, access_token: &str) -> Result<Self> {
-        let endpoint = Url::parse(endpoint)?;
-        let access_token = access_token.to_string();
-        Ok(Self {
-            endpoint,
-            access_token,
-            ws_sender: None,
-        })
-    }
-
-    /// Starts the WebSocket connection
-    ///
-    /// # Returns
-    /// Returns a stream of responses from the server
-    pub async fn start(&mut self) -> Result<impl futures_util::Stream<Item = Result<AnyResponse>>> {
-        let url = self.endpoint.clone();
-        let mut request =
-            tokio_tungstenite::tungstenite::client::IntoClientRequest::into_client_request(url)?;
-
-        request.headers_mut().insert(
-            "Authorization",
-            format!("Bearer {}", self.access_token).parse().unwrap(),
-        );
-
-        let (ws_stream, _) = connect_async(request).await?;
-        let (ws_sender, ws_receiver) = ws_stream.split();
-
-        self.ws_sender = Some(ws_sender);
-        let response_stream =
-            ws_receiver
-                .map_err(anyhow::Error::from)
-                .try_filter_map(|msg| async {
-                    let r: Result<Option<AnyResponse>> = match msg {
-                        Message::Text(text) => {
-                            Ok(Some(serde_json::from_str::<Response>(&text)?.into()))
-                        }
-                        Message::Binary(data) => {
-                            Ok(Some(BinaryWsUpdate::deserialize_slice(&data)?.into()))
-                        }
-                        Message::Close(_) => Ok(Some(
-                            Response::Error(ErrorResponse {
-                                error: "WebSocket connection closed".to_string(),
-                            })
-                            .into(),
-                        )),
-                        _ => Ok(None),
-                    };
-                    r
-                });
-
-        Ok(response_stream)
-    }
-
-    /// Subscribes to price feed updates
-    ///
-    /// # Arguments
-    /// * `request` - A subscription request containing feed IDs and parameters
-    pub async fn subscribe(&mut self, request: Request) -> Result<()> {
-        if let Some(sender) = &mut self.ws_sender {
-            let msg = serde_json::to_string(&request)?;
-            sender.send(Message::Text(msg)).await?;
-            Ok(())
-        } else {
-            anyhow::bail!("WebSocket connection not started")
-        }
-    }
-
-    /// Unsubscribes from a previously subscribed feed
-    ///
-    /// # Arguments
-    /// * `subscription_id` - The ID of the subscription to cancel
-    pub async fn unsubscribe(&mut self, subscription_id: SubscriptionId) -> Result<()> {
-        if let Some(sender) = &mut self.ws_sender {
-            let request = Request::Unsubscribe(UnsubscribeRequest { subscription_id });
-            let msg = serde_json::to_string(&request)?;
-            sender.send(Message::Text(msg)).await?;
-            Ok(())
-        } else {
-            anyhow::bail!("WebSocket connection not started")
-        }
-    }
-
-    /// Closes the WebSocket connection
-    pub async fn close(&mut self) -> Result<()> {
-        if let Some(sender) = &mut self.ws_sender {
-            sender.send(Message::Close(None)).await?;
-            self.ws_sender = None;
-            Ok(())
-        } else {
-            anyhow::bail!("WebSocket connection not started")
-        }
-    }
-}
+pub mod client;
+pub mod resilient_ws_connection;
+pub mod ws_connection;