retry logic for dropped ws

drift-labs · Mar 22, 2024 · d2582c1 · d2582c1
1 parent 058cd51
commit d2582c1
Show file tree

Hide file tree

Showing 6 changed files with 105 additions and 149 deletions.
diff --git a/.DS_Store b/.DS_Store
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -41,7 +41,6 @@ tokio-tungstenite = { version = "0.21.0", features = ["native-tls"] }
 regex = "1.10.2"
 dashmap = "5.5.3"
 rayon = "1.9.0"
-clippy = "0.0.302"
 
 [dev-dependencies]
 pyth = { git = "https://github.com/drift-labs/protocol-v2.git", tag = "v2.67.0", features = [

diff --git a/src/types.rs b/src/types.rs
@@ -249,6 +249,8 @@ pub enum SdkError {
     BorrowError(#[from] BorrowError),
     #[error("{0}")]
     Generic(String),
+    #[error("max connection attempts reached")]
+    MaxReconnectionAttemptsReached,
 }
 
 impl SdkError {

diff --git a/src/websocket_account_subscriber.rs b/src/websocket_account_subscriber.rs
@@ -76,49 +76,69 @@ impl WebsocketAccountSubscriber {
         let (unsub_tx, mut unsub_rx) = tokio::sync::mpsc::channel::<()>(1);
         self.unsubscriber = Some(unsub_tx);
 
+        let mut attempt = 0;
+        let max_reconnection_attempts = 20;
+        let base_delay = tokio::time::Duration::from_secs(2);
+
         tokio::spawn({
             let event_emitter = self.event_emitter.clone();
             let mut latest_slot = 0;
-            let susbscription_name = self.subscription_name;
+            let subscription_name = self.subscription_name;
             let pubkey = self.pubkey.clone();
             async move {
-                let (mut account_updates, account_unsubscribe) = pubsub
-                    .account_subscribe(&pubkey, Some(account_config))
-                    .await
-                    .unwrap();
-
                 loop {
-                    tokio::select! {
-                        message = account_updates.next() => {
-                            match message {
-                                Some(message) => {
-                                    let slot = message.context.slot;
-                                    if slot >= latest_slot {
-                                        latest_slot = slot;
-                                        let account_update = AccountUpdate {
-                                            pubkey: pubkey.to_string(),
-                                            data: message.value,
-                                            slot,
-                                        };
-                                        event_emitter.emit(susbscription_name, Box::new(account_update));
+                    let (mut account_updates, account_unsubscribe) = pubsub
+                        .account_subscribe(&pubkey, Some(account_config.clone()))
+                        .await
+                        .unwrap();
+
+                    loop {
+                        tokio::select! {
+                            message = account_updates.next() => {
+                                match message {
+                                    Some(message) => {
+                                        let slot = message.context.slot;
+                                        if slot >= latest_slot {
+                                            latest_slot = slot;
+                                            let account_update = AccountUpdate {
+                                                pubkey: pubkey.to_string(),
+                                                data: message.value,
+                                                slot,
+                                            };
+                                            event_emitter.emit(subscription_name, Box::new(account_update));
+                                        }
+                                    }
+                                    None => {
+                                        log::warn!("{}: Account stream interrupted", subscription_name);
+                                        account_unsubscribe().await;
+                                        break;
                                     }
                                 }
-                                None => {
-                                    log::warn!("{}: Account stream interrupted", susbscription_name);
+                            }
+                            unsub = unsub_rx.recv() => {
+                                if let Some(_) = unsub {
+                                    log::debug!("{}: Unsubscribing from account stream", subscription_name);
                                     account_unsubscribe().await;
-                                    break;
+                                    return Ok(());
+
                                 }
                             }
                         }
-                        unsub = unsub_rx.recv() => {
-                            if let Some(_) = unsub {
-                                log::debug!("{}: Unsubscribing from account stream", susbscription_name);
-                                account_unsubscribe().await;
-                                break;
+                    }
 
-                            }
-                        }
+                    if attempt >= max_reconnection_attempts {
+                        log::error!("{}: Max reconnection attempts reached", subscription_name);
+                        return Err(crate::SdkError::MaxReconnectionAttemptsReached);
                     }
+
+                    let delay_duration = base_delay * 2_u32.pow(attempt);
+                    log::debug!(
+                        "{}: Reconnecting in {:?}",
+                        subscription_name,
+                        delay_duration
+                    );
+                    tokio::time::sleep(delay_duration).await;
+                    attempt += 1;
                 }
             }
         });