ab_cli_utils/
lib.rs

1//! Utilities used in various CLI applications
2
3use std::panic;
4use std::process::exit;
5use tokio::signal;
6use tracing::level_filters::LevelFilter;
7use tracing_subscriber::layer::SubscriberExt;
8use tracing_subscriber::util::SubscriberInitExt;
9use tracing_subscriber::{EnvFilter, Layer, fmt};
10
11/// Install a panic handler which exits on panics, rather than unwinding. Unwinding can hang the
12/// tokio runtime waiting for stuck tasks or threads.
13pub fn set_exit_on_panic() {
14    let default_panic_hook = panic::take_hook();
15    panic::set_hook(Box::new(move |panic_info| {
16        default_panic_hook(panic_info);
17        exit(1);
18    }));
19}
20
21/// Initialize logger with typical settings
22pub fn init_logger() {
23    tracing_subscriber::registry()
24        .with(
25            fmt::layer().with_filter(
26                EnvFilter::builder()
27                    .with_default_directive(LevelFilter::INFO.into())
28                    .from_env_lossy(),
29            ),
30        )
31        .init();
32}
33
34/// Raise soft file descriptor limit to the hard limit, if possible
35pub fn raise_fd_limit() {
36    match fdlimit::raise_fd_limit() {
37        Ok(fdlimit::Outcome::LimitRaised { from, to }) => {
38            tracing::debug!(
39                "Increased file descriptor limit from previous (most likely soft) limit {} to \
40                new (most likely hard) limit {}",
41                from,
42                to
43            );
44        }
45        Ok(fdlimit::Outcome::Unsupported) => {
46            // Unsupported platform (a platform other than Linux or macOS)
47        }
48        Err(error) => {
49            tracing::warn!(
50                "Failed to increase file descriptor limit for the process due to an error: {}.",
51                error
52            );
53        }
54    }
55}
56
57/// Create a future that waits for `SIGINT` or `SIGTERM` to be sent to the process
58pub async fn shutdown_signal() {
59    #[cfg(unix)]
60    {
61        use futures::FutureExt;
62        use std::pin::pin;
63
64        let mut sigint = signal::unix::signal(signal::unix::SignalKind::interrupt())
65            .expect("Setting signal handlers must never fail");
66        let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())
67            .expect("Setting signal handlers must never fail");
68
69        futures::future::select(
70            pin!(sigint.recv().map(|_| {
71                tracing::info!("Received SIGINT, shutting down farmer...");
72            }),),
73            pin!(sigterm.recv().map(|_| {
74                tracing::info!("Received SIGTERM, shutting down farmer...");
75            }),),
76        )
77        .await;
78    }
79    #[cfg(not(unix))]
80    {
81        signal::ctrl_c()
82            .await
83            .expect("Setting signal handlers must never fail");
84
85        tracing::info!("Received Ctrl+C, shutting down farmer...");
86    }
87}
ab_cli_utils/lib.rs

ab_cli_utils/
lib.rs
