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); } } }