From 661c469f386dd500d65e8f756d5ba1e70f6a3234 Mon Sep 17 00:00:00 2001
From: Arkadiy Paronyan <arkady.paronyan@gmail.com>
Date: Fri, 21 Feb 2020 11:47:02 +0100
Subject: [PATCH] Wait for RPC server to cleanup on shutdown (#5004)

---
 substrate/client/service/src/lib.rs | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/substrate/client/service/src/lib.rs b/substrate/client/service/src/lib.rs
index c45d44dfca..8c5d5decca 100644
--- a/substrate/client/service/src/lib.rs
+++ b/substrate/client/service/src/lib.rs
@@ -534,6 +534,30 @@ impl<TBl, TCl, TSc, TNetStatus, TNet, TTxPool, TOc> Drop for
 	}
 }
 
+#[cfg(not(target_os = "unknown"))]
+// Wrapper for HTTP and WS servers that makes sure they are properly shut down.
+mod waiting {
+	pub struct HttpServer(pub Option<sc_rpc_server::HttpServer>);
+	impl Drop for HttpServer {
+		fn drop(&mut self) {
+			if let Some(server) = self.0.take() {
+				server.close_handle().close();
+				server.wait();
+			}
+		}
+	}
+
+	pub struct WsServer(pub Option<sc_rpc_server::WsServer>);
+	impl Drop for WsServer {
+		fn drop(&mut self) {
+			if let Some(server) = self.0.take() {
+				server.close_handle().close();
+				let _ = server.wait();
+			}
+		}
+	}
+}
+
 /// Starts RPC servers that run in their own thread, and returns an opaque object that keeps them alive.
 #[cfg(not(target_os = "unknown"))]
 fn start_rpc_servers<G, E, H: FnMut() -> sc_rpc_server::RpcHandler<sc_rpc::Metadata>>(
@@ -562,7 +586,7 @@ fn start_rpc_servers<G, E, H: FnMut() -> sc_rpc_server::RpcHandler<sc_rpc::Metad
 		maybe_start_server(
 			config.rpc_http,
 			|address| sc_rpc_server::start_http(address, config.rpc_cors.as_ref(), gen_handler()),
-		)?,
+		)?.map(|s| waiting::HttpServer(Some(s))),
 		maybe_start_server(
 			config.rpc_ws,
 			|address| sc_rpc_server::start_ws(
@@ -571,7 +595,7 @@ fn start_rpc_servers<G, E, H: FnMut() -> sc_rpc_server::RpcHandler<sc_rpc::Metad
 				config.rpc_cors.as_ref(),
 				gen_handler(),
 			),
-		)?.map(Mutex::new),
+		)?.map(|s| waiting::WsServer(Some(s))).map(Mutex::new),
 	)))
 }