From 4a09fed36ca8f2b9f3ccd2e1c518e777a4c1da6a Mon Sep 17 00:00:00 2001
From: Ivan Sorokin <ivan@sorokin.io>
Date: Thu, 25 Mar 2021 14:29:24 -0400
Subject: [PATCH] Allow rest api to shutdown (#3614)

Co-authored-by: Quentin Le Sceller <q.lesceller@gmail.com>
---
 api/src/rest.rs   | 33 +++++++++++++++++++--------------
 api/tests/rest.rs |  3 +++
 2 files changed, 22 insertions(+), 14 deletions(-)

diff --git a/api/src/rest.rs b/api/src/rest.rs
index 7139cd4421..f74e583396 100644
--- a/api/src/rest.rs
+++ b/api/src/rest.rs
@@ -30,6 +30,7 @@ use rustls::internal::pemfile;
 use std::convert::Infallible;
 use std::fmt::{self, Display};
 use std::fs::File;
+use std::mem;
 use std::net::SocketAddr;
 use std::sync::Arc;
 use std::{io, thread};
@@ -199,25 +200,26 @@ impl ApiServer {
 		addr: SocketAddr,
 		router: Router,
 	) -> Result<thread::JoinHandle<()>, Error> {
-		if self.shutdown_sender.is_some() {
+		if self.is_running() {
 			return Err(ErrorKind::Internal(
 				"Can't start HTTP API server, it's running already".to_string(),
 			)
 			.into());
 		}
-		let (tx, _rx) = oneshot::channel::<()>();
+		let (tx, rx) = oneshot::channel::<()>();
 		self.shutdown_sender = Some(tx);
 		thread::Builder::new()
 			.name("apis".to_string())
 			.spawn(move || {
 				let server = async move {
-					let server = Server::bind(&addr).serve(make_service_fn(move |_| {
-						let router = router.clone();
-						async move { Ok::<_, Infallible>(router) }
-					}));
-					// TODO graceful shutdown is unstable, investigate
-					//.with_graceful_shutdown(rx)
-
+					let server = Server::bind(&addr)
+						.serve(make_service_fn(move |_| {
+							let router = router.clone();
+							async move { Ok::<_, Infallible>(router) }
+						}))
+						.with_graceful_shutdown(async {
+							rx.await.ok();
+						});
 					server.await
 				};
 
@@ -239,7 +241,7 @@ impl ApiServer {
 		router: Router,
 		conf: TLSConfig,
 	) -> Result<thread::JoinHandle<()>, Error> {
-		if self.shutdown_sender.is_some() {
+		if self.is_running() {
 			return Err(ErrorKind::Internal(
 				"Can't start HTTPS API server, it's running already".to_string(),
 			)
@@ -280,10 +282,9 @@ impl ApiServer {
 
 	/// Stops the API server, it panics in case of error
 	pub fn stop(&mut self) -> bool {
-		if self.shutdown_sender.is_some() {
-			// TODO re-enable stop after investigation
-			//let tx = mem::replace(&mut self.shutdown_sender, None).unwrap();
-			//tx.send(()).expect("Failed to stop API server");
+		if self.is_running() {
+			let tx = mem::replace(&mut self.shutdown_sender, None).unwrap();
+			tx.send(()).expect("Failed to stop API server");
 			info!("API server has been stopped");
 			true
 		} else {
@@ -291,6 +292,10 @@ impl ApiServer {
 			false
 		}
 	}
+
+	pub fn is_running(&self) -> bool {
+		self.shutdown_sender.is_some()
+	}
 }
 
 pub struct LoggingMiddleware {}
diff --git a/api/tests/rest.rs b/api/tests/rest.rs
index 4f583ab2b8..f94668d6b2 100644
--- a/api/tests/rest.rs
+++ b/api/tests/rest.rs
@@ -78,6 +78,7 @@ fn test_start_api() {
 	assert_eq!(counter.value(), 1);
 	assert!(server.stop());
 	thread::sleep(time::Duration::from_millis(1_000));
+	assert!(!server.is_running());
 }
 
 // To enable this test you need a trusted PKCS12 (p12) certificate bundle
@@ -100,6 +101,8 @@ fn test_start_api_tls() {
 	let index = request_with_retry("https://yourdomain.com:14444/v1/").unwrap();
 	assert_eq!(index.len(), 2);
 	assert!(!server.stop());
+	thread::sleep(time::Duration::from_millis(1_000));
+	assert!(!server.is_running());
 }
 
 fn request_with_retry(url: &str) -> Result<Vec<String>, api::Error> {