reth_cli_runner/
lib.rs

1//! A tokio based CLI runner.
2
3#![doc(
4    html_logo_url = "https://raw.githubusercontent.com/paradigmxyz/reth/main/assets/reth-docs.png",
5    html_favicon_url = "https://avatars0.githubusercontent.com/u/97369466?s=256",
6    issue_tracker_base_url = "https://github.com/paradigmxyz/reth/issues/"
7)]
8#![cfg_attr(not(test), warn(unused_crate_dependencies))]
9#![cfg_attr(docsrs, feature(doc_cfg))]
10
11//! Entrypoint for running commands.
12
13use reth_tasks::{TaskExecutor, TaskManager};
14use std::{future::Future, pin::pin, sync::mpsc, time::Duration};
15use tracing::{debug, error, trace};
16
17/// Executes CLI commands.
18///
19/// Provides utilities for running a cli command to completion.
20#[derive(Debug)]
21#[non_exhaustive]
22pub struct CliRunner {
23    tokio_runtime: tokio::runtime::Runtime,
24}
25
26impl CliRunner {
27    /// Attempts to create a new [`CliRunner`] using the default tokio
28    /// [`Runtime`](tokio::runtime::Runtime).
29    ///
30    /// The default tokio runtime is multi-threaded, with both I/O and time drivers enabled.
31    pub fn try_default_runtime() -> Result<Self, std::io::Error> {
32        Ok(Self { tokio_runtime: tokio_runtime()? })
33    }
34
35    /// Create a new [`CliRunner`] from a provided tokio [`Runtime`](tokio::runtime::Runtime).
36    pub const fn from_runtime(tokio_runtime: tokio::runtime::Runtime) -> Self {
37        Self { tokio_runtime }
38    }
39
40    /// Executes an async block on the runtime and blocks until completion.
41    pub fn block_on<F, T>(&self, fut: F) -> T
42    where
43        F: Future<Output = T>,
44    {
45        self.tokio_runtime.block_on(fut)
46    }
47
48    /// Executes the given _async_ command on the tokio runtime until the command future resolves or
49    /// until the process receives a `SIGINT` or `SIGTERM` signal.
50    ///
51    /// Tasks spawned by the command via the [`TaskExecutor`] are shut down and an attempt is made
52    /// to drive their shutdown to completion after the command has finished.
53    pub fn run_command_until_exit<F, E>(
54        self,
55        command: impl FnOnce(CliContext) -> F,
56    ) -> Result<(), E>
57    where
58        F: Future<Output = Result<(), E>>,
59        E: Send + Sync + From<std::io::Error> + From<reth_tasks::PanickedTaskError> + 'static,
60    {
61        let AsyncCliRunner { context, mut task_manager, tokio_runtime } =
62            AsyncCliRunner::new(self.tokio_runtime);
63
64        // Executes the command until it finished or ctrl-c was fired
65        let command_res = tokio_runtime.block_on(run_to_completion_or_panic(
66            &mut task_manager,
67            run_until_ctrl_c(command(context)),
68        ));
69
70        if command_res.is_err() {
71            error!(target: "reth::cli", "shutting down due to error");
72        } else {
73            debug!(target: "reth::cli", "shutting down gracefully");
74            // after the command has finished or exit signal was received we shutdown the task
75            // manager which fires the shutdown signal to all tasks spawned via the task
76            // executor and awaiting on tasks spawned with graceful shutdown
77            task_manager.graceful_shutdown_with_timeout(Duration::from_secs(5));
78        }
79
80        // `drop(tokio_runtime)` would block the current thread until its pools
81        // (including blocking pool) are shutdown. Since we want to exit as soon as possible, drop
82        // it on a separate thread and wait for up to 5 seconds for this operation to
83        // complete.
84        let (tx, rx) = mpsc::channel();
85        std::thread::Builder::new()
86            .name("tokio-runtime-shutdown".to_string())
87            .spawn(move || {
88                drop(tokio_runtime);
89                let _ = tx.send(());
90            })
91            .unwrap();
92
93        let _ = rx.recv_timeout(Duration::from_secs(5)).inspect_err(|err| {
94            debug!(target: "reth::cli", %err, "tokio runtime shutdown timed out");
95        });
96
97        command_res
98    }
99
100    /// Executes a regular future until completion or until external signal received.
101    pub fn run_until_ctrl_c<F, E>(self, fut: F) -> Result<(), E>
102    where
103        F: Future<Output = Result<(), E>>,
104        E: Send + Sync + From<std::io::Error> + 'static,
105    {
106        self.tokio_runtime.block_on(run_until_ctrl_c(fut))?;
107        Ok(())
108    }
109
110    /// Executes a regular future as a spawned blocking task until completion or until external
111    /// signal received.
112    ///
113    /// See [`Runtime::spawn_blocking`](tokio::runtime::Runtime::spawn_blocking) .
114    pub fn run_blocking_until_ctrl_c<F, E>(self, fut: F) -> Result<(), E>
115    where
116        F: Future<Output = Result<(), E>> + Send + 'static,
117        E: Send + Sync + From<std::io::Error> + 'static,
118    {
119        let tokio_runtime = self.tokio_runtime;
120        let handle = tokio_runtime.handle().clone();
121        let fut = tokio_runtime.handle().spawn_blocking(move || handle.block_on(fut));
122        tokio_runtime
123            .block_on(run_until_ctrl_c(async move { fut.await.expect("Failed to join task") }))?;
124
125        // drop the tokio runtime on a separate thread because drop blocks until its pools
126        // (including blocking pool) are shutdown. In other words `drop(tokio_runtime)` would block
127        // the current thread but we want to exit right away.
128        std::thread::Builder::new()
129            .name("tokio-runtime-shutdown".to_string())
130            .spawn(move || drop(tokio_runtime))
131            .unwrap();
132
133        Ok(())
134    }
135}
136
137/// [`CliRunner`] configuration when executing commands asynchronously
138struct AsyncCliRunner {
139    context: CliContext,
140    task_manager: TaskManager,
141    tokio_runtime: tokio::runtime::Runtime,
142}
143
144// === impl AsyncCliRunner ===
145
146impl AsyncCliRunner {
147    /// Given a tokio [`Runtime`](tokio::runtime::Runtime), creates additional context required to
148    /// execute commands asynchronously.
149    fn new(tokio_runtime: tokio::runtime::Runtime) -> Self {
150        let task_manager = TaskManager::new(tokio_runtime.handle().clone());
151        let task_executor = task_manager.executor();
152        Self { context: CliContext { task_executor }, task_manager, tokio_runtime }
153    }
154}
155
156/// Additional context provided by the [`CliRunner`] when executing commands
157#[derive(Debug)]
158pub struct CliContext {
159    /// Used to execute/spawn tasks
160    pub task_executor: TaskExecutor,
161}
162
163/// Creates a new default tokio multi-thread [Runtime](tokio::runtime::Runtime) with all features
164/// enabled
165pub fn tokio_runtime() -> Result<tokio::runtime::Runtime, std::io::Error> {
166    tokio::runtime::Builder::new_multi_thread().enable_all().build()
167}
168
169/// Runs the given future to completion or until a critical task panicked.
170///
171/// Returns the error if a task panicked, or the given future returned an error.
172async fn run_to_completion_or_panic<F, E>(tasks: &mut TaskManager, fut: F) -> Result<(), E>
173where
174    F: Future<Output = Result<(), E>>,
175    E: Send + Sync + From<reth_tasks::PanickedTaskError> + 'static,
176{
177    {
178        let fut = pin!(fut);
179        tokio::select! {
180            task_manager_result = tasks => {
181                if let Err(panicked_error) = task_manager_result {
182                    return Err(panicked_error.into());
183                }
184            },
185            res = fut => res?,
186        }
187    }
188    Ok(())
189}
190
191/// Runs the future to completion or until:
192/// - `ctrl-c` is received.
193/// - `SIGTERM` is received (unix only).
194async fn run_until_ctrl_c<F, E>(fut: F) -> Result<(), E>
195where
196    F: Future<Output = Result<(), E>>,
197    E: Send + Sync + 'static + From<std::io::Error>,
198{
199    let ctrl_c = tokio::signal::ctrl_c();
200
201    #[cfg(unix)]
202    {
203        let mut stream = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
204        let sigterm = stream.recv();
205        let sigterm = pin!(sigterm);
206        let ctrl_c = pin!(ctrl_c);
207        let fut = pin!(fut);
208
209        tokio::select! {
210            _ = ctrl_c => {
211                trace!(target: "reth::cli", "Received ctrl-c");
212            },
213            _ = sigterm => {
214                trace!(target: "reth::cli", "Received SIGTERM");
215            },
216            res = fut => res?,
217        }
218    }
219
220    #[cfg(not(unix))]
221    {
222        let ctrl_c = pin!(ctrl_c);
223        let fut = pin!(fut);
224
225        tokio::select! {
226            _ = ctrl_c => {
227                trace!(target: "reth::cli", "Received ctrl-c");
228            },
229            res = fut => res?,
230        }
231    }
232
233    Ok(())
234}
reth_cli_runner/lib.rs

reth_cli_runner/
lib.rs
