fix(backend:rpc): reconnect on connection error (#538)

* fix(rpc): reconnect on connection error

* fix(health): return error status on unhealthy

* fix: undo side-effect on rpc connection check

* fix: into_response

Author: Karrq

backend/lib/src/api/handlers/mod.rs

@@ -32,8 +32,7 @@ pub async fn value_props(State(services): State<Services>) -> Result<impl IntoRe
 
 pub async fn msp_health(State(services): State<Services>) -> Result<impl IntoResponse, Error> {
     debug!("GET health check");
-    let response = services.health.check_health().await;
-    Ok(Json(response))
+    Ok(services.health.check_health().await)
 }
 
 // ==================== Payment Handler ====================

backend/lib/src/data/rpc/client.rs

@@ -27,12 +27,22 @@ impl StorageHubRpcClient {
         self.connection.is_connected().await
     }
 
+    /// Attempts to reconnect if the connection isn't connected
+    async fn ensure_connected(&self) {
+        if !self.is_connected().await {
+            // TODO: More robust reconnection mechanism, like we do for the original connection
+            _ = self.connection.reconnect().await;
+        }
+    }
+
     /// Call a JSON-RPC method on the connected node
     pub async fn call<P, R>(&self, method: &str, params: P) -> RpcResult<R>
     where
         P: ToRpcParams + Send,
         R: DeserializeOwned,
     {
+        self.ensure_connected().await;
+
         self.connection.call(method, params).await
     }
 
@@ -41,6 +51,8 @@ impl StorageHubRpcClient {
     where
         R: DeserializeOwned,
     {
+        self.ensure_connected().await;
+
         self.connection.call_no_params(method).await
     }
 
@@ -54,15 +66,14 @@ impl StorageHubRpcClient {
     pub async fn get_current_price_per_giga_unit_per_tick(&self) -> RpcResult<u128> {
         debug!(target: "rpc::client::get_current_price_per_giga_unit_per_tick", "RPC call: get_current_price_per_giga_unit_per_tick");
 
-        self.connection.call_no_params(methods::CURRENT_PRICE).await
+        self.call_no_params(methods::CURRENT_PRICE).await
     }
 
     /// Returns whether the given `file_key` is expected to be received by the MSP node
     pub async fn is_file_key_expected(&self, file_key: &str) -> RpcResult<bool> {
         debug!(target: "rpc::client::is_file_key_expected", file_key = %file_key, "RPC call: is_file_key_expected");
 
-        self.connection
-            .call(methods::FILE_KEY_EXPECTED, jsonrpsee::rpc_params![file_key])
+        self.call(methods::FILE_KEY_EXPECTED, jsonrpsee::rpc_params![file_key])
             .await
     }
 
@@ -77,12 +88,11 @@ impl StorageHubRpcClient {
         &self,
         file_key: &str,
     ) -> RpcResult<GetFileFromFileStorageResult> {
-        self.connection
-            .call(
-                methods::IS_FILE_IN_FILE_STORAGE,
-                jsonrpsee::rpc_params![file_key],
-            )
-            .await
+        self.call(
+            methods::IS_FILE_IN_FILE_STORAGE,
+            jsonrpsee::rpc_params![file_key],
+        )
+        .await
     }
 
     /// Request the MSP node to export the given `file_key` to the given URL
@@ -94,12 +104,11 @@ impl StorageHubRpcClient {
             "RPC call: save_file_to_disk"
         );
 
-        self.connection
-            .call(
-                methods::SAVE_FILE_TO_DISK,
-                jsonrpsee::rpc_params![file_key, url],
-            )
-            .await
+        self.call(
+            methods::SAVE_FILE_TO_DISK,
+            jsonrpsee::rpc_params![file_key, url],
+        )
+        .await
     }
 
     /// Request the MSP to accept a FileKeyProof (`proof`) for the given `file_key`
@@ -111,33 +120,32 @@ impl StorageHubRpcClient {
             "RPC call: receive_file_chunks"
         );
 
-        self.connection
-            .call(
-                methods::RECEIVE_FILE_CHUNKS,
-                jsonrpsee::rpc_params![file_key, proof],
-            )
-            .await
+        self.call(
+            methods::RECEIVE_FILE_CHUNKS,
+            jsonrpsee::rpc_params![file_key, proof],
+        )
+        .await
     }
 
     /// Retrieve the Onchain Provider ID of the MSP Node (therefore the MSP ID)
     pub async fn get_provider_id(&self) -> RpcResult<RpcProviderId> {
         debug!(target: "rpc::client::get_provider_id", "RPC call: get_provider_id");
 
-        self.connection.call_no_params(methods::PROVIDER_ID).await
+        self.call_no_params(methods::PROVIDER_ID).await
     }
 
     /// Retrieve the list of value propositions of the MSP Node
     pub async fn get_value_props(&self) -> RpcResult<GetValuePropositionsResult> {
         debug!(target: "rpc::client::get_value_props", "RPC call: get_value_props");
 
-        self.connection.call_no_params(methods::VALUE_PROPS).await
+        self.call_no_params(methods::VALUE_PROPS).await
     }
 
     /// Retrieve the list of multiaddresses associated with the MSP Node
     pub async fn get_multiaddresses(&self) -> RpcResult<Vec<String>> {
         debug!(target: "rpc::client::get_multiaddresses", "RPC call: get_multiaddresses");
 
-        self.connection.call_no_params(methods::PEER_IDS).await
+        self.call_no_params(methods::PEER_IDS).await
     }
 }
 
@@ -175,6 +183,27 @@ mod tests {
         StorageHubRpcClient::new(connection)
     }
 
+    #[tokio::test]
+    async fn reconnect_automatically() {
+        let conn = MockConnection::new();
+        conn.disconnect().await;
+        let conn = Arc::new(AnyRpcConnection::Mock(conn));
+        let client = StorageHubRpcClient::new(conn);
+
+        assert!(
+            !client.is_connected().await,
+            "Should not be connected initially"
+        );
+
+        let result = client.get_provider_id().await;
+        assert!(
+            result.is_ok(),
+            "Should reconnect and be able to retrieve provider id"
+        );
+
+        assert!(client.is_connected().await, "Should be connected now");
+    }
+
     #[tokio::test]
     async fn get_current_price_per_unit_per_tick() {
         let client = mock_rpc();

backend/lib/src/data/rpc/connection/mod.rs

@@ -71,11 +71,11 @@ impl RpcConnection for AnyRpcConnection {
         }
     }
 
-    async fn close(&self) -> RpcResult<()> {
+    async fn reconnect(&self) -> RpcResult<()> {
         match self {
-            AnyRpcConnection::Real(conn) => conn.close().await,
+            AnyRpcConnection::Real(conn) => conn.reconnect().await,
             #[cfg(feature = "mocks")]
-            AnyRpcConnection::Mock(conn) => conn.close().await,
+            AnyRpcConnection::Mock(conn) => conn.reconnect().await,
         }
     }
 }

backend/lib/src/data/rpc/mock_connection.rs

@@ -102,7 +102,7 @@ impl MockConnection {
     }
 
     /// Simulate reconnection
-    pub async fn reconnect(&self) {
+    pub async fn connect(&self) {
         let mut connected = self.connected.write().await;
         *connected = true;
     }
@@ -291,8 +291,8 @@ impl RpcConnection for MockConnection {
         *connected
     }
 
-    async fn close(&self) -> RpcResult<()> {
-        self.disconnect().await;
+    async fn reconnect(&self) -> RpcResult<()> {
+        self.connect().await;
         Ok(())
     }
 }
@@ -325,8 +325,8 @@ mod tests {
         // Test connection status
         assert!(conn.is_connected().await);
 
-        // Test close
-        conn.close().await.unwrap();
+        // Test disconnection
+        conn.disconnect().await;
         assert!(!conn.is_connected().await);
     }
 
@@ -434,9 +434,6 @@ mod tests {
     async fn test_connection_disconnect_reconnect() {
         let conn = MockConnection::new();
 
-        // Initially connected
-        assert!(conn.is_connected().await);
-
         // Disconnect
         conn.disconnect().await;
         assert!(!conn.is_connected().await);
@@ -446,7 +443,7 @@ mod tests {
         assert!(matches!(result, Err(RpcConnectionError::ConnectionClosed)));
 
         // Reconnect
-        conn.reconnect().await;
+        conn.connect().await;
         assert!(conn.is_connected().await);
 
         // Call should work now

backend/lib/src/data/rpc/mod.rs

@@ -43,6 +43,6 @@ pub trait RpcConnection: Send + Sync {
     /// Check if the connection is currently active
     async fn is_connected(&self) -> bool;
 
-    /// Close the connection gracefully
-    async fn close(&self) -> RpcResult<()>;
+    /// Create a new instance of the connection, closing the current one gracefully
+    async fn reconnect(&self) -> RpcResult<()>;
 }

backend/lib/src/data/rpc/ws_connection.rs

@@ -128,9 +128,10 @@ impl RpcConnection for WsConnection {
         }
     }
 
-    async fn close(&self) -> RpcResult<()> {
-        self.client.write().await.take();
+    async fn reconnect(&self) -> RpcResult<()> {
+        let new_client = Self::build_client(&self.config).await?;
 
+        self.client.write().await.replace(Arc::new(new_client));
         Ok(())
     }
 }

backend/lib/src/services/health.rs

@@ -4,6 +4,12 @@
 
 use std::sync::Arc;
 
+use axum::{
+    body::Body,
+    http::StatusCode,
+    response::{IntoResponse, Response},
+    Json,
+};
 use serde::Serialize;
 use shc_rpc::RpcProviderId;
 use tracing::{debug, error};
@@ -18,6 +24,24 @@ pub struct DetailedHealthStatus {
     pub components: HealthComponents,
 }
 
+impl DetailedHealthStatus {
+    pub fn is_healthy(&self) -> bool {
+        self.status == HealthService::HEALTHY
+    }
+}
+
+impl IntoResponse for DetailedHealthStatus {
+    fn into_response(self) -> Response<Body> {
+        let status = if self.is_healthy() {
+            StatusCode::OK
+        } else {
+            StatusCode::SERVICE_UNAVAILABLE
+        };
+
+        (status, Json(self)).into_response()
+    }
+}
+
 #[derive(Serialize)]
 pub struct HealthComponents {
     pub storage: ComponentHealth,