fix: handle heartbeat shutdown gracefully

Signed-off-by: tison <[email protected]>
GreptimeTeam · Dec 6, 2023 · a4a4682 · a4a4682
1 parent ea80570
commit a4a4682
Showing 1 changed file with 12 additions and 6 deletions.
diff --git a/src/meta-srv/src/service/heartbeat.rs b/src/meta-srv/src/service/heartbeat.rs
@@ -55,10 +55,13 @@ impl heartbeat_server::Heartbeat for MetaSrv {
                             Some(header) => header,
                             None => {
                                 let err = error::MissingRequestHeaderSnafu {}.build();
-                                tx.send(Err(err.into())).await.expect("working rx");
-                                break;
+                                if tx.send(Err(err.into())).await.is_err() {
+                                    // response was dropped; shutting down
+                                    break;
+                                }
                             }
                         };
+
                         debug!("Receiving heartbeat request: {:?}", req);
 
                         if pusher_key.is_none() {
@@ -78,7 +81,10 @@ impl heartbeat_server::Heartbeat for MetaSrv {
                         is_not_leader = res.as_ref().map_or(false, |r| r.is_not_leader());
 
                         debug!("Sending heartbeat response: {:?}", res);
-                        tx.send(res).await.expect("working rx");
+                        if tx.send(res).await.is_err() {
+                            // response was dropped; shutting down
+                            break;
+                        }
                     }
                     Err(err) => {
                         if let Some(io_err) = error::match_for_io_error(&err) {
@@ -89,9 +95,9 @@ impl heartbeat_server::Heartbeat for MetaSrv {
                             }
                         }
 
-                        match tx.send(Err(err)).await {
-                            Ok(_) => (),
-                            Err(_err) => break, // response was dropped
+                        if tx.send(Err(err)).await.is_err() {
+                            // response was dropped; shutting down
+                            break;
                         }
                     }
                 }