fix: correctly handle graceful shutdown after losing leader (risingwa…

…velabs#8734)
Little-Wallace · Mar 24, 2023 · 58f29a6 · 58f29a6
1 parent 1c6b246
commit 58f29a6
Showing 1 changed file with 29 additions and 10 deletions.
diff --git a/src/meta/src/lib.rs b/src/meta/src/lib.rs
@@ -254,18 +254,37 @@ pub fn start(opts: MetaNodeOpts) -> Pin<Box<dyn Future<Output = ()> + Send>> {
         .await
         .unwrap();
 
-        let res = tokio::select! {
-            _ = tokio::signal::ctrl_c() => {
-                tracing::info!("receive ctrl+c");
-                shutdown_send.send(()).unwrap();
-                join_handle.await
+        match leader_lost_handle {
+            None => {
+                tokio::select! {
+                    _ = tokio::signal::ctrl_c() => {
+                        tracing::info!("receive ctrl+c");
+                        shutdown_send.send(()).unwrap();
+                        join_handle.await.unwrap()
+                    }
+                    res = &mut join_handle => res.unwrap(),
+                };
+            }
+            Some(mut handle) => {
+                tokio::select! {
+                    _ = &mut handle => {
+                        tracing::info!("receive leader lost signal");
+                        shutdown_send.send(()).unwrap();
+                        join_handle.await.unwrap()
+                    }
+                    _ = tokio::signal::ctrl_c() => {
+                        tracing::info!("receive ctrl+c");
+                        shutdown_send.send(()).unwrap();
+                        join_handle.await.unwrap();
+                        handle.abort();
+                    }
+                    res = &mut join_handle => {
+                        res.unwrap();
+                        handle.abort();
+                    },
+                };
             }
-            res = &mut join_handle => res,
         };
-        res.unwrap();
-        if let Some(leader_lost_handle) = leader_lost_handle {
-            leader_lost_handle.abort();
-        }
     })
 }