feat: v0.2.0 - add restart command

src/main.rs

@@ -1,20 +1,10 @@
-#[macro_use]
-extern crate rust_util;
+#[macro_use] extern crate lazy_static;
+#[macro_use] extern crate rust_util;
 
-use std::{
-    fs,
-    panic,
-    thread,
-    time::Duration,
-    process::Command,
-    sync::Arc,
-};
+use std::{ collections::HashMap, fs, panic, thread, time::Duration, process::Command, sync::{ Arc, Mutex } };
 use chrono::prelude::*;
 use serde::{ Deserialize, Serialize };
-use rust_util::{
-    util_str::read_str_to_lines,
-    util_file::locate_file,
-};
+use rust_util::{ util_str::read_str_to_lines, util_file::locate_file };
 use dingtalk::DingTalk;
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -29,6 +19,11 @@ struct KeepRunningConfig {
 struct KeepRunningConfigItem {
     title: String,
     grep_tokens: Vec<String>,
+    restart_command: Option<Vec<String>>,
+}
+
+lazy_static!{
+    static ref CHECK_FAIL_MAP: Mutex<HashMap<String, u64>> = Mutex::new(HashMap::new());
 }
 
 fn main() {
@@ -94,6 +89,22 @@ fn keep_runningd(keep_running_config: Arc<KeepRunningConfig>) {
             let check_lines = ps_aux_lines.iter().filter(|ln| {
                 keep_running_config_item.grep_tokens.iter().all(|t| ln.contains(t))
             }).collect::<Vec<_>>();
+
+            let mut fail_count = 0;
+            { // limit `CHECK_FAIL_MAP` lock scope
+                let cloned_title = keep_running_config_item.title.clone();
+                let mut check_fail_map = CHECK_FAIL_MAP.lock().unwrap();
+                let check_success = !check_lines.is_empty();
+                if check_success {
+                    check_fail_map.insert(cloned_title, 0);
+                } else {
+                    match check_fail_map.get_mut(&cloned_title) {
+                        None => { check_fail_map.insert(cloned_title, 1); fail_count = 1; },
+                        Some(n) => { *n += 1; fail_count = *n; },
+                    }
+                }
+            }
+
             if check_lines.is_empty() { // if check fail!
                 information!("Send DingTalk notification!");
                 use tokio::runtime;
@@ -101,11 +112,28 @@ fn keep_runningd(keep_running_config: Arc<KeepRunningConfig>) {
                     Err(err) => failure!("Prepare tokio runtime error: {}", err),
                     Ok(mut rt) => {
                         let mut sb = String::with_capacity(1024);
-                        sb.push_str(&format!("Check failed: {}", &keep_running_config_item.title));
+                        sb.push_str(&format!("Check failed: {}, fail count: {}", &keep_running_config_item.title, fail_count));
                         sb.push_str(&format!("\ncheck time: {:?}", Local::now()));
                         rt.block_on(send_notify(&keep_running_config.notify_token, &sb));
                     },
                 }
+                match &keep_running_config_item.restart_command {
+                    Some(restart_command) if fail_count > 1 => {
+                        information!("Fail count is {} > 1, try restart command: {:?}", fail_count, restart_command);
+                        let mut restart_command_iter = restart_command.iter();
+                        if let Some(program) = restart_command_iter.next() {
+                            let mut cmd = Command::new(program);
+                            while let Some(arg) = restart_command_iter.next() {
+                                cmd.arg(arg);
+                            }
+                            match cmd.spawn() {
+                                Ok(child) => success!("Ran process success, process id: {}", child.id()),
+                                Err(e) => failure!("Run restart command failed: {}", e),
+                            }
+                        }
+                    },
+                    _ => ( /* IGNORE */ ),
+                }
             } else if keep_running_config.show_debug_output.unwrap_or(false) {
                 debugging!("Find: {:?}", &check_lines);
             }