From feb8c2dbfa9613deb78d683a3107c0440d0d4a1c Mon Sep 17 00:00:00 2001
From: Yujia Qiao <rapiz3142@gmail.com>
Date: Mon, 28 Mar 2022 16:42:09 +0800
Subject: [PATCH] fix: restart when heartbeat times out (#147)

---
 src/client.rs | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 657e666..8df8389 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -33,9 +33,11 @@ pub async fn run_client(
     shutdown_rx: broadcast::Receiver<bool>,
     service_rx: mpsc::Receiver<ServiceChange>,
 ) -> Result<()> {
-    let config = config.client.ok_or_else(|| anyhow!(
+    let config = config.client.ok_or_else(|| {
+        anyhow!(
         "Try to run as a client, but the configuration is missing. Please add the `[client]` block"
-    ))?;
+    )
+    })?;
 
     match config.transport.transport_type {
         TransportType::Tcp => {
@@ -459,8 +461,7 @@ impl<T: 'static + Transport> ControlChannel<T> {
                     }
                 },
                 _ = time::sleep(Duration::from_secs(self.heartbeat_timeout)), if self.heartbeat_timeout != 0 => {
-                    warn!("Heartbeat timed out");
-                    break;
+                    return Err(anyhow!("Heartbeat timed out"))
                 }
                 _ = &mut self.shutdown_rx => {
                     break;