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::{PanickedTaskError, TaskExecutor};
14use std::{
15    future::Future,
16    pin::{pin, Pin},
17    sync::mpsc,
18    time::Duration,
19};
20use tokio::task::JoinHandle;
21use tracing::{debug, error, info};
22
23/// Executes CLI commands.
24///
25/// Provides utilities for running a cli command to completion.
26#[derive(Debug)]
27pub struct CliRunner {
28    config: CliRunnerConfig,
29    runtime: reth_tasks::Runtime,
30}
31
32impl CliRunner {
33    /// Attempts to create a new [`CliRunner`] using the default
34    /// [`Runtime`](reth_tasks::Runtime).
35    ///
36    /// The default runtime is multi-threaded, with both I/O and time drivers enabled.
37    pub fn try_default_runtime() -> Result<Self, reth_tasks::RuntimeBuildError> {
38        Self::try_with_runtime_config(reth_tasks::RuntimeConfig::default())
39    }
40
41    /// Creates a new [`CliRunner`] with the given [`RuntimeConfig`](reth_tasks::RuntimeConfig).
42    pub fn try_with_runtime_config(
43        config: reth_tasks::RuntimeConfig,
44    ) -> Result<Self, reth_tasks::RuntimeBuildError> {
45        let runtime = reth_tasks::RuntimeBuilder::new(config).build()?;
46        Ok(Self { config: CliRunnerConfig::default(), runtime })
47    }
48
49    /// Sets the [`CliRunnerConfig`] for this runner.
50    pub const fn with_config(mut self, config: CliRunnerConfig) -> Self {
51        self.config = config;
52        self
53    }
54
55    /// Returns a clone of the underlying [`Runtime`](reth_tasks::Runtime).
56    pub fn runtime(&self) -> reth_tasks::Runtime {
57        self.runtime.clone()
58    }
59
60    /// Executes an async block on the runtime and blocks until completion.
61    pub fn block_on<F, T>(&self, fut: F) -> T
62    where
63        F: Future<Output = T>,
64    {
65        self.runtime.handle().block_on(Box::pin(fut) as Pin<Box<dyn Future<Output = T> + '_>>)
66    }
67
68    /// Executes the given _async_ command on the tokio runtime until the command future resolves or
69    /// until the process receives a `SIGINT` or `SIGTERM` signal.
70    ///
71    /// Tasks spawned by the command via the [`TaskExecutor`] are shut down and an attempt is made
72    /// to drive their shutdown to completion after the command has finished.
73    pub fn run_command_until_exit<F, E>(
74        self,
75        command: impl FnOnce(CliContext) -> F,
76    ) -> Result<(), E>
77    where
78        F: Future<Output = Result<(), E>>,
79        E: Send
80            + Sync
81            + std::fmt::Display
82            + From<std::io::Error>
83            + From<reth_tasks::PanickedTaskError>
84            + 'static,
85    {
86        let (context, task_manager_handle) = cli_context(&self.runtime);
87        self.run_until_exit(task_manager_handle, Box::pin(command(context)))
88    }
89
90    /// Runs the type-erased command until it finished or ctrl-c was fired, then shuts down.
91    fn run_until_exit<E>(
92        self,
93        task_manager_handle: JoinHandle<Result<(), PanickedTaskError>>,
94        command: Pin<Box<dyn Future<Output = Result<(), E>> + '_>>,
95    ) -> Result<(), E>
96    where
97        E: Send
98            + Sync
99            + std::fmt::Display
100            + From<std::io::Error>
101            + From<reth_tasks::PanickedTaskError>
102            + 'static,
103    {
104        // Executes the command until it finished or ctrl-c was fired
105        let command_res = self
106            .runtime
107            .handle()
108            .block_on(run_to_completion_or_panic(task_manager_handle, run_until_ctrl_c(command)));
109
110        if let Err(err) = &command_res {
111            error!(target: "reth::cli", %err, "shutting down due to error");
112        } else {
113            debug!(target: "reth::cli", "shutting down gracefully");
114            // after the command has finished or exit signal was received we shutdown the
115            // runtime which fires the shutdown signal to all tasks spawned via the task
116            // executor and awaiting on tasks spawned with graceful shutdown
117            self.runtime.graceful_shutdown_with_timeout(self.config.graceful_shutdown_timeout);
118        }
119
120        runtime_shutdown(self.runtime, true);
121
122        command_res
123    }
124
125    /// Executes a command in a blocking context with access to `CliContext`.
126    ///
127    /// See [`Runtime::spawn_blocking`](tokio::runtime::Runtime::spawn_blocking).
128    pub fn run_blocking_command_until_exit<F, E>(
129        self,
130        command: impl FnOnce(CliContext) -> F + Send + 'static,
131    ) -> Result<(), E>
132    where
133        F: Future<Output = Result<(), E>> + Send + 'static,
134        E: Send
135            + Sync
136            + std::fmt::Display
137            + From<std::io::Error>
138            + From<reth_tasks::PanickedTaskError>
139            + 'static,
140    {
141        self.run_blocking_until_exit(Box::new(move |context| Box::pin(command(context))))
142    }
143
144    /// Runs the type-erased command on the blocking thread pool until it finished or ctrl-c was
145    /// fired, then shuts down.
146    fn run_blocking_until_exit<E>(
147        self,
148        command: Box<dyn FnOnce(CliContext) -> BoxedCommand<E> + Send>,
149    ) -> Result<(), E>
150    where
151        E: Send
152            + Sync
153            + std::fmt::Display
154            + From<std::io::Error>
155            + From<reth_tasks::PanickedTaskError>
156            + 'static,
157    {
158        let (context, task_manager_handle) = cli_context(&self.runtime);
159
160        // Spawn the command on the blocking thread pool
161        let handle = self.runtime.handle().clone();
162        let command_handle = self.runtime.spawn_blocking(move || handle.block_on(command(context)));
163
164        // Wait for the command to complete or ctrl-c
165        let command_res = self.runtime.handle().block_on(run_to_completion_or_panic(
166            task_manager_handle,
167            run_until_ctrl_c(
168                async move { command_handle.await.expect("Failed to join blocking task") },
169            ),
170        ));
171
172        if let Err(err) = &command_res {
173            error!(target: "reth::cli", %err, "shutting down due to error");
174        } else {
175            debug!(target: "reth::cli", "shutting down gracefully");
176            self.runtime.graceful_shutdown_with_timeout(self.config.graceful_shutdown_timeout);
177        }
178
179        runtime_shutdown(self.runtime, true);
180
181        command_res
182    }
183
184    /// Executes a regular future until completion or until external signal received.
185    pub fn run_until_ctrl_c<F, E>(self, fut: F) -> Result<(), E>
186    where
187        F: Future<Output = Result<(), E>>,
188        E: Send + Sync + From<std::io::Error> + 'static,
189    {
190        self.runtime.handle().block_on(run_until_ctrl_c(
191            Box::pin(fut) as Pin<Box<dyn Future<Output = Result<(), E>> + '_>>
192        ))?;
193        Ok(())
194    }
195
196    /// Executes a regular future as a spawned blocking task until completion or until external
197    /// signal received.
198    ///
199    /// See [`Runtime::spawn_blocking`](tokio::runtime::Runtime::spawn_blocking).
200    pub fn run_blocking_until_ctrl_c<F, E>(self, fut: F) -> Result<(), E>
201    where
202        F: Future<Output = Result<(), E>> + Send + 'static,
203        E: Send + Sync + From<std::io::Error> + 'static,
204    {
205        self.run_blocking_until_ctrl_c_boxed(Box::pin(fut))
206    }
207
208    fn run_blocking_until_ctrl_c_boxed<E>(self, fut: BoxedCommand<E>) -> Result<(), E>
209    where
210        E: Send + Sync + From<std::io::Error> + 'static,
211    {
212        let handle = self.runtime.handle().clone();
213        let fut = self.runtime.spawn_blocking(move || handle.block_on(fut));
214        self.runtime
215            .handle()
216            .block_on(run_until_ctrl_c(async move { fut.await.expect("Failed to join task") }))?;
217
218        runtime_shutdown(self.runtime, false);
219
220        Ok(())
221    }
222}
223
224/// A type-erased command future, so that code driving it is not instantiated per command.
225type BoxedCommand<E> = Pin<Box<dyn Future<Output = Result<(), E>> + Send>>;
226
227/// Extracts the task manager handle from the runtime and creates the [`CliContext`].
228fn cli_context(
229    runtime: &reth_tasks::Runtime,
230) -> (CliContext, JoinHandle<Result<(), PanickedTaskError>>) {
231    let handle =
232        runtime.take_task_manager_handle().expect("Runtime must contain a TaskManager handle");
233    let context = CliContext { task_executor: runtime.clone() };
234    (context, handle)
235}
236
237/// Additional context provided by the [`CliRunner`] when executing commands
238#[derive(Debug)]
239pub struct CliContext {
240    /// Used to execute/spawn tasks
241    pub task_executor: TaskExecutor,
242}
243
244/// Default timeout for graceful shutdown of tasks.
245const DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
246
247/// Configuration for [`CliRunner`].
248#[derive(Debug, Clone)]
249pub struct CliRunnerConfig {
250    /// Timeout for graceful shutdown of tasks.
251    ///
252    /// After the command completes, this is the maximum time to wait for spawned tasks
253    /// to finish before forcefully terminating them.
254    pub graceful_shutdown_timeout: Duration,
255}
256
257impl Default for CliRunnerConfig {
258    fn default() -> Self {
259        Self::new()
260    }
261}
262
263impl CliRunnerConfig {
264    /// Creates a new config with default values.
265    pub const fn new() -> Self {
266        Self { graceful_shutdown_timeout: DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT }
267    }
268
269    /// Sets the graceful shutdown timeout.
270    pub const fn with_graceful_shutdown_timeout(mut self, timeout: Duration) -> Self {
271        self.graceful_shutdown_timeout = timeout;
272        self
273    }
274}
275
276/// Runs the given future to completion or until a critical task panicked.
277///
278/// Returns the error if a task panicked, or the given future returned an error.
279async fn run_to_completion_or_panic<F, E>(
280    task_manager_handle: JoinHandle<Result<(), PanickedTaskError>>,
281    fut: F,
282) -> Result<(), E>
283where
284    F: Future<Output = Result<(), E>>,
285    E: Send + Sync + From<reth_tasks::PanickedTaskError> + 'static,
286{
287    let fut = pin!(fut);
288    tokio::select! {
289        task_manager_result = task_manager_handle => {
290            if let Ok(Err(panicked_error)) = task_manager_result {
291                return Err(panicked_error.into());
292            }
293        },
294        res = fut => res?,
295    }
296    Ok(())
297}
298
299/// Runs the future to completion or until:
300/// - `ctrl-c` is received.
301/// - `SIGTERM` is received (unix only).
302async fn run_until_ctrl_c<F, E>(fut: F) -> Result<(), E>
303where
304    F: Future<Output = Result<(), E>>,
305    E: Send + Sync + 'static + From<std::io::Error>,
306{
307    let ctrl_c = tokio::signal::ctrl_c();
308
309    #[cfg(unix)]
310    {
311        let mut stream = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
312        let sigterm = stream.recv();
313        let sigterm = pin!(sigterm);
314        let ctrl_c = pin!(ctrl_c);
315        let fut = pin!(fut);
316
317        tokio::select! {
318            _ = ctrl_c => {
319                info!(target: "reth::cli", "Received ctrl-c");
320            },
321            _ = sigterm => {
322                info!(target: "reth::cli", "Received SIGTERM");
323            },
324            res = fut => res?,
325        }
326    }
327
328    #[cfg(not(unix))]
329    {
330        let ctrl_c = pin!(ctrl_c);
331        let fut = pin!(fut);
332
333        tokio::select! {
334            _ = ctrl_c => {
335                info!(target: "reth::cli", "Received ctrl-c");
336            },
337            res = fut => res?,
338        }
339    }
340
341    Ok(())
342}
343
344/// Default timeout for waiting on the tokio runtime to shut down.
345const DEFAULT_RUNTIME_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(5);
346
347/// Shut down the given [`Runtime`](reth_tasks::Runtime), and wait for it if `wait` is set.
348///
349/// Shuts down the owned tokio runtime even if tasks retain clones of the shared handle.
350/// Pool teardown can block, so it runs on a separate thread.
351fn runtime_shutdown(rt: reth_tasks::Runtime, wait: bool) {
352    let (tx, rx) = mpsc::channel();
353    std::thread::Builder::new()
354        .name("rt-shutdown".to_string())
355        .spawn(move || {
356            rt.shutdown_timeout(DEFAULT_RUNTIME_SHUTDOWN_TIMEOUT);
357            let _ = tx.send(());
358        })
359        .unwrap();
360
361    if wait {
362        let _ = rx.recv_timeout(DEFAULT_RUNTIME_SHUTDOWN_TIMEOUT).inspect_err(|err| {
363            tracing::warn!(target: "reth::cli", %err, "runtime shutdown timed out");
364        });
365    }
366}
367
368#[cfg(test)]
369mod tests {
370    use super::*;
371
372    #[test]
373    fn runtime_shutdown_stops_tasks_with_live_runtime_clones() {
374        let runtime = reth_tasks::Runtime::test();
375        let retained = runtime.clone();
376        let task_runtime = runtime.clone();
377        let (started, ready) = mpsc::channel();
378        let task = runtime.handle().spawn(async move {
379            let _runtime = task_runtime;
380            started.send(()).unwrap();
381            std::future::pending::<()>().await;
382        });
383        ready.recv_timeout(Duration::from_secs(5)).unwrap();
384
385        runtime_shutdown(runtime, true);
386
387        assert!(task.is_finished(), "shutdown returned while a task still owned the runtime");
388        assert!(retained.handle().block_on(task).unwrap_err().is_cancelled());
389    }
390}
reth_cli_runner/lib.rs

reth_cli_runner/
lib.rs
