From 4e50293c076ba06c0d4a4ac4bfa8c8eea4659973 Mon Sep 17 00:00:00 2001
From: Hatter Jiang <jht5945@gmail.com>
Date: Tue, 6 Oct 2020 23:05:04 +0800
Subject: [PATCH] feat: v0.2.0 - add restart command

---
 Cargo.lock  |  3 ++-
 Cargo.toml  |  3 ++-
 README.md   |  4 ++++
 src/main.rs | 58 +++++++++++++++++++++++++++++++++++++++--------------
 4 files changed, 51 insertions(+), 17 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 4bf4fb0..acb4646 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -517,11 +517,12 @@ dependencies = [
 
 [[package]]
 name = "keeprunningd"
-version = "0.1.0"
+version = "0.2.0"
 dependencies = [
  "argparse",
  "chrono",
  "dingtalk",
+ "lazy_static",
  "rust_util",
  "serde",
  "serde_json",
diff --git a/Cargo.toml b/Cargo.toml
index fb2eb2d..479c4b8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "keeprunningd"
-version = "0.1.0"
+version = "0.2.0"
 authors = ["Hatter Jiang <jht5945@gmail.com>"]
 edition = "2018"
 
@@ -14,6 +14,7 @@ argparse = "0.2.2"
 rust_util = "0.6.3"
 dingtalk = "1.3.2"
 chrono = "0.4.11"
+lazy_static = "1.4.0"
 # log = "0.4.8"
 
 [profile.release]
diff --git a/README.md b/README.md
index 3ce16bc..80f9906 100644
--- a/README.md
+++ b/README.md
@@ -1,2 +1,6 @@
 # keeprunningd
 
+* Changelog
+    * v0.2.0 - add restart command
+    * v0.1.0 - first version
+
diff --git a/src/main.rs b/src/main.rs
index 5619387..8a1c5e9 100644
--- a/src/main.rs
+++ b/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);
             }