add graceful shutdown support for lcp service

Signed-off-by: Jun Kimura <jun.kimura@datachain.jp>

Author: bluele

Cargo.lock

@@ -61,7 +61,7 @@ impl ServiceCmd {
                 let srv = AppService::new(opts.get_home(), enclave);
 
                 info!("start service: addr={addr} mrenclave={mrenclave}");
-                run_service(srv, rt, addr)
+                rt.block_on(async { run_service(srv, addr).await })
             }
         }
     }

app/src/commands/service.rs

@@ -8,6 +8,7 @@ tonic = { version = "0.9", default-features = false }
 tonic-reflection = { version = "0.9" }
 tokio = { version = "1.0", features = ["full"] }
 anyhow = { version = "1.0.56" }
+log = { version = "0.4.8" }
 
 lcp-types = { path = "../types" }
 crypto = { path = "../crypto" }

modules/service/Cargo.toml

@@ -4,9 +4,10 @@ use lcp_proto::lcp::service::{
     elc::v1::{msg_server::MsgServer as ELCMsgServer, query_server::QueryServer as ELCQueryServer},
     enclave::v1::query_server::QueryServer as EnclaveQueryServer,
 };
+use log::*;
 use std::{marker::PhantomData, net::SocketAddr, path::PathBuf, sync::Arc};
 use store::transaction::CommitStore;
-use tokio::runtime::Runtime;
+use tokio::signal::unix::{signal, SignalKind};
 use tonic::transport::Server;
 
 pub struct AppService<E, S>
@@ -47,7 +48,7 @@ where
     }
 }
 
-pub fn run_service<E, S>(srv: AppService<E, S>, rt: Arc<Runtime>, addr: SocketAddr) -> Result<()>
+pub async fn run_service<E, S>(srv: AppService<E, S>, addr: SocketAddr) -> Result<()>
 where
     S: CommitStore,
     E: EnclaveProtoAPI<S>,
@@ -59,15 +60,23 @@ where
         .register_encoded_file_descriptor_set(lcp_proto::FILE_DESCRIPTOR_SET)
         .build()
         .expect("failed to create gRPC reflection servicer");
-    rt.block_on(async {
-        Server::builder()
-            .add_service(elc_msg_srv)
-            .add_service(elc_query_srv)
-            .add_service(enclave_srv)
-            .add_service(reflection)
-            .serve(addr)
-            .await
-            .unwrap();
-    });
+
+    let mut sigint = signal(SignalKind::interrupt()).expect("failed to set SIGINT handler");
+    let mut sigterm = signal(SignalKind::terminate()).expect("failed to set SIGTERM handler");
+    let shutdown_signal = async {
+        let signal_type = tokio::select! {
+            _ = sigint.recv() => "SIGINT",
+            _ = sigterm.recv() => "SIGTERM",
+        };
+        info!("shutdown signal ({}) received, stopping server", signal_type);
+    };
+    Server::builder()
+        .add_service(elc_msg_srv)
+        .add_service(elc_query_srv)
+        .add_service(enclave_srv)
+        .add_service(reflection)
+        .serve_with_shutdown(addr, shutdown_signal)
+        .await?;
+    info!("server stopped");
     Ok(())
 }