From f72932d84bcdfa5f8a00838785bd985dd1fcc0cf Mon Sep 17 00:00:00 2001
From: gabriele-0201 <gabriele.miotti2@gmail.com>
Date: Thu, 11 Jan 2024 17:16:05 +0100
Subject: [PATCH] shim: add no_retry flag

implement Connector::try_connect to handle a single connetion attempt
---
 sugondat/shim/src/cli.rs               |  6 ++++++
 sugondat/shim/src/cmd/query/mod.rs     |  2 +-
 sugondat/shim/src/cmd/serve.rs         |  6 +++---
 sugondat/shim/src/sugondat_rpc/conn.rs | 25 +++++++++++++++++++++++++
 sugondat/shim/src/sugondat_rpc/mod.rs  |  9 +++++++--
 5 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/sugondat/shim/src/cli.rs b/sugondat/shim/src/cli.rs
index 89776241..586c67f0 100644
--- a/sugondat/shim/src/cli.rs
+++ b/sugondat/shim/src/cli.rs
@@ -87,6 +87,12 @@ pub struct SugondatRpcParams {
     /// The address of the sugondat-node to connect to.
     #[clap(long, default_value = "ws://localhost:9988", env = ENV_SUGONDAT_NODE_URL)]
     pub node_url: String,
+
+    /// By default the first connection to the node is retried until it is properly connected.
+    ///
+    /// This flag avoids this behavior by attempting to connect only once
+    #[clap(long)]
+    pub no_retry: bool,
 }
 
 impl DockParams {
diff --git a/sugondat/shim/src/cmd/query/mod.rs b/sugondat/shim/src/cmd/query/mod.rs
index d17b2afc..7b8980e8 100644
--- a/sugondat/shim/src/cmd/query/mod.rs
+++ b/sugondat/shim/src/cmd/query/mod.rs
@@ -19,5 +19,5 @@ pub async fn run(params: Params) -> anyhow::Result<()> {
 async fn connect_rpc(
     conn_params: crate::cli::SugondatRpcParams,
 ) -> anyhow::Result<sugondat_rpc::Client> {
-    sugondat_rpc::Client::new(conn_params.node_url).await
+    sugondat_rpc::Client::new(conn_params.node_url, conn_params.no_retry).await
 }
diff --git a/sugondat/shim/src/cmd/serve.rs b/sugondat/shim/src/cmd/serve.rs
index 83dfb639..8f4e6452 100644
--- a/sugondat/shim/src/cmd/serve.rs
+++ b/sugondat/shim/src/cmd/serve.rs
@@ -16,7 +16,7 @@ Pass --submit-dev-alice or --submit-private-key=<..> to fix."
         );
     }
     let server = Server::builder().build(listen_on).await?;
-    let client = connect_client(&params.rpc.node_url).await?;
+    let client = connect_client(&params.rpc.node_url, params.rpc.no_retry).await?;
     let methods = dock::init(dock::Config {
         // TODO: whenever there are more docks, the logic of checking if any at least one is enabled
         //       and other similar stuff should be in CLI.
@@ -30,7 +30,7 @@ Pass --submit-dev-alice or --submit-private-key=<..> to fix."
     Ok(())
 }
 
-async fn connect_client(url: &str) -> anyhow::Result<Client> {
-    let client = Client::new(url.to_string()).await?;
+async fn connect_client(url: &str, no_retry: bool) -> anyhow::Result<Client> {
+    let client = Client::new(url.to_string(), no_retry).await?;
     Ok(client)
 }
diff --git a/sugondat/shim/src/sugondat_rpc/conn.rs b/sugondat/shim/src/sugondat_rpc/conn.rs
index db0951cb..4c3b039c 100644
--- a/sugondat/shim/src/sugondat_rpc/conn.rs
+++ b/sugondat/shim/src/sugondat_rpc/conn.rs
@@ -87,6 +87,31 @@ impl Connector {
         }
     }
 
+    /// Attempt a single connection. Returns the connection handle if successful, otherwise returns the reason for failure
+    pub async fn try_connect(&self) -> anyhow::Result<Arc<Conn>> {
+        let mut state = self.state.lock().await;
+        match &mut *state {
+            State::Connected(conn) => Ok(conn.clone()),
+            State::Connecting { .. } => Err(anyhow::anyhow!(
+                "The connector is already attempting to connect"
+            )),
+            State::Disconnected => {
+                let conn_id = self.gen_conn_id();
+                let rpc_url = self.rpc_url.clone();
+                match Conn::connect(conn_id, &rpc_url).await {
+                    Ok(conn) => {
+                        *state = State::Connected(conn.clone());
+                        Ok(conn)
+                    }
+                    Err(e) => Err(anyhow::anyhow!(
+                        "failed to connect to sugondat node: {}\n",
+                        e
+                    )),
+                }
+            }
+        }
+    }
+
     /// Makes sure that the client is connected. Returns the connection handle.
     pub async fn ensure_connected(&self) -> Arc<Conn> {
         let mut state = self.state.lock().await;
diff --git a/sugondat/shim/src/sugondat_rpc/mod.rs b/sugondat/shim/src/sugondat_rpc/mod.rs
index 2baa6966..5adee878 100644
--- a/sugondat/shim/src/sugondat_rpc/mod.rs
+++ b/sugondat/shim/src/sugondat_rpc/mod.rs
@@ -38,7 +38,7 @@ impl Client {
     /// The RPC URL must be a valid URL pointing to a sugondat node. If it's not a malformed URL,
     /// returns an error.
     #[tracing::instrument(level = Level::DEBUG)]
-    pub async fn new(rpc_url: String) -> anyhow::Result<Self> {
+    pub async fn new(rpc_url: String, no_retry: bool) -> anyhow::Result<Self> {
         anyhow::ensure!(
             url::Url::parse(&rpc_url).is_ok(),
             "invalid RPC URL: {}",
@@ -50,7 +50,12 @@ impl Client {
         let me = Self {
             connector: Arc::new(conn::Connector::new(rpc_url)),
         };
-        me.connector.ensure_connected().await;
+
+        match no_retry {
+            true => me.connector.try_connect().await?,
+            false => me.connector.ensure_connected().await,
+        };
+
         Ok(me)
     }