fix: graceful shutdown of ws conns (#1473)

Author: iduartgomez

crates/core/src/client_events.rs

@@ -185,16 +185,14 @@ where
                         node_controller.send(NodeEvent::Disconnect { cause: None }).await.ok();
                         anyhow::bail!("shutdown event");
                     }
+                    Err(error) if matches!(error.kind(), ErrorKind::TransportProtocolDisconnect) => {
+                        return Err(anyhow::anyhow!(error));
+                    }
                     Err(error) => {
                         tracing::debug!(%error, "client error");
                         continue;
                     }
                 };
-                // fixme: only allow in certain modes (e.g. while testing)
-                if let ClientRequest::Disconnect { cause } = &*req.request {
-                    node_controller.send(NodeEvent::Disconnect { cause: cause.clone() }).await.ok();
-                    anyhow::bail!("shutdown event");
-                }
                 let cli_id = req.client_id;
                 let res = process_open_request(req, op_manager.clone()).await;
                 results.push(async move {
@@ -207,6 +205,10 @@ where
                             }
                         }
                         Ok(None) => (cli_id, Ok(None)),
+                        Err(Error::Disconnected) => {
+                            tracing::debug!("client disconnected");
+                            (cli_id, Err(ClientError::from(ErrorKind::Disconnect)))
+                        }
                         Err(err) => (cli_id, Err(ErrorKind::OperationError { cause: format!("{err}").into() }.into())),
                     }
                 });
@@ -544,6 +546,9 @@ async fn process_open_request(
 
                 return Ok(Some(Either::Right(callback_rx.unwrap())));
             }
+            ClientRequest::Close => {
+                return Err(Error::Disconnected);
+            }
             _ => {
                 tracing::error!("Op not supported");
             }

crates/core/src/client_events/combinator.rs

@@ -171,7 +171,7 @@ async fn client_fn(
             }
         }
     }
-    tracing::error!("Client shut down");
+    tracing::error!("Peer client interface shut down");
 }
 
 #[cfg(test)]

crates/core/tests/operations.rs

@@ -201,22 +201,19 @@ async fn test_put_contract() -> TestResult {
         make_put(&mut client_api_a, wrapped_state.clone(), contract.clone()).await?;
 
         // Wait for put response
-        loop {
-            let resp = tokio::time::timeout(Duration::from_secs(30), client_api_a.recv()).await;
-            match resp {
-                Ok(Ok(HostResponse::ContractResponse(ContractResponse::PutResponse { key }))) => {
-                    assert_eq!(key, contract_key);
-                    break;
-                }
-                Ok(Ok(other)) => {
-                    tracing::warn!("unexpected response while waiting for put: {:?}", other);
-                }
-                Ok(Err(e)) => {
-                    bail!("Error receiving put response: {}", e);
-                }
-                Err(_) => {
-                    bail!("Timeout waiting for put response");
-                }
+        let resp = tokio::time::timeout(Duration::from_secs(30), client_api_a.recv()).await;
+        match resp {
+            Ok(Ok(HostResponse::ContractResponse(ContractResponse::PutResponse { key }))) => {
+                assert_eq!(key, contract_key);
+            }
+            Ok(Ok(other)) => {
+                tracing::warn!("unexpected response while waiting for put: {:?}", other);
+            }
+            Ok(Err(e)) => {
+                bail!("Error receiving put response: {}", e);
+            }
+            Err(_) => {
+                bail!("Timeout waiting for put response");
             }
         }
 
@@ -362,22 +359,19 @@ async fn test_update_contract() -> TestResult {
         make_put(&mut client_api_a, wrapped_state.clone(), contract.clone()).await?;
 
         // Wait for put response
-        loop {
-            let resp = tokio::time::timeout(Duration::from_secs(30), client_api_a.recv()).await;
-            match resp {
-                Ok(Ok(HostResponse::ContractResponse(ContractResponse::PutResponse { key }))) => {
-                    assert_eq!(key, contract_key, "Contract key mismatch in PUT response");
-                    break;
-                }
-                Ok(Ok(other)) => {
-                    tracing::warn!("unexpected response while waiting for put: {:?}", other);
-                }
-                Ok(Err(e)) => {
-                    bail!("Error receiving put response: {}", e);
-                }
-                Err(_) => {
-                    bail!("Timeout waiting for put response");
-                }
+        let resp = tokio::time::timeout(Duration::from_secs(30), client_api_a.recv()).await;
+        match resp {
+            Ok(Ok(HostResponse::ContractResponse(ContractResponse::PutResponse { key }))) => {
+                assert_eq!(key, contract_key, "Contract key mismatch in PUT response");
+            }
+            Ok(Ok(other)) => {
+                tracing::warn!("unexpected response while waiting for put: {:?}", other);
+            }
+            Ok(Err(e)) => {
+                bail!("Error receiving put response: {}", e);
+            }
+            Err(_) => {
+                bail!("Timeout waiting for put response");
             }
         }