wixet_bootstrap/
lib.rs

1#![warn(rustdoc::invalid_rust_codeblocks)]
2#![warn(missing_docs)]
3#![warn(rustdoc::missing_crate_level_docs)]
4
5//! This library just starts and configures a logger and provides a friendly shutdown process for tokio apps.
6//! Check [init] to see how to use it
7use std::collections::HashMap;
8
9use signal_hook::consts::signal::*;
10use signal_hook_tokio::{Signals, Handle};
11use futures::stream::StreamExt;
12use log::info;
13use anyhow::Result;
14use tokio::task::JoinHandle;
15
16/// Signal handler. It publish a message into the shutdown channel
17async fn handle_signals(mut signals: Signals, tx: flume::Sender<i32>) {
18    while let Some(signal) = signals.next().await {
19        match signal {
20            SIGHUP => {
21                info!("Reloading");
22            }
23            SIGTERM | SIGINT | SIGQUIT => {
24                info!("Exit signal received, doing friendly shutdown...");
25                tx.send_async(0).await.unwrap();
26                // Shutdown the system;
27            },
28            _ => unreachable!(),
29        }
30    }
31}
32
33/// Configures the logger format. If outfile is none, no log file will be written
34pub fn setup_logger(outfile: Option<&str>, level: Option<log::LevelFilter>, extra_levels:  Option<HashMap<&str, log::LevelFilter>>) -> Result<()>{
35    let mut chain = fern::Dispatch::new()
36    // Perform allocation-free log formatting
37    .format(|out, message, record| {
38        out.finish(format_args!(
39            "{}[{}][{}] {}",
40            chrono::Local::now().format("[%Y-%m-%d][%H:%M:%S]"),
41            record.target(),
42            record.level(),
43            message
44        ))
45    })
46    // Add blanket level filter -
47    .level(level.unwrap_or(log::LevelFilter::Info));
48
49    if let Some(extra_levels) = extra_levels {
50        for (module, level) in extra_levels.into_iter() {
51            chain = chain.level_for(module.to_owned(), level);
52        }
53    }
54    
55    
56    if let Some(outfile) = outfile {
57        chain = chain.chain(fern::log_file(outfile)?);
58    }
59    chain.chain(std::io::stdout())
60    
61    .apply()?;
62    Ok(())
63
64}
65
66/// The shutdown context. It stores information about the signal handler task
67pub struct InitContext {
68    handle: Handle,
69    join: JoinHandle<()>
70}
71
72impl InitContext {
73    /// When the program ends, you should call this to friendly stop the signal handler
74    pub async fn stop(self) -> Result<()> {
75        self.handle.close();
76        self.join.await?;
77        Ok(())
78    }
79}
80
81/// The main function. Call it as soon as possible in your program. It returns the context
82/// and the channel to listen the exit requests
83/// 
84/// ```
85/// use wixet_bootstrap::init;
86/// use log::info;
87/// 
88/// #[tokio::main]
89/// async fn main() {
90///     info!("This log line will be ignored because the logger is not configured yet");
91///     let (closer, exit) = init(Some("output.log")).await?; //If you provide None, it simple will not write a log file (just output)
92///     info!("Hello to my application!")
93/// 
94///     // Do may awesome stuff spawing tokio tasks
95/// 
96///     // I use select here because it is common to listen for multiple signals, but you can just await the `exit` if not
97///     tokio::select!{
98///         _ = exit.recv_async() => {
99///         info!("Shutdown process started");
100///         // Do your friendly stop process here
101///         // This code is run when ctrl+c or any other kill interrupt is received
102///     }
103/// };
104/// 
105/// // A friendly shutdown by deinitializing all "init" stuff.
106/// closer.stop().await?;
107/// info!("Bye");
108/// 
109/// }
110/// ```
111pub async fn init(log_file: Option<&str>, level: Option<log::LevelFilter>, extra_levels:  Option<HashMap<&str, log::LevelFilter>>) -> Result<(InitContext, flume::Receiver<i32>)> {
112    setup_logger(log_file, level, extra_levels)?;
113    // Signals
114    let signals = Signals::new([
115        SIGHUP,
116        SIGTERM,
117        SIGINT,
118        SIGQUIT,
119    ])?;
120    let handle = signals.handle();
121    let (tx, rx) = flume::unbounded();
122    let signals_task = tokio::spawn(handle_signals(signals, tx));
123    Ok((InitContext{handle, join: signals_task}, rx))
124}