Graceful shutdown for the task manager (#6654)

* Initial commit

Forked at: 74655edebc
Parent branch: origin/master

* Move task_manager.rs to mod.rs

* Graceful shutdown for the task manager

* Await all background task JoinHandle at the same time

* Add tests

* Make future() wait also for exit signal + fix essential task failed

Probably related to https://github.com/paritytech/cumulus/issues/111

* add comments for non-obvious code

* Use clean_shutdown() in sc-cli

* Adapt code and upgrade tokio in sc-cli

* cleanup spacing in doc

* Add license

* I guess actually running the clean shutdown would be a good idea

* fix tests

* Update client/cli/src/runner.rs

Co-authored-by: Benjamin Kampmann <ben@gnunicorn.org>

* Improve error logging

* disable other tests (can't reproduce on my machine)

* Revert "disable other tests (can't reproduce on my machine)"

This reverts commit c133c590f33c253123ba0555ce719a71ededd60d.

* It is possible that the tasks are ended first

* Revert "It is possible that the tasks are ended first"

This reverts commit 502aba4a49fb5d892e704c412b8a81768a3f2c71.

* Use single threaded scheduler for more predictability

* enable_time

* Revert "enable_time"

This reverts commit 4e152140764a4bddeedff06a4e36ec701909e8c7.

* Revert "Use single threaded scheduler for more predictability"

This reverts commit ee5e13c5f13ff71e012dcda13579dffeb15f8ffc.

* Revert "Revert "It is possible that the tasks are ended first""

This reverts commit 1b91a8ca3eebbdc18be199c8ca188e88669ae649.

* This cannot be verified either with a threaded pool

* Apply suggestions from code review

Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

Co-authored-by: Benjamin Kampmann <ben@parity.io>
Co-authored-by: Benjamin Kampmann <ben@gnunicorn.org>
Co-authored-by: Bastian Köcher <bkchr@users.noreply.github.com>

substrate/client/service/src/task_manager/tests.rs

@@ -0,0 +1,210 @@
+// This file is part of Substrate.
+
+// Copyright (C) 2020 Parity Technologies (UK) Ltd.
+// SPDX-License-Identifier: GPL-3.0-or-later WITH Classpath-exception-2.0
+
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+// GNU General Public License for more details.
+
+// You should have received a copy of the GNU General Public License
+// along with this program. If not, see <https://www.gnu.org/licenses/>.
+
+use crate::config::TaskExecutor;
+use crate::task_manager::TaskManager;
+use futures::future::FutureExt;
+use parking_lot::Mutex;
+use std::any::Any;
+use std::sync::Arc;
+use std::time::Duration;
+
+#[derive(Clone, Debug)]
+struct DropTester(Arc<Mutex<usize>>);
+
+struct DropTesterRef(DropTester);
+
+impl DropTester {
+	fn new() -> DropTester {
+		DropTester(Arc::new(Mutex::new(0)))
+	}
+
+	fn new_ref(&self) -> DropTesterRef {
+		*self.0.lock() += 1;
+		DropTesterRef(self.clone())
+	}
+}
+
+impl PartialEq<usize> for DropTester {
+	fn eq(&self, other: &usize) -> bool {
+		&*self.0.lock() == other
+	}
+}
+
+impl Drop for DropTesterRef {
+	fn drop(&mut self) {
+		*(self.0).0.lock() -= 1;
+	}
+}
+
+#[test]
+fn ensure_drop_tester_working() {
+	let drop_tester = DropTester::new();
+	assert_eq!(drop_tester, 0);
+	let drop_tester_ref_1 = drop_tester.new_ref();
+	assert_eq!(drop_tester, 1);
+	let drop_tester_ref_2 = drop_tester.new_ref();
+	assert_eq!(drop_tester, 2);
+	drop(drop_tester_ref_1);
+	assert_eq!(drop_tester, 1);
+	drop(drop_tester_ref_2);
+	assert_eq!(drop_tester, 0);
+}
+
+async fn run_background_task(_keep_alive: impl Any) {
+	loop {
+		tokio::time::delay_for(Duration::from_secs(1)).await;
+	}
+}
+
+async fn run_background_task_blocking(duration: Duration, _keep_alive: impl Any) {
+	loop {
+		// block for X sec (not interruptible)
+		std::thread::sleep(duration);
+		// await for 1 sec (interruptible)
+		tokio::time::delay_for(Duration::from_secs(1)).await;
+	}
+}
+
+#[test]
+fn ensure_futures_are_awaited_on_shutdown() {
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor: TaskExecutor = (move |future, _| handle.spawn(future).map(|_| ())).into();
+
+	let task_manager = TaskManager::new(task_executor, None).unwrap();
+	let spawn_handle = task_manager.spawn_handle();
+	let drop_tester = DropTester::new();
+	spawn_handle.spawn("task1", run_background_task(drop_tester.new_ref()));
+	spawn_handle.spawn("task2", run_background_task(drop_tester.new_ref()));
+	assert_eq!(drop_tester, 2);
+	// allow the tasks to even start
+	runtime.block_on(async { tokio::time::delay_for(Duration::from_secs(1)).await });
+	assert_eq!(drop_tester, 2);
+	runtime.block_on(task_manager.clean_shutdown());
+	assert_eq!(drop_tester, 0);
+}
+
+#[test]
+fn ensure_keep_alive_during_shutdown() {
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor: TaskExecutor = (move |future, _| handle.spawn(future).map(|_| ())).into();
+
+	let mut task_manager = TaskManager::new(task_executor, None).unwrap();
+	let spawn_handle = task_manager.spawn_handle();
+	let drop_tester = DropTester::new();
+	task_manager.keep_alive(drop_tester.new_ref());
+	spawn_handle.spawn("task1", run_background_task(()));
+	assert_eq!(drop_tester, 1);
+	// allow the tasks to even start
+	runtime.block_on(async { tokio::time::delay_for(Duration::from_secs(1)).await });
+	assert_eq!(drop_tester, 1);
+	runtime.block_on(task_manager.clean_shutdown());
+	assert_eq!(drop_tester, 0);
+}
+
+#[test]
+fn ensure_blocking_futures_are_awaited_on_shutdown() {
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor: TaskExecutor = (move |future, _| handle.spawn(future).map(|_| ())).into();
+
+	let task_manager = TaskManager::new(task_executor, None).unwrap();
+	let spawn_handle = task_manager.spawn_handle();
+	let drop_tester = DropTester::new();
+	spawn_handle.spawn(
+		"task1",
+		run_background_task_blocking(Duration::from_secs(3), drop_tester.new_ref()),
+	);
+	spawn_handle.spawn(
+		"task2",
+		run_background_task_blocking(Duration::from_secs(3), drop_tester.new_ref()),
+	);
+	assert_eq!(drop_tester, 2);
+	// allow the tasks to even start
+	runtime.block_on(async { tokio::time::delay_for(Duration::from_secs(1)).await });
+	assert_eq!(drop_tester, 2);
+	runtime.block_on(task_manager.clean_shutdown());
+	assert_eq!(drop_tester, 0);
+}
+
+#[test]
+fn ensure_no_task_can_be_spawn_after_terminate() {
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor: TaskExecutor = (move |future, _| handle.spawn(future).map(|_| ())).into();
+
+	let mut task_manager = TaskManager::new(task_executor, None).unwrap();
+	let spawn_handle = task_manager.spawn_handle();
+	let drop_tester = DropTester::new();
+	spawn_handle.spawn("task1", run_background_task(drop_tester.new_ref()));
+	spawn_handle.spawn("task2", run_background_task(drop_tester.new_ref()));
+	assert_eq!(drop_tester, 2);
+	// allow the tasks to even start
+	runtime.block_on(async { tokio::time::delay_for(Duration::from_secs(1)).await });
+	assert_eq!(drop_tester, 2);
+	task_manager.terminate();
+	spawn_handle.spawn("task3", run_background_task(drop_tester.new_ref()));
+	runtime.block_on(task_manager.clean_shutdown());
+	assert_eq!(drop_tester, 0);
+}
+
+#[test]
+fn ensure_task_manager_future_ends_when_task_manager_terminated() {
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor: TaskExecutor = (move |future, _| handle.spawn(future).map(|_| ())).into();
+
+	let mut task_manager = TaskManager::new(task_executor, None).unwrap();
+	let spawn_handle = task_manager.spawn_handle();
+	let drop_tester = DropTester::new();
+	spawn_handle.spawn("task1", run_background_task(drop_tester.new_ref()));
+	spawn_handle.spawn("task2", run_background_task(drop_tester.new_ref()));
+	assert_eq!(drop_tester, 2);
+	// allow the tasks to even start
+	runtime.block_on(async { tokio::time::delay_for(Duration::from_secs(1)).await });
+	assert_eq!(drop_tester, 2);
+	task_manager.terminate();
+	runtime.block_on(task_manager.future()).expect("future has ended without error");
+	runtime.block_on(task_manager.clean_shutdown());
+	assert_eq!(drop_tester, 0);
+}
+
+#[test]
+fn ensure_task_manager_future_ends_with_error_when_essential_task_ends() {
+	let mut runtime = tokio::runtime::Runtime::new().unwrap();
+	let handle = runtime.handle().clone();
+	let task_executor: TaskExecutor = (move |future, _| handle.spawn(future).map(|_| ())).into();
+
+	let mut task_manager = TaskManager::new(task_executor, None).unwrap();
+	let spawn_handle = task_manager.spawn_handle();
+	let spawn_essential_handle = task_manager.spawn_essential_handle();
+	let drop_tester = DropTester::new();
+	spawn_handle.spawn("task1", run_background_task(drop_tester.new_ref()));
+	spawn_handle.spawn("task2", run_background_task(drop_tester.new_ref()));
+	assert_eq!(drop_tester, 2);
+	// allow the tasks to even start
+	runtime.block_on(async { tokio::time::delay_for(Duration::from_secs(1)).await });
+	assert_eq!(drop_tester, 2);
+	spawn_essential_handle.spawn("task3", async { panic!("task failed") });
+	runtime.block_on(task_manager.future()).expect_err("future()'s Result must be Err");
+	assert_eq!(drop_tester, 2);
+	runtime.block_on(task_manager.clean_shutdown());
+	assert_eq!(drop_tester, 0);
+}