From f48464eee9ac83a7786df3a3f78285465f805a3f Mon Sep 17 00:00:00 2001
From: iGxnon <igxnon@gmail.com>
Date: Mon, 17 Jul 2023 14:23:46 +0800
Subject: [PATCH] feat: support grpc health checking protocol

Signed-off-by: iGxnon <igxnon@gmail.com>
---
 Cargo.lock                       | 15 +++++++++++++++
 xline/Cargo.toml                 |  1 +
 xline/src/server/xline_server.rs | 13 ++++++++++++-
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index 154f7e77d..ce446c890 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2598,6 +2598,20 @@ dependencies = [
  "syn 1.0.107",
 ]
 
+[[package]]
+name = "tonic-health"
+version = "0.8.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a88aee666ef3a4d1ee46218bbc8e5f69bcf9cc27bf2e871d6b724d83f56d179f"
+dependencies = [
+ "async-stream",
+ "bytes",
+ "prost",
+ "tokio",
+ "tokio-stream 0.1.12 (registry+https://github.com/rust-lang/crates.io-index)",
+ "tonic",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"
@@ -3158,6 +3172,7 @@ dependencies = [
  "tokio-stream 0.1.12 (git+https://github.com/madsim-rs/tokio.git?rev=ab251ad)",
  "tokio-util",
  "toml 0.5.10",
+ "tonic-health",
  "tracing",
  "tracing-appender",
  "tracing-opentelemetry",
diff --git a/xline/Cargo.toml b/xline/Cargo.toml
index ddab8afe6..0609e779d 100644
--- a/xline/Cargo.toml
+++ b/xline/Cargo.toml
@@ -47,6 +47,7 @@ tokio = { version = "0.2", package = "madsim-tokio", features = [
 ] }
 tokio-stream = { git = "https://github.com/madsim-rs/tokio.git", rev = "ab251ad" }
 tonic = { version = "0.2", package = "madsim-tonic" }
+tonic-health = "0.8.0"
 tracing = "0.1.37"
 tracing-opentelemetry = "0.18.0"
 tracing-subscriber = { version = "0.3.16", features = ["env-filter"] }
diff --git a/xline/src/server/xline_server.rs b/xline/src/server/xline_server.rs
index 26a327b73..9bc4b1084 100644
--- a/xline/src/server/xline_server.rs
+++ b/xline/src/server/xline_server.rs
@@ -8,6 +8,7 @@ use jsonwebtoken::{DecodingKey, EncodingKey};
 use tokio::net::TcpListener;
 use tokio_stream::wrappers::TcpListenerStream;
 use tonic::transport::Server;
+use tonic_health::ServingStatus;
 use utils::config::{ClientTimeout, CurpConfig, ServerTimeout, StorageConfig};
 
 use super::{
@@ -85,7 +86,7 @@ impl XlineServer {
             client_timeout,
             storage_cfg: storage_config,
             server_timeout,
-            shutdown_trigger: Arc::new(event_listener::Event::new()),
+            shutdown_trigger: Arc::new(Event::new()),
         }
     }
 
@@ -195,6 +196,10 @@ impl XlineServer {
             maintenance_server,
             curp_server,
         ) = self.init_servers(persistent, key_pair).await?;
+        let (mut reporter, health_server) = tonic_health::server::health_reporter();
+        reporter
+            .set_service_status("", ServingStatus::Serving)
+            .await;
         Ok(Server::builder()
             .add_service(RpcLockServer::new(lock_server))
             .add_service(RpcKvServer::new(kv_server))
@@ -203,6 +208,7 @@ impl XlineServer {
             .add_service(RpcWatchServer::new(watch_server))
             .add_service(RpcMaintenanceServer::new(maintenance_server))
             .add_service(ProtocolServer::new(curp_server))
+            .add_service(health_server)
             .serve(addr)
             .await?)
     }
@@ -232,6 +238,10 @@ impl XlineServer {
             maintenance_server,
             curp_server,
         ) = self.init_servers(persistent, key_pair).await?;
+        let (mut reporter, health_server) = tonic_health::server::health_reporter();
+        reporter
+            .set_service_status("", ServingStatus::Serving)
+            .await;
         Ok(Server::builder()
             .add_service(RpcLockServer::new(lock_server))
             .add_service(RpcKvServer::new(kv_server))
@@ -240,6 +250,7 @@ impl XlineServer {
             .add_service(RpcWatchServer::new(watch_server))
             .add_service(RpcMaintenanceServer::new(maintenance_server))
             .add_service(ProtocolServer::new(curp_server))
+            .add_service(health_server)
             .serve_with_incoming_shutdown(TcpListenerStream::new(xline_listener), signal)
             .await?)
     }