add timeout

Author: keyvankhademi

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

@@ -1,3 +1,5 @@
+use std::time::Duration;
+
 use base64::Engine;
 use pyth_lazer_client::backoff::PythLazerExponentialBackoffBuilder;
 use pyth_lazer_client::client::PythLazerClient;
@@ -41,6 +43,7 @@ async fn main() -> anyhow::Result<()> {
         get_lazer_access_token(),
         4,
         PythLazerExponentialBackoffBuilder::default().build(),
+        Duration::from_secs(5), // Timeout for each connection
     )?;
 
     let stream = client

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

@@ -21,6 +21,7 @@ pub struct PythLazerClient {
     num_connections: usize,
     ws_connections: Vec<PythLazerResilientWSConnection>,
     backoff: ExponentialBackoff,
+    timeout: Duration,
 }
 
 impl PythLazerClient {
@@ -35,6 +36,7 @@ impl PythLazerClient {
         access_token: String,
         num_connections: usize,
         backoff: ExponentialBackoff,
+        timeout: Duration,
     ) -> Result<Self> {
         if backoff.max_elapsed_time.is_some() {
             bail!("max_elapsed_time is not supported in Pyth Lazer client");
@@ -48,6 +50,7 @@ impl PythLazerClient {
             num_connections,
             ws_connections: Vec::with_capacity(num_connections),
             backoff,
+            timeout,
         })
     }
 
@@ -62,6 +65,7 @@ impl PythLazerClient {
                 endpoint,
                 self.access_token.clone(),
                 self.backoff.clone(),
+                self.timeout,
                 ws_connection_sender.clone(),
             );
             self.ws_connections.push(connection);

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

@@ -35,10 +35,12 @@ impl PythLazerResilientWSConnection {
         endpoint: Url,
         access_token: String,
         backoff: ExponentialBackoff,
+        timeout: Duration,
         sender: mpsc::Sender<AnyResponse>,
     ) -> Self {
         let (request_sender, mut request_receiver) = mpsc::channel(CHANNEL_CAPACITY);
-        let mut task = PythLazerResilientWSConnectionTask::new(endpoint, access_token, backoff);
+        let mut task =
+            PythLazerResilientWSConnectionTask::new(endpoint, access_token, backoff, timeout);
 
         tokio::spawn(async move {
             if let Err(e) = task.run(sender, &mut request_receiver).await {
@@ -71,15 +73,22 @@ struct PythLazerResilientWSConnectionTask {
     access_token: String,
     subscriptions: Vec<SubscribeRequest>,
     backoff: ExponentialBackoff,
+    timeout: Duration,
 }
 
 impl PythLazerResilientWSConnectionTask {
-    pub fn new(endpoint: Url, access_token: String, backoff: ExponentialBackoff) -> Self {
+    pub fn new(
+        endpoint: Url,
+        access_token: String,
+        backoff: ExponentialBackoff,
+        timeout: Duration,
+    ) -> Self {
         Self {
             endpoint,
             access_token,
             subscriptions: Vec::new(),
             backoff,
+            timeout,
         }
     }
 
@@ -130,10 +139,12 @@ impl PythLazerResilientWSConnectionTask {
                 .await?;
         }
         loop {
+            let timeout_response = tokio::time::timeout(self.timeout, stream.next());
+
             select! {
-                response = stream.next() => {
+                response = timeout_response => {
                     match response {
-                        Some(response) => match response {
+                        Ok(Some(response)) => match response {
                             Ok(response) => {
                                 sender
                                     .send(response)
@@ -144,9 +155,12 @@ impl PythLazerResilientWSConnectionTask {
                                 bail!("WebSocket stream error: {}", e);
                             }
                         },
-                        None => {
+                        Ok(None) => {
                             bail!("WebSocket stream ended unexpectedly");
                         }
+                        Err(_elapsed) => {
+                            bail!("WebSocket stream timed out");
+                        }
                     }
                 }
                 Some(request) = request_receiver.recv() => {