chore(sc-cli): improve runner and signals (#13688)

* chore(sc-cli): improve runner and signals

* Update client/cli/src/runner.rs

* fmt

---------

Co-authored-by: Bastian Köcher <git@kchr.de>
This commit is contained in:
yjh
2023-03-27 06:07:17 +08:00
committed by GitHub
parent 0043ff00e2
commit 3a3905ccb1
3 changed files with 106 additions and 89 deletions
+12 -89
View File
@@ -16,80 +16,15 @@
// 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::{error::Error as CliError, Result, SubstrateCli};
use crate::{error::Error as CliError, Result, Signals, SubstrateCli};
use chrono::prelude::*;
use futures::{
future::{self, BoxFuture, FutureExt},
pin_mut, select, Future,
};
use futures::{future::FutureExt, Future};
use log::info;
use sc_service::{Configuration, Error as ServiceError, TaskManager};
use sc_utils::metrics::{TOKIO_THREADS_ALIVE, TOKIO_THREADS_TOTAL};
use std::{marker::PhantomData, time::Duration};
/// Abstraction over OS signals to handle the shutdown of the node smoothly.
///
/// On `unix` this represents `SigInt` and `SigTerm`.
pub struct Signals(BoxFuture<'static, ()>);
impl Signals {
/// Capture the relevant signals to handle shutdown of the node smoothly.
///
/// Needs to be called in a Tokio context to have access to the tokio reactor.
#[cfg(target_family = "unix")]
pub fn capture() -> std::result::Result<Self, ServiceError> {
use tokio::signal::unix::{signal, SignalKind};
let mut stream_int = signal(SignalKind::interrupt()).map_err(ServiceError::Io)?;
let mut stream_term = signal(SignalKind::terminate()).map_err(ServiceError::Io)?;
Ok(Signals(
async move {
future::select(stream_int.recv().boxed(), stream_term.recv().boxed()).await;
}
.boxed(),
))
}
/// Capture the relevant signals to handle shutdown of the node smoothly.
///
/// Needs to be called in a Tokio context to have access to the tokio reactor.
#[cfg(not(unix))]
pub fn capture() -> std::result::Result<Self, ServiceError> {
use tokio::signal::ctrl_c;
Ok(Signals(
async move {
let _ = ctrl_c().await;
}
.boxed(),
))
}
/// A dummy signal that never returns.
pub fn dummy() -> Self {
Self(future::pending().boxed())
}
}
async fn main<F, E>(func: F, signals: impl Future<Output = ()>) -> std::result::Result<(), E>
where
F: Future<Output = std::result::Result<(), E>> + future::FusedFuture,
E: std::error::Error + Send + Sync + 'static,
{
let signals = signals.fuse();
pin_mut!(func, signals);
select! {
_ = signals => {},
res = func => res?,
}
Ok(())
}
/// Build a tokio runtime with all features
/// Build a tokio runtime with all features.
pub fn build_runtime() -> std::result::Result<tokio::runtime::Runtime, std::io::Error> {
tokio::runtime::Builder::new_multi_thread()
.on_thread_start(|| {
@@ -103,25 +38,6 @@ pub fn build_runtime() -> std::result::Result<tokio::runtime::Runtime, std::io::
.build()
}
fn run_until_exit<F, E>(
tokio_runtime: tokio::runtime::Runtime,
future: F,
task_manager: TaskManager,
signals: impl Future<Output = ()>,
) -> std::result::Result<(), E>
where
F: Future<Output = std::result::Result<(), E>> + future::Future,
E: std::error::Error + Send + Sync + 'static + From<ServiceError>,
{
let f = future.fuse();
pin_mut!(f);
tokio_runtime.block_on(main(f, signals))?;
drop(task_manager);
Ok(())
}
/// A Substrate CLI runtime that can be used to run a node or a command
pub struct Runner<C: SubstrateCli> {
config: Configuration,
@@ -171,7 +87,10 @@ impl<C: SubstrateCli> Runner<C> {
self.print_node_infos();
let mut task_manager = self.tokio_runtime.block_on(initialize(self.config))?;
let res = self.tokio_runtime.block_on(main(task_manager.future().fuse(), self.signals.0));
let res = self
.tokio_runtime
.block_on(self.signals.run_until_signal(task_manager.future().fuse()));
// We need to drop the task manager here to inform all tasks that they should shut down.
//
// This is important to be done before we instruct the tokio runtime to shutdown. Otherwise
@@ -234,7 +153,11 @@ impl<C: SubstrateCli> Runner<C> {
E: std::error::Error + Send + Sync + 'static + From<ServiceError> + From<CliError>,
{
let (future, task_manager) = runner(self.config)?;
run_until_exit::<_, E>(self.tokio_runtime, future, task_manager, self.signals.0)
self.tokio_runtime.block_on(self.signals.run_until_signal(future.fuse()))?;
// Drop the task manager before dropping the rest, to ensure that all futures were informed
// about the shut down.
drop(task_manager);
Ok(())
}
/// Get an immutable reference to the node Configuration