zakura 1.7.0

//! The Tokio runtime for server commands.
//!
//! The tokio runtime is used for:
//! - non-blocking async tasks, via [`Future`]s and
//! - blocking network and file tasks, via [`spawn_blocking`](tokio::task::spawn_blocking).
//!
//! The rayon thread pool is used for:
//! - long-running CPU-bound tasks like cryptography, via [`rayon::spawn_fifo`].

use std::{future::Future, time::Duration};

use eyre::Report;
use tokio::runtime::Runtime;
use tokio_util::sync::CancellationToken;

/// When Zebra is shutting down, wait this long for tokio tasks to finish.
const TOKIO_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(20);

/// Returns a new multi-threaded Tokio runtime with all drivers enabled.
pub(crate) fn new_runtime() -> Runtime {
    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .expect("runtime building should not fail")
}

/// Zakurad's graceful shutdown function, blocks until one of the supported
/// shutdown signals is received.
async fn shutdown() {
    imp::shutdown().await;
}

/// Extension trait for running the server command on the Tokio runtime.
pub(crate) trait RuntimeRun {
    /// Runs the future returned by `run` until it finishes or a shutdown signal arrives,
    /// then stops the runtime, waiting up to [`TOKIO_SHUTDOWN_TIMEOUT`] for its tasks.
    ///
    /// Returns the future's result, or `Ok` if a signal arrived during startup.
    fn run_with_graceful_shutdown<F>(
        self,
        run: impl FnOnce(CancellationToken, CancellationToken) -> F,
    ) -> Result<(), Report>
    where
        F: Future<Output = Result<(), Report>>;
}

impl RuntimeRun for Runtime {
    fn run_with_graceful_shutdown<F>(
        self,
        run: impl FnOnce(CancellationToken, CancellationToken) -> F,
    ) -> Result<(), Report>
    where
        F: Future<Output = Result<(), Report>>,
    {
        let shutdown_token = CancellationToken::new();
        // cancelled means shutdown must await the root future's cleanup.
        let shutdown_cleanup_required = CancellationToken::new();
        let fut = run(shutdown_token.clone(), shutdown_cleanup_required.clone());
        let result = self.block_on(run_until_shutdown(
            fut,
            shutdown(),
            shutdown_token,
            shutdown_cleanup_required,
        ));

        finish_runtime(self, result)
    }
}

pub(crate) async fn run_until_shutdown(
    fut: impl Future<Output = Result<(), Report>>,
    shutdown_signal: impl Future<Output = ()>,
    shutdown_token: CancellationToken,
    shutdown_cleanup_required: CancellationToken,
) -> Result<(), Report> {
    tokio::pin!(fut);
    tokio::pin!(shutdown_signal);

    tokio::select! {
        biased;
        _ = &mut shutdown_signal => {
            shutdown_token.cancel();
            if shutdown_cleanup_required.is_cancelled() {
                fut.await
            } else {
                Ok(())
            }
        }
        result = &mut fut => result,
    }
}

fn finish_runtime(runtime: Runtime, result: Result<(), Report>) -> Result<(), Report> {
    // Don't wait for long blocking tasks before shutting down
    info!(
        ?TOKIO_SHUTDOWN_TIMEOUT,
        "waiting for async tokio tasks to shut down"
    );
    runtime.shutdown_timeout(TOKIO_SHUTDOWN_TIMEOUT);

    match &result {
        Ok(()) => {
            info!("shutting down Zakura");
        }
        Err(error) => {
            warn!(?error, "shutting down Zakura due to an error");
        }
    }

    result
}

#[cfg(test)]
mod tests {
    use std::future;

    use tokio::sync::oneshot;

    use super::*;

    #[tokio::test]
    async fn graceful_shutdown_waits_for_root_future_cleanup() {
        let shutdown = CancellationToken::new();
        let future_shutdown = shutdown.clone();
        let shutdown_cleanup_required = CancellationToken::new();
        shutdown_cleanup_required.cancel();
        let (cleanup_tx, cleanup_rx) = oneshot::channel();

        let fut = async move {
            future_shutdown.cancelled().await;
            cleanup_tx.send(()).expect("cleanup receiver is open");
            Ok(())
        };

        run_until_shutdown(fut, future::ready(()), shutdown, shutdown_cleanup_required)
            .await
            .expect("shutdown should succeed");
        cleanup_rx.await.expect("root future ran cleanup");
    }

    #[tokio::test]
    async fn shutdown_during_startup_does_not_wait_for_root_future() {
        let shutdown = CancellationToken::new();
        let shutdown_cleanup_required = CancellationToken::new();

        run_until_shutdown(
            future::pending(),
            future::ready(()),
            shutdown.clone(),
            shutdown_cleanup_required,
        )
        .await
        .expect("shutdown should succeed");

        assert!(shutdown.is_cancelled());
    }
}

#[cfg(unix)]
mod imp {
    use tokio::signal::unix::{signal, SignalKind};

    pub(super) async fn shutdown() {
        // If both signals are received, select! chooses one of them at random.
        tokio::select! {
            // SIGINT  - Terminal interrupt signal. Typically generated by shells in response to Ctrl-C.
            _ = sig(SignalKind::interrupt(), "SIGINT") => {}
            // SIGTERM - Standard shutdown signal used by process launchers.
            _ = sig(SignalKind::terminate(), "SIGTERM") => {}
        };
    }

    #[instrument]
    async fn sig(kind: SignalKind, name: &'static str) {
        // Create a Future that completes the first
        // time the process receives 'sig'.
        signal(kind)
            .expect("Failed to register signal handler")
            .recv()
            .await;

        zakura_chain::shutdown::set_shutting_down();

        info!(
            // use target to remove 'imp' from output
            target: "zakurad::signal",
            "received {}, starting shutdown",
            name,
        );
    }
}

#[cfg(not(unix))]
mod imp {

    pub(super) async fn shutdown() {
        //  Wait for Ctrl-C in Windows terminals.
        // (Zebra doesn't support NT Service control messages. Use a service wrapper for long-running instances.)
        tokio::signal::ctrl_c()
            .await
            .expect("listening for ctrl-c signal should never fail");

        zakura_chain::shutdown::set_shutting_down();

        info!(
            // use target to remove 'imp' from output
            target: "zakurad::signal",
            "received Ctrl-C, starting shutdown",
        );
    }
}