Procházet zdrojové kódy

drk: gracefully handle rpc disconnects with a resetable client

skoupidi před 9 měsíci
rodič
revize
035a53c500
3 změnil soubory, kde provedl 51 přidání a 27 odebrání
  1. 9 21
      bin/drk/src/interactive.rs
  2. 4 3
      bin/drk/src/lib.rs
  3. 38 3
      bin/drk/src/rpc.rs

+ 9 - 21
bin/drk/src/interactive.rs

@@ -536,13 +536,11 @@ pub async fn interactive(drk: &DrkPtr, endpoint: &Url, history_path: &str, ex: &
                         &subscription_tasks,
                         &shell_sender,
                         ex,
-                        &mut output,
                     )
                     .await
                 }
                 "unsubscribe" => {
-                    handle_unsubscribe(&mut subscription_active, &subscription_tasks, &mut output)
-                        .await
+                    handle_unsubscribe(&mut subscription_active, &subscription_tasks).await
                 }
                 "snooze" => snooze_active = true,
                 "unsnooze" => snooze_active = false,
@@ -610,10 +608,8 @@ pub async fn interactive(drk: &DrkPtr, endpoint: &Url, history_path: &str, ex: &
     }
 
     // Stop the subscription tasks if they are active
-    if subscription_active {
-        subscription_tasks[0].stop().await;
-        subscription_tasks[1].stop().await;
-    }
+    subscription_tasks[0].stop_nowait();
+    subscription_tasks[1].stop_nowait();
 
     // Write history file
     let _ = linenoise_history_save(history_file);
@@ -2353,12 +2349,11 @@ async fn handle_subscribe(
     subscription_tasks: &[StoppableTaskPtr; 2],
     shell_sender: &Sender<Vec<String>>,
     ex: &ExecutorPtr,
-    output: &mut Vec<String>,
 ) {
-    if *subscription_active {
-        output.push(String::from("Subscription is already active!"));
-        return
-    }
+    // Kill zombie tasks if they failed
+    subscription_tasks[0].stop_nowait();
+    subscription_tasks[1].stop_nowait();
+    *subscription_active = true;
 
     // Start the subcristion task
     let drk_ = drk.clone();
@@ -2377,22 +2372,15 @@ async fn handle_subscribe(
         Error::DetachedTaskStopped,
         ex.clone(),
     );
-
-    *subscription_active = true;
 }
 
 /// Auxiliary function to define the unsubscribe command handling.
 async fn handle_unsubscribe(
     subscription_active: &mut bool,
     subscription_tasks: &[StoppableTaskPtr; 2],
-    output: &mut Vec<String>,
 ) {
-    if !*subscription_active {
-        output.push(String::from("Subscription is already inactive!"));
-        return
-    }
-    subscription_tasks[0].stop().await;
-    subscription_tasks[1].stop().await;
+    subscription_tasks[0].stop_nowait();
+    subscription_tasks[1].stop_nowait();
     *subscription_active = false;
 }
 

+ 4 - 3
bin/drk/src/lib.rs

@@ -21,7 +21,7 @@ use std::{fs::create_dir_all, sync::Arc};
 use smol::lock::RwLock;
 use url::Url;
 
-use darkfi::{rpc::client::RpcClient, system::ExecutorPtr, util::path::expand_path, Error, Result};
+use darkfi::{system::ExecutorPtr, util::path::expand_path, Error, Result};
 
 /// Error codes
 pub mod error;
@@ -29,6 +29,7 @@ use error::{WalletDbError, WalletDbResult};
 
 /// darkfid JSON-RPC related methods
 pub mod rpc;
+use rpc::DarkfidRpcClient;
 
 /// Payment methods
 pub mod transfer;
@@ -78,7 +79,7 @@ pub struct Drk {
     /// Wallet database operations handler
     pub wallet: WalletPtr,
     /// JSON-RPC client to execute requests to darkfid daemon
-    pub rpc_client: Option<RpcClient>,
+    pub rpc_client: Option<RwLock<DarkfidRpcClient>>,
     /// Flag indicating if fun stuff are enabled
     pub fun: bool,
 }
@@ -112,7 +113,7 @@ impl Drk {
 
         // Initialize rpc client
         let rpc_client = if let Some(endpoint) = endpoint {
-            Some(RpcClient::new(endpoint, ex.clone()).await?)
+            Some(RwLock::new(DarkfidRpcClient::new(endpoint, ex.clone()).await))
         } else {
             None
         };

+ 38 - 3
bin/drk/src/rpc.rs

@@ -59,6 +59,28 @@ use crate::{
     Drk, DrkPtr,
 };
 
+/// Structure to hold a JSON-RPC client and its config,
+/// so we can recreate it in case of an error.
+pub struct DarkfidRpcClient {
+    endpoint: Url,
+    ex: ExecutorPtr,
+    client: Option<RpcClient>,
+}
+
+impl DarkfidRpcClient {
+    pub async fn new(endpoint: Url, ex: ExecutorPtr) -> Self {
+        let client = RpcClient::new(endpoint.clone(), ex.clone()).await.ok();
+        Self { endpoint, ex, client }
+    }
+
+    /// Stop the client.
+    pub async fn stop(&self) {
+        if let Some(ref client) = self.client {
+            client.stop().await
+        }
+    }
+}
+
 /// Auxiliary structure holding various in memory caches to use during scan
 pub struct ScanCache {
     /// The Money Merkle tree containing coins
@@ -559,15 +581,28 @@ impl Drk {
         params: &JsonValue,
     ) -> Result<JsonValue> {
         let Some(ref rpc_client) = self.rpc_client else { return Err(Error::RpcClientStopped) };
+        let mut lock = rpc_client.write().await;
+        let Some(ref client) = lock.client else { return Err(Error::RpcClientStopped) };
         let req = JsonRequest::new(method, params.clone());
-        let rep = rpc_client.request(req).await?;
+
+        // Execute request
+        if let Ok(rep) = client.request(req.clone()).await {
+            drop(lock);
+            return Ok(rep)
+        }
+
+        // Reset the rpc client in case of an error and try again
+        let client = RpcClient::new(lock.endpoint.clone(), lock.ex.clone()).await?;
+        let rep = client.request(req).await?;
+        lock.client = Some(client);
+        drop(lock);
         Ok(rep)
     }
 
     /// Auxiliary function to stop current JSON-RPC client, if its initialized.
     pub async fn stop_rpc_client(&self) -> Result<()> {
         if let Some(ref rpc_client) = self.rpc_client {
-            rpc_client.stop().await;
+            rpc_client.read().await.stop().await;
         };
         Ok(())
     }
@@ -621,6 +656,7 @@ pub async fn subscribe_blocks(
             return Err(Error::Custom(err_msg))
         }
     };
+    drop(lock);
 
     // Check if other blocks have been created
     if last_confirmed_height != last_scanned_height || last_confirmed_hash != last_scanned_hash {
@@ -670,7 +706,6 @@ pub async fn subscribe_blocks(
     shell_message.push(String::from("Detached subscription to background"));
     shell_message.push(String::from("All is good. Waiting for block notifications..."));
     shell_sender.send(shell_message).await?;
-    drop(lock);
 
     let e = 'outer: loop {
         match subscription.receive().await {