1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
use std::path::PathBuf;

use anyhow::Result;
use async_trait::async_trait;
use clap::{App, Arg};
use tokio::runtime;

use crate::logging;

const VERSION: &str = env!("CARGO_PKG_VERSION");
const MINIMUM_WORKER_THREAD_COUNT: usize = 6;

#[async_trait]
pub trait Daemon {
    type Config;

    fn load_config(&self, path_maybe: &Option<PathBuf>) -> Result<Self::Config>;

    async fn start(&mut self, cfg: Self::Config) -> Result<()>;
    async fn stop(&mut self) -> Result<()>;
}

fn worker_thread_count() -> usize {
    let core_count = num_cpus::get();
    core_count.max(MINIMUM_WORKER_THREAD_COUNT)
}

fn init_runtime() -> Result<runtime::Runtime> {
    let rt = runtime::Builder::new_multi_thread()
        .enable_io()
        .enable_time()
        .worker_threads(worker_thread_count())
        .build()?;
    Ok(rt)
}

fn main_loop<D: Daemon>(
    cfg_path: Option<PathBuf>,
    log_config: Option<PathBuf>,
    mut daemon: D,
) -> Result<()> {
    logging::init_logger(&log_config)?;
    let cfg = daemon.load_config(&cfg_path)?;

    let rt = init_runtime()?;
    rt.block_on(async move {
        daemon.start(cfg).await?;
        tokio::signal::ctrl_c().await?;
        daemon.stop().await?;

        Ok(())
    })
}

pub struct DaemonProcess {}

impl DaemonProcess {
    pub fn start<D: Daemon>(name: &str, about: &str, daemon: D) {
        let matches = App::new(name)
            .version(VERSION)
            .about(about)
            .arg(
                Arg::new("config")
                    .short('c')
                    .long("cfg")
                    .value_name("FILE")
                    .about("Sets a config file")
                    .takes_value(true),
            )
            .arg(
                Arg::new("log_config")
                    .short('l')
                    .long("log")
                    .about("Sets the log config file")
                    .takes_value(true),
            )
            .get_matches();

        let cfg: Option<PathBuf> = matches.value_of_t("config").ok();

        let log_config: Option<PathBuf> = matches.value_of_t("log_config").ok();

        if let Err(e) = main_loop(cfg, log_config, daemon) {
            log::error!("fatal: {}", e);
        }
    }
}