Skip to main content

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        #[expect(clippy::exit, reason = "Exit on panic is intentional")]
18        exit(1);
19    }));
20}
21
22/// Initialize logger with typical settings
23pub fn init_logger() {
24    tracing_subscriber::registry()
25        .with(
26            fmt::layer().with_filter(
27                EnvFilter::builder()
28                    .with_default_directive(LevelFilter::INFO.into())
29                    .from_env_lossy(),
30            ),
31        )
32        .init();
33}
34
35/// Raise soft file descriptor limit to the hard limit, if possible
36pub fn raise_fd_limit() {
37    match fdlimit::raise_fd_limit() {
38        Ok(fdlimit::Outcome::LimitRaised { from, to }) => {
39            tracing::debug!(
40                "Increased file descriptor limit from previous (most likely soft) limit {} to \
41                new (most likely hard) limit {}",
42                from,
43                to
44            );
45        }
46        Ok(fdlimit::Outcome::Unsupported) => {
47            // Unsupported platform (a platform other than Linux or macOS)
48        }
49        Err(error) => {
50            tracing::warn!(
51                "Failed to increase file descriptor limit for the process due to an error: {}.",
52                error
53            );
54        }
55    }
56}
57
58/// Create a future that waits for `SIGINT` or `SIGTERM` to be sent to the process
59pub async fn shutdown_signal() {
60    #[cfg(unix)]
61    {
62        use futures::FutureExt;
63        use std::pin::pin;
64
65        let mut sigint = signal::unix::signal(signal::unix::SignalKind::interrupt())
66            .expect("Setting signal handlers must never fail");
67        let mut sigterm = signal::unix::signal(signal::unix::SignalKind::terminate())
68            .expect("Setting signal handlers must never fail");
69
70        futures::future::select(
71            pin!(sigint.recv().map(|_| {
72                tracing::info!("Received SIGINT, shutting down farmer...");
73            }),),
74            pin!(sigterm.recv().map(|_| {
75                tracing::info!("Received SIGTERM, shutting down farmer...");
76            }),),
77        )
78        .await;
79    }
80    #[cfg(not(unix))]
81    {
82        signal::ctrl_c()
83            .await
84            .expect("Setting signal handlers must never fail");
85
86        tracing::info!("Received Ctrl+C, shutting down farmer...");
87    }
88}