Graceful shutdown for the task manager

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.
 ///
@@ -281,15 +283,16 @@ type TaskExecutorInner = Arc<dyn Fn(Pin<Box<dyn Future<Output = ()> + Send>>, Ta
 /// # 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>>) {}
+/// # pub fn spawn(&self, _: std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>>)
+/// # -> std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>> { Box::pin(async {}) }
 /// # }
 /// # } }
 /// 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);
+///		handle.spawn(future)
 ///	}).into();
 /// ```
 ///
@@ -298,14 +301,15 @@ type TaskExecutorInner = Arc<dyn Fn(Pin<Box<dyn Future<Output = ()> + Send>>, Ta
 /// ```
 /// # use sc_service::TaskExecutor;
 /// # mod async_std { pub mod task {
-/// # pub fn spawn(_: std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>>) {}
+/// # pub fn spawn(_: std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>>)
+/// # -> std::pin::Pin<Box<dyn futures::future::Future<Output = ()> + Send>> { Box::pin(async {}) }
 /// # } }
 /// let task_executor: TaskExecutor = (|future, _task_type| {
-///		async_std::task::spawn(future);
+///		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 {
@@ -315,7 +319,7 @@ impl std::fmt::Debug for TaskExecutor {
 
 impl<F> std::convert::From<F> for TaskExecutor
 where
-	F: Fn(Pin<Box<dyn Future<Output = ()> + Send>>, TaskType) + Send + Sync + 'static,
+	F: Fn(SomeFuture, TaskType) -> JoinFuture + Send + Sync + 'static,
 {
 	fn from(x: F) -> Self {
 		Self(Arc::new(x))
@@ -324,7 +328,7 @@ where
 
 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)
 	}
 }

client/service/src/task_manager.rs → client/service/src/task_manager/mod.rs

@@ -15,7 +15,7 @@
 
 use std::{panic, result::Result, pin::Pin};
 use exit_future::Signal;
-use log::debug;
+use log::{debug, error};
 use futures::{
 	Future, FutureExt, StreamExt,
 	future::{select, Either, BoxFuture},
@@ -30,7 +30,7 @@ use prometheus_endpoint::{
 };
 use sc_client_api::CloneableSpawn;
 use sp_utils::mpsc::{TracingUnboundedSender, TracingUnboundedReceiver, tracing_unbounded};
-use crate::{config::{TaskExecutor, TaskType}, Error};
+use crate::{config::{TaskExecutor, TaskType, JoinFuture}, Error};
 
 mod prometheus_future;
 
@@ -40,6 +40,7 @@ pub struct SpawnTaskHandle {
 	on_exit: exit_future::Exit,
 	executor: TaskExecutor,
 	metrics: Option<Metrics>,
+	task_notifier: TracingUnboundedSender<JoinFuture>,
 }
 
 impl SpawnTaskHandle {
@@ -67,6 +68,11 @@ impl SpawnTaskHandle {
 		task: impl Future<Output = ()> + Send + 'static,
 		task_type: TaskType,
 	) {
+		if self.task_notifier.is_closed() {
+			debug!("Attempt to spawn a new task has been prevented: {}", name);
+			return;
+		}
+
 		let on_exit = self.on_exit.clone();
 		let metrics = self.metrics.clone();
 
@@ -111,7 +117,15 @@ impl SpawnTaskHandle {
 			}
 		};
 
-		self.executor.spawn(Box::pin(future), task_type);
+		let join_handle = self.executor.spawn(Box::pin(future), task_type);
+		let mut task_notifier = self.task_notifier.clone();
+		self.executor.spawn(Box::pin(async move {
+				if let Err(err) = task_notifier.send(join_handle).await {
+					error!("Could not send spawned task handle to queue: {}", err);
+				}
+			}),
+			TaskType::Async,
+		);
 	}
 }
 
@@ -223,6 +237,8 @@ pub struct TaskManager {
 	essential_failed_rx: TracingUnboundedReceiver<()>,
 	/// Things to keep alive until the task manager is dropped.
 	keep_alive: Box<dyn std::any::Any + Send + Sync>,
+	task_notifier: TracingUnboundedSender<JoinFuture>,
+	completion_future: JoinFuture,
 }
 
 impl TaskManager {
@@ -233,11 +249,19 @@ impl TaskManager {
 		prometheus_registry: Option<&Registry>
 	) -> Result<Self, PrometheusError> {
 		let (signal, on_exit) = exit_future::signal();
+
 		// A side-channel for essential tasks to communicate shutdown.
 		let (essential_failed_tx, essential_failed_rx) = tracing_unbounded("mpsc_essential_tasks");
 
 		let metrics = prometheus_registry.map(Metrics::register).transpose()?;
 
+		let (task_notifier, mut background_tasks) = tracing_unbounded("mpsc_background_tasks");
+		let completion_future = executor.spawn(Box::pin(async move {
+			while let Some(handle) = background_tasks.next().await {
+				handle.await;
+			}
+		}), TaskType::Async);
+
 		Ok(Self {
 			on_exit,
 			signal: Some(signal),
@@ -246,6 +270,8 @@ impl TaskManager {
 			essential_failed_tx,
 			essential_failed_rx,
 			keep_alive: Box::new(()),
+			task_notifier,
+			completion_future,
 		})
 	}
 
@@ -256,6 +282,7 @@ impl TaskManager {
 			on_exit: self.on_exit.clone(),
 			executor: self.executor.clone(),
 			metrics: self.metrics.clone(),
+			task_notifier: self.task_notifier.clone(),
 		}
 	}
 
@@ -264,6 +291,19 @@ impl TaskManager {
 		SpawnEssentialTaskHandle::new(self.essential_failed_tx.clone(), self.spawn_handle())
 	}
 
+	/// Send the signal for termination, prevent new tasks to be created, await for all the existing
+	/// tasks to finished.
+	pub fn clean_shutdown(mut self) -> Pin<Box<dyn Future<Output = ()> + Send>> {
+		self.terminate();
+		let keep_alive = self.keep_alive;
+		let completion_future = self.completion_future;
+
+		Box::pin(async move {
+			completion_future.await;
+			drop(keep_alive)
+		})
+	}
+
 	/// Return a future that will end if an essential task fails.
 	pub fn future<'a>(&'a mut self) -> Pin<Box<dyn Future<Output = Result<(), Error>> + Send + 'a>> {
 		Box::pin(async move {
@@ -277,6 +317,7 @@ impl TaskManager {
 	pub fn terminate(&mut self) {
 		if let Some(signal) = self.signal.take() {
 			let _ = signal.fire();
+			self.task_notifier.close_channel();
 		}
 	}
 
@@ -286,13 +327,6 @@ impl TaskManager {
 	}
 }
 
-impl Drop for TaskManager {
-	fn drop(&mut self) {
-		debug!(target: "service", "Tasks manager shutdown");
-		self.terminate();
-	}
-}
-
 #[derive(Clone)]
 struct Metrics {
 	// This list is ordered alphabetically