newton_cli_runner/
lib.rs

1//! cli crate
2use std::{future::Future, pin::pin};
3use tracing::trace;
4
5/// Newton Runner
6#[derive(Clone, Debug, Default)]
7#[non_exhaustive]
8pub struct NewtonRunner;
9
10impl NewtonRunner {
11    /// Run until ctrl c
12    pub fn run_blocking_until_ctrl_c<F, E>(self, fut: F) -> Result<(), E>
13    where
14        F: Future<Output = Result<(), E>> + Send + 'static,
15        E: Send + Sync + From<std::io::Error> + 'static,
16    {
17        let tokio_runtime = tokio_runtime()?;
18        let handle = tokio_runtime.handle().clone();
19        let fut = tokio_runtime.handle().spawn_blocking(move || handle.block_on(fut));
20        tokio_runtime.block_on(run_until_ctrl_c(async move { fut.await.expect("Failed to join task") }))?;
21
22        // drop the tokio runtime on a separate thread because drop blocks until its pools
23        // (including blocking pool) are shutdown. In other words `drop(tokio_runtime)` would block
24        // the current thread but we want to exit right away.
25        std::thread::Builder::new()
26            .name("tokio-runtime-shutdown".to_string())
27            .spawn(move || drop(tokio_runtime))
28            .unwrap();
29
30        Ok(())
31    }
32}
33
34/// Creates a new default tokio multi-thread [Runtime](tokio::runtime::Runtime) with all features
35/// enabled
36pub fn tokio_runtime() -> Result<tokio::runtime::Runtime, std::io::Error> {
37    tokio::runtime::Builder::new_multi_thread().enable_all().build()
38}
39
40/// Runs the future to completion or until:
41/// - `ctrl-c` is received.
42/// - `SIGTERM` is received (unix only).
43async fn run_until_ctrl_c<F, E>(fut: F) -> Result<(), E>
44where
45    F: Future<Output = Result<(), E>>,
46    E: Send + Sync + 'static + From<std::io::Error>,
47{
48    let ctrl_c = tokio::signal::ctrl_c();
49
50    #[cfg(unix)]
51    {
52        let mut stream = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
53        let sigterm = stream.recv();
54        let sigterm = pin!(sigterm);
55        let ctrl_c = pin!(ctrl_c);
56        let fut = pin!(fut);
57
58        tokio::select! {
59            _ = ctrl_c => {
60                trace!(target: "newton_prover::cli", "Received ctrl-c");
61            },
62            _ = sigterm => {
63                trace!(target: "newton_prover::cli", "Received SIGTERM");
64            },
65            res = fut => res?,
66        }
67    }
68
69    #[cfg(not(unix))]
70    {
71        let ctrl_c = pin!(ctrl_c);
72        let fut = pin!(fut);
73
74        tokio::select! {
75            _ = ctrl_c => {
76                trace!(target: "newton_prover::cli", "Received ctrl-c");
77            },
78            res = fut => res?,
79        }
80    }
81
82    Ok(())
83}
newton_cli_runner/lib.rs

newton_cli_runner/
lib.rs
