refactor `services.rs`

1 files changed, 105 insertions, 116 deletions

diff --git a/init/src/services/units.rs b/init/src/services/units.rs
index 20561e9..ffb7900 100644
--- a/init/src/services/units.rs
+++ b/init/src/services/units.rs
@@ -10,12 +10,11 @@
 // or logs in such init systems are just taking the stdout+stderr
 // of service and showing its output? idk 4 now, look into how its supposed to be
 
-use crate::{log_warning, RUNLEVEL_STATE};
+use crate::{log_success, log_warning, RUNLEVEL_STATE};
 use serde::Deserialize;
 use std::{
     fs::{read_dir, read_to_string},
     process::Child,
-    time::Duration,
 };
 
 /*
@@ -26,8 +25,8 @@ use std::{
 
    [config]
    exec = "path to exec"
-   runlevel =
-   restart=always|never
+   runlevel = enum Runlevel
+   restart=always|never|on-failure
 
 */
 
@@ -37,9 +36,11 @@ enum Restart {
     Always,
     #[serde(alias = "never")]
     Never,
+    #[serde(alias = "on-failure")]
+    OnFailure,
 }
 
-#[derive(Deserialize, PartialEq, Eq)]
+#[derive(Deserialize, PartialEq, Eq, Clone, Copy, Debug)]
 pub enum Runlevel {
     /// The system is shutting down, runlevel int: 0
     #[serde(alias = "shutdown")]
@@ -54,14 +55,17 @@ pub enum Runlevel {
     #[serde(alias = "multiuser-network")]
     MultiNetwork,
     /// Multi-user mode with GUI, runlevel int: 5
-    #[serde(alias = "multiuser-network-gui")] 
+    #[serde(alias = "multiuser-network-gui")]
     MultiGUINetwork,
     /// Runlevel is not critical for running the service, runlevel int: 4
-    #[serde(alias = "undefiled")] 
+    #[serde(alias = "undefined")]
     Undefined,
     /// System going too reboot, runlevel int: 6
     #[serde(alias = "reboot")]
     Reboot,
+    /// System going to halt
+    #[serde(alias = "halt")]
+    Halt,
 }
 
 #[allow(dead_code)]
@@ -94,7 +98,7 @@ pub struct Unit {
     config: ServiceConfig,
 }
 
-/* 
+/*
 impl Restart {
     pub fn as_str(&self) -> &'static str {
         match self {
@@ -103,8 +107,7 @@ impl Restart {
         }
     }
 }
-*/ 
-
+*/
 
 /// Function:
 /// 1. starting services declared in `/etc/vigil/units/`
@@ -115,131 +118,116 @@ impl Restart {
 pub fn services_mainloop() -> Result<(), Box<dyn std::error::Error + Send>> {
     let unit_list = parse_all_units()?;
     let mut pids: Vec<(Child, String, Restart, u8)> = vec![];
-
-    let guard = match RUNLEVEL_STATE.lock() {
-        Ok(g) => g,
-        Err(poisoned) => {
-            log_warning("Mutex was poisoned. Recovering");
-            poisoned.into_inner()
-        }
-    };
+    let mut last_runlevel = Runlevel::Undefined;
 
     loop {
-        /* *RUNLEVEL_STATE.lock().unwrap() */
-        match *guard {
-            Runlevel::Undefined => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::Undefined {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
-                    }
+        let current_runlevel = {
+            let guard = match RUNLEVEL_STATE.lock() {
+                Ok(g) => g,
+                Err(poisoned) => {
+                    log_warning("Mutex was poisoned. Recovering");
+                    poisoned.into_inner()
                 }
-            }
+            };
+            *guard
+        };
 
-            Runlevel::Reboot => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::Reboot {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
-                    }
-                }
-            }
+        if current_runlevel != last_runlevel {
+            log_success(&format!(
+                "Runlevel changed: {:?} -> {:?}",
+                last_runlevel, current_runlevel
+            ));
 
-            Runlevel::Shutdown => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::Shutdown {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
+            // Stops other runlevel services
+            pids.retain(|(child, exec, _, _)| {
+                // TODO: Correct stop with SIGTERM + timeout
+                match child.try_wait() {
+                    Ok(Some(_)) => {
+                        log_success(&format!("Service exited: {}", exec));
+                        false
                     }
+                    _ => true,
                 }
-            }
+            });
 
-            Runlevel::OneUser => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::OneUser {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
+            for unit in &unit_list {
+                if unit.config.runlevel == current_runlevel {
+                    match std::process::Command::new(&unit.config.exec).spawn() {
+                        Ok(child) => {
+                            pids.push((
+                                child,
+                                unit.config.exec.clone(),
+                                unit.config.restart.clone(),
+                                0,
+                            ));
+                            log_success(&format!("Started: {}", unit.config.exec));
+                        }
+                        Err(e) => {
+                            log_warning(&format!("Failed to start {}: {}", unit.config.exec, e));
+                        }
                     }
                 }
             }
 
-            Runlevel::MultiNoNetwork => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::MultiNoNetwork {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
-                    }
-                }
-            }
+            last_runlevel = current_runlevel;
 
-            Runlevel::MultiNetwork => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::MultiNetwork {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
-                    }
-                }
+            if matches!(current_runlevel, Runlevel::Shutdown | Runlevel::Halt) {
+                break Ok(());
             }
+        }
 
-            Runlevel::MultiGUINetwork => {
-                for unit in &unit_list {
-                    if unit.config.runlevel == Runlevel::MultiGUINetwork {
-                        let child = std::process::Command::new(unit.config.exec.clone())
-                            .spawn()
-                            .map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send>)?;
-                        pids.push((
-                            child,
-                            unit.config.exec.clone(),
-                            unit.config.restart.clone(),
-                            0,
-                        ));
+        pids.retain_mut(|entry| {
+            let (child, exec, restart_policy, restart_count) = entry;
+            match child.try_wait() {
+                Ok(Some(status)) => {
+                    log_warning(&format!("Service exited: {} (status: {:?})", exec, status));
+
+                    let should_restart = match *restart_policy {
+                        Restart::Always => true,
+                        Restart::OnFailure => !status.success(),
+                        Restart::Never => false,
+                    };
+
+                    if should_restart {
+                        if *restart_count >= 3 {
+                            log_warning(&format!(
+                                "Maximum restart attempts (3) reached for {}, giving up",
+                                exec
+                            ));
+                            return false;
+                        }
+
+                        match std::process::Command::new(exec.as_str()).spawn() {
+                            Ok(new_child) => {
+                                *child = new_child;
+                                *restart_count += 1;
+                                log_success(&format!(
+                                    "Restarted {} (attempt {})",
+                                    exec, restart_count
+                                ));
+                                true
+                            }
+                            Err(e) => {
+                                log_warning(&format!("Failed to restart {}: {}", exec, e));
+                                false
+                            }
+                        }
+                    } else {
+
+                        false
                     }
                 }
+                Ok(None) => true, 
+                Err(e) => {
+                    log_warning(&format!("Failed to check status of {}: {}", exec, e));
+                    false
+                }
             }
-        }
+        });
+    }
+}
 
+/*
         for i in 0..pids.len() {
             match pids[i].0.try_wait() {
                 Ok(Some(status)) => {
@@ -276,6 +264,7 @@ pub fn services_mainloop() -> Result<(), Box<dyn std::error::Error + Send>> {
         std::thread::sleep(Duration::from_secs(15));
     }
 }
+*/
 
 fn parse_all_units() -> Result<Vec<Unit>, Box<dyn std::error::Error + Send>> {
     let mut units: Vec<Unit> = Vec::new();