From 12e7d9ff244e17e4b0d9a9d461646e21bc2a04ea Mon Sep 17 00:00:00 2001
From: Dmitry <b4tm4n@mail.ru>
Date: Fri, 13 Sep 2024 13:11:33 +0300
Subject: [PATCH] + log level arg

---
 src/logger.rs | 39 +++++++++++++++++++++++++++++++++++++++
 src/main.rs   | 32 ++++++++++++++------------------
 2 files changed, 53 insertions(+), 18 deletions(-)
 create mode 100644 src/logger.rs

diff --git a/src/logger.rs b/src/logger.rs
new file mode 100644
index 0000000..add0303
--- /dev/null
+++ b/src/logger.rs
@@ -0,0 +1,39 @@
+use anyhow::Context;
+use clap::ValueEnum;
+use flexi_logger::{AdaptiveFormat, Logger, LoggerHandle};
+
+#[derive(Debug, Clone, ValueEnum)]
+pub enum LogLevel {
+    Error,
+    Warn,
+    Info,
+    Debug,
+    Trace,
+}
+
+impl AsRef<str> for LogLevel {
+    fn as_ref(&self) -> &'static str {
+        match &self {
+            LogLevel::Error => "error",
+            LogLevel::Warn => "warn",
+            LogLevel::Info => "info",
+            LogLevel::Debug => "debug",
+            LogLevel::Trace => "trace",
+        }
+    }
+}
+
+pub fn create_logger<T>(log_level: T) -> anyhow::Result<LoggerHandle>
+where
+    T: AsRef<str>,
+{
+    let logger = Logger::try_with_str(log_level)
+        .context("default logging level invalid")?
+        .format(flexi_logger::detailed_format)
+        .adaptive_format_for_stdout(AdaptiveFormat::Detailed)
+        .log_to_stdout()
+        .start()
+        .context("can't start logger");
+    log_panics::init();
+    logger
+}
diff --git a/src/main.rs b/src/main.rs
index 67920f8..324a340 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -10,8 +10,6 @@ use std::default::Default;
 
 use tokio::time::{timeout, Instant};
 
-use flexi_logger::{AdaptiveFormat, Logger, LoggerHandle};
-
 use std::sync::Arc;
 use std::time::Duration;
 
@@ -19,6 +17,9 @@ use clap::Parser;
 
 use parse_duration::parse as parse_duration;
 
+mod logger;
+use logger::{create_logger, LogLevel};
+
 #[derive(Parser, Debug)]
 #[clap(author, version, about, long_about = None)]
 #[clap(propagate_version = true)]
@@ -34,6 +35,10 @@ struct Cli {
     /// unhealthy status timeout
     #[arg(short, long, default_value = "35s", value_parser = parse_duration)]
     pub unhealthy_timeout: Duration,
+
+    /// log level
+    #[arg(long, value_enum, default_value_t=LogLevel::Info)]
+    pub log_level: LogLevel,
 }
 
 fn get_query_options(args: &Cli) -> ListContainersOptions<String> {
@@ -48,18 +53,6 @@ fn get_query_options(args: &Cli) -> ListContainersOptions<String> {
     }
 }
 
-fn create_logger() -> anyhow::Result<LoggerHandle> {
-    let logger = Logger::try_with_str("info")
-        .context("default logging level invalid")?
-        .format(flexi_logger::detailed_format)
-        .adaptive_format_for_stdout(AdaptiveFormat::Detailed)
-        .log_to_stdout()
-        .start()
-        .context("can't start logger");
-    log_panics::init();
-    logger
-}
-
 async fn query_containers<T>(
     connection: &Docker,
     query_options: ListContainersOptions<T>,
@@ -91,7 +84,9 @@ async fn query_task(
     args: Cli,
 ) {
     let query_options = get_query_options(&args);
+    let max_query_time = interval - Duration::from_millis(1);
     let mut query_time = Duration::new(0, 0);
+
     log::debug!("query_task -> start recv");
     while (timeout(interval - query_time, shutdown_rx.recv()).await).is_err() {
         let start = Instant::now();
@@ -101,7 +96,7 @@ async fn query_task(
             .unwrap_or_default();
 
         let end = Instant::now();
-        query_time = min(end - start, interval - Duration::from_millis(1));
+        query_time = min(end - start, max_query_time);
 
         let res = tx.send(containers).await;
         if res.is_err() {
@@ -119,7 +114,7 @@ async fn filter_task(
     while let Some(containers) = in_rx.recv().await {
         let now = Instant::now();
 
-        log::info!("filter -> found unhealthy: {}", containers.len());
+        log::debug!("filter -> found unhealthy: {}", containers.len());
 
         let prev_times = unhealthy_time.take().unwrap();
         let mut new_times: HashMap<String, Instant> = prev_times
@@ -139,7 +134,7 @@ async fn filter_task(
 
         let _ = unhealthy_time.replace(new_times);
 
-        log::info!("filter -> filtered unhealthy: {}", containers.len());
+        log::debug!("filter -> filtered unhealthy: {}", containers.len());
 
         if containers.is_empty() {
             continue;
@@ -181,9 +176,10 @@ fn shutdown_control(shutdown: Option<mpsc::Sender<()>>) {
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let logger = create_logger()?;
     let cli = Cli::parse();
 
+    let logger = create_logger(&cli.log_level)?;
+
     let connection = Arc::new(Docker::connect_with_defaults()?);
     let _ = connection
         .as_ref()