From 4cc20ae2fe5598374fd4fa878980efad632a60d3 Mon Sep 17 00:00:00 2001
From: 0x5459 <0x5459@protonmail.com>
Date: Wed, 8 May 2024 23:59:00 +0800
Subject: [PATCH] feat: Add request_timeout to rcp_client

---
 crates/cli/src/main.rs             | 14 +++++++++--
 crates/node/src/rpc_client/mod.rs  | 37 ++++++++++++++++++++++++++----
 crates/tests/e2e-tests/src/main.rs |  4 ++--
 3 files changed, 47 insertions(+), 8 deletions(-)

diff --git a/crates/cli/src/main.rs b/crates/cli/src/main.rs
index 4c278a4d..7b3bbda8 100644
--- a/crates/cli/src/main.rs
+++ b/crates/cli/src/main.rs
@@ -1,13 +1,14 @@
 use clap::Parser;
 use clap::Subcommand;
 use clap_num::number_range;
-use gevulot_node::rpc_client::RpcClient;
+use gevulot_node::rpc_client::RpcClientBuilder;
 use gevulot_node::types::program::ResourceRequest;
 use gevulot_node::types::Hash;
 use gevulot_node::types::TransactionTree;
 use libsecp256k1::PublicKey;
 use std::net::SocketAddr;
 use std::path::PathBuf;
+use std::time::Duration;
 
 #[derive(Parser, Debug)]
 #[clap(author = "Gevulot Team", version, about, long_about = None)]
@@ -20,6 +21,9 @@ pub struct ArgConfiguration {
         value_name = "URL"
     )]
     json_url: String,
+    /// Timeout duration for rpc request (in seconds).
+    #[clap(long = "rpctimeout", value_name = "RPC TIMEOUT")]
+    rpc_timeout: Option<u64>,
     /// Private key file path to sign Tx.
     #[clap(
         short,
@@ -185,7 +189,13 @@ async fn main() {
 
     let args = ArgConfiguration::parse();
 
-    let client = RpcClient::new(args.json_url);
+    let mut client_builder = RpcClientBuilder::default();
+    if let Some(rpc_timeout) = args.rpc_timeout {
+        client_builder = client_builder.request_timeout(Duration::from_secs(rpc_timeout));
+    }
+    let client = client_builder
+        .build(args.json_url)
+        .expect("build rpc client");
 
     match args.command {
         ConfCommands::GenerateKey => match gevulot_cli::keyfile::create_key_file(&args.keyfile) {
diff --git a/crates/node/src/rpc_client/mod.rs b/crates/node/src/rpc_client/mod.rs
index 6255f6fc..d7279533 100644
--- a/crates/node/src/rpc_client/mod.rs
+++ b/crates/node/src/rpc_client/mod.rs
@@ -9,17 +9,46 @@ use jsonrpsee::{
     http_client::{HttpClient, HttpClientBuilder},
 };
 use std::error::Error;
+use std::time::Duration;
+
+/// A RPC client builder for connecting to the Gevulot network
+#[derive(Debug)]
+pub struct RpcClientBuilder {
+    request_timeout: Duration,
+}
+
+impl Default for RpcClientBuilder {
+    fn default() -> Self {
+        Self {
+            request_timeout: Duration::from_secs(60),
+        }
+    }
+}
+
+impl RpcClientBuilder {
+    /// Set request timeout (default is 60 seconds).
+    pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
+        self.request_timeout = request_timeout;
+        self
+    }
+
+    /// Returns a [RpcClient] connected to the Gevulot network running at the URI provided.
+    pub fn build(self, url: impl AsRef<str>) -> Result<RpcClient, Box<dyn Error>> {
+        let client = HttpClientBuilder::default()
+            .request_timeout(self.request_timeout)
+            .build(url)?;
+        Ok(RpcClient { client })
+    }
+}
 
 pub struct RpcClient {
     client: HttpClient,
 }
 
 impl RpcClient {
+    #[deprecated(note = "Please use `RpcClientBuilder` instead.")]
     pub fn new(url: impl AsRef<str>) -> Self {
-        let client = HttpClientBuilder::default()
-            .build(url)
-            .expect("http client");
-        RpcClient { client }
+        RpcClientBuilder::default().build(url).expect("http client")
     }
 
     pub async fn send_transaction(&self, tx: &Transaction<Created>) -> Result<(), Box<dyn Error>> {
diff --git a/crates/tests/e2e-tests/src/main.rs b/crates/tests/e2e-tests/src/main.rs
index c16c2856..1db6b78f 100644
--- a/crates/tests/e2e-tests/src/main.rs
+++ b/crates/tests/e2e-tests/src/main.rs
@@ -1,6 +1,6 @@
 use clap::Parser;
 use gevulot_node::{
-    rpc_client::RpcClient,
+    rpc_client::{RpcClient, RpcClientBuilder},
     types::{
         program::ResourceRequest,
         transaction::{Payload, ProgramData, ProgramMetadata, Workflow, WorkflowStep},
@@ -39,7 +39,7 @@ pub struct ArgConfiguration {
 #[tokio::main]
 async fn main() -> Result<()> {
     let cfg = ArgConfiguration::parse();
-    let client = RpcClient::new(cfg.json_rpc_url);
+    let client = RpcClientBuilder::default().build(cfg.json_rpc_url)?;
     let file_server = Arc::new(FileServer::new(cfg.listen_addr).await);
 
     let bs = std::fs::read(cfg.key_file)?;