Graceful shutdown for the task manager

client/cli/Cargo.toml

@@ -20,7 +20,7 @@ regex = "1.3.1"
 time = "0.1.42"
 ansi_term = "0.12.1"
 lazy_static = "1.4.0"
-tokio = { version = "0.2.9", features = [ "signal", "rt-core", "rt-threaded" ] }
+tokio = { version = "0.2.21", features = [ "signal", "rt-core", "rt-threaded", "blocking" ] }
 futures = "0.3.4"
 fdlimit = "0.1.4"
 serde_json = "1.0.41"

client/cli/src/runner.rs

@@ -94,8 +94,8 @@ pub fn build_runtime() -> std::result::Result<tokio::runtime::Runtime, std::io::
 }
 
 fn run_until_exit<FUT, ERR>(
-	mut tokio_runtime: tokio::runtime::Runtime, 
-	future: FUT, 
+	mut tokio_runtime: tokio::runtime::Runtime,
+	future: FUT,
 	mut task_manager: TaskManager,
 ) -> Result<()>
 where
@@ -127,15 +127,13 @@ impl<C: SubstrateCli> Runner<C> {
 		let runtime_handle = tokio_runtime.handle().clone();
 
 		let task_executor = move |fut, task_type| {
-			match task_type {
-				TaskType::Async => { runtime_handle.spawn(fut); }
-				TaskType::Blocking => {
-					runtime_handle.spawn(async move {
-						// `spawn_blocking` is looking for the current runtime, and as such has to
-						// be called from within `spawn`.
-						tokio::task::spawn_blocking(move || futures::executor::block_on(fut))
-					});
-				}
+			let map_result = |_| ();
+
+			match task_type{
+				TaskType::Async => runtime_handle.spawn(fut).map(map_result),
+				TaskType::Blocking =>
+					runtime_handle.spawn_blocking(move || futures::executor::block_on(fut))
+						.map(map_result),
 			}
 		};
 
@@ -235,8 +233,7 @@ impl<C: SubstrateCli> Runner<C> {
 		let mut task_manager = initialise(self.config)?;
 		self.tokio_runtime.block_on(main(task_manager.future().fuse()))
 			.map_err(|e| e.to_string())?;
-		task_manager.terminate();
-		drop(self.tokio_runtime);
+		self.tokio_runtime.block_on(task_manager.clean_shutdown());
 		Ok(())
 	}
 

client/service/Cargo.toml

@@ -90,3 +90,5 @@ substrate-test-runtime-client = { version = "2.0.0-rc4", path = "../../test-util
 sp-consensus-babe = { version = "0.8.0-rc4", path = "../../primitives/consensus/babe" }
 grandpa = { version = "0.8.0-rc4", package = "sc-finality-grandpa", path = "../finality-grandpa" }
 grandpa-primitives = { version = "2.0.0-rc4", package = "sp-finality-grandpa", path = "../../primitives/finality-grandpa" }
+tokio = { version = "0.2", default-features = false }
+async-std = { version = "1.6", default-features = false }

client/service/src/config.rs

@@ -263,7 +263,9 @@ impl std::convert::From<PathBuf> for BasePath {
 	}
 }
 
-type TaskExecutorInner = Arc<dyn Fn(Pin<Box<dyn Future<Output = ()> + Send>>, TaskType) + Send + Sync>;
+// NOTE: here for code readability.
+pub(crate) type SomeFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
+pub(crate) type JoinFuture = Pin<Box<dyn Future<Output = ()> + Send>>;
 
 /// Callable object that execute tasks.
 ///
@@ -275,56 +277,47 @@ type TaskExecutorInner = Arc<dyn Fn(Pin<Box<dyn Future<Output = ()> + Send>>, Ta
 ///
 /// ```
 /// # use sc_service::TaskExecutor;
-/// # mod tokio { pub mod runtime {
-/// # #[derive(Clone)]
-/// # pub struct Runtime;
-/// # impl Runtime {
-/// # pub fn new() -> Result<Self, ()> { Ok(Runtime) }
-/// # pub fn handle(&self) -> &Self { &self }
-/// # pub fn spawn(&self, _: std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>>) {}
-/// # }
-/// # } }
+/// use futures::future::FutureExt;
 /// use tokio::runtime::Runtime;
 ///
 /// let runtime = Runtime::new().unwrap();
 /// let handle = runtime.handle().clone();
 /// let task_executor: TaskExecutor = (move |future, _task_type| {
-///		handle.spawn(future);
-///	}).into();
+///     handle.spawn(future).map(|_| ())
+/// }).into();
 /// ```
 ///
 /// ## Using async-std
 ///
 /// ```
 /// # use sc_service::TaskExecutor;
-/// # mod async_std { pub mod task {
-/// # pub fn spawn(_: std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>>) {}
-/// # } }
 /// let task_executor: TaskExecutor = (|future, _task_type| {
-///		async_std::task::spawn(future);
-///	}).into();
+///     // NOTE: async-std's JoinHandle is not a Result so we don't need to map the result
+///     async_std::task::spawn(future)
+/// }).into();
 /// ```
 #[derive(Clone)]
-pub struct TaskExecutor(TaskExecutorInner);
+pub struct TaskExecutor(Arc<dyn Fn(SomeFuture, TaskType) -> JoinFuture + Send + Sync>);
 
 impl std::fmt::Debug for TaskExecutor {
 	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
 		write!(f, "TaskExecutor")
 	}
 }
 
-impl<F> std::convert::From<F> for TaskExecutor
+impl<F, FUT> std::convert::From<F> for TaskExecutor
 where
-	F: Fn(Pin<Box<dyn Future<Output = ()> + Send>>, TaskType) + Send + Sync + 'static,
+	F: Fn(SomeFuture, TaskType) -> FUT + Send + Sync + 'static,
+	FUT: Future<Output = ()> + Send + 'static,
 {
-	fn from(x: F) -> Self {
-		Self(Arc::new(x))
+	fn from(func: F) -> Self {
+		Self(Arc::new(move |fut, tt| Box::pin(func(fut, tt))))
 	}
 }
 
 impl TaskExecutor {
 	/// Spawns a new asynchronous task.
-	pub fn spawn(&self, future: Pin<Box<dyn Future<Output = ()> + Send>>, task_type: TaskType) {
+	pub fn spawn(&self, future: SomeFuture, task_type: TaskType) -> JoinFuture {
 		self.0(future, task_type)
 	}
 }