wip: add wait until connectec to cometbft-rpc

unionlabs · Sep 12, 2024 · 998c980 · 998c980
1 parent a1fc398
commit 998c980
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 0 deletions.
diff --git a/lib/cometbft-rpc/src/lib.rs b/lib/cometbft-rpc/src/lib.rs
@@ -2,6 +2,7 @@ use core::fmt;
 use std::{
     fmt::Debug,
     num::{NonZeroU32, NonZeroU64, NonZeroU8},
+    time::Duration,
 };
 
 use ::serde::de::DeserializeOwned;
@@ -51,6 +52,12 @@ impl Client {
                         .instrument(debug_span!("cometbft_rpc_client", %url))
                 });
 
+                // TODO: Config
+                client
+                    .wait_until_connected(Duration::from_secs(5))
+                    .await
+                    .map_err(|e| JsonRpcError::Custom(e.to_string()))?;
+
                 ClientInner::Ws(client)
             }
             Some(("http" | "https", _)) => {

diff --git a/lib/reconnecting-jsonrpc-ws-client/src/lib.rs b/lib/reconnecting-jsonrpc-ws-client/src/lib.rs
@@ -110,12 +110,35 @@ impl Client {
         self.healthy.load(Ordering::SeqCst)
     }
 
+    pub async fn wait_until_connected(
+        &self,
+        timeout: Duration,
+    ) -> Result<(), ConnectionTimeoutError> {
+        tokio::time::timeout(timeout, async {
+            loop {
+                if self.is_healthy() {
+                    return;
+                } else {
+                    sleep(Duration::from_secs(1)).await
+                }
+            }
+        })
+        .await
+        .map_err(|_| ConnectionTimeoutError { timeout })
+    }
+
     pub fn shutdown(&self) {
         self.cancellation_token.cancel();
         self.handle.abort()
     }
 }
 
+#[derive(Debug, thiserror::Error)]
+#[error("websocket connection timed out after {}.{}s", timeout.as_secs(), timeout.subsec_nanos())]
+pub struct ConnectionTimeoutError {
+    pub timeout: Duration,
+}
+
 impl Drop for Client {
     fn drop(&mut self) {
         self.cancellation_token.cancel();