[Feature] Gracefully reboot

airup-sdk/src/system.rs

@@ -143,7 +143,10 @@ pub trait ConnectionExt {
     ) -> impl Future<Output = anyhow::Result<Result<Vec<LogRecord>, Error>>>;
 
     /// Enters the specific milestone.
-    fn enter_milestone(&mut self, name: &str) -> impl Future<Output = anyhow::Result<Result<(), Error>>>;
+    fn enter_milestone(
+        &mut self,
+        name: &str,
+    ) -> impl Future<Output = anyhow::Result<Result<(), Error>>>;
 }
 impl ConnectionExt for super::Connection {
     async fn sideload_service(

airup/src/reboot.rs

@@ -37,13 +37,13 @@ pub async fn main(cmdline: Cmdline) -> anyhow::Result<()> {
     let mut conn = super::connect().await?;
 
     if cmdline.reboot {
-        conn.reboot().await??;
+        conn.enter_milestone("reboot").await?.ok();
     }
     if cmdline.poweroff {
-        conn.poweroff().await??;
+        conn.enter_milestone("poweroff").await?.ok();
     }
     if cmdline.halt {
-        conn.halt().await??;
+        conn.enter_milestone("halt").await?.ok();
     }
 
     Ok(())

airupd/src/lifetime.rs

@@ -73,12 +73,15 @@ pub enum Event {
 }
 impl Event {
     /// Handles the event.
-    pub fn handle(&self) -> ! {
+    pub async fn handle(&self) -> ! {
         match self {
             Self::Exit(code) => std::process::exit(*code),
-            Self::Poweroff => power_manager().poweroff().unwrap_log("poweroff() failed"),
-            Self::Reboot => power_manager().reboot().unwrap_log("reboot() failed"),
-            Self::Halt => power_manager().halt().unwrap_log("halt() failed"),
+            Self::Poweroff => power_manager()
+                .poweroff()
+                .await
+                .unwrap_log("poweroff() failed"),
+            Self::Reboot => power_manager().reboot().await.unwrap_log("reboot() failed"),
+            Self::Halt => power_manager().halt().await.unwrap_log("halt() failed"),
             Self::ReloadImage => {
                 airupfx::process::reload_image().unwrap_log("reload_image() failed")
             }

airupd/src/main.rs

@@ -55,6 +55,6 @@ async fn main() {
     let mut lifetime = app::airupd().lifetime.subscribe();
     if let Ok(event) = lifetime.recv().await {
         drop(_lock);
-        event.handle();
+        event.handle().await;
     }
 }

airupd/src/milestones/reboot.rs

@@ -3,8 +3,8 @@
 use crate::app::airupd;
 use ahash::AHashSet;
 use airup_sdk::Error;
-use tokio::task::JoinHandle;
 use std::time::Duration;
+use tokio::task::JoinHandle;
 
 pub const PRESETS: &[&str] = &["reboot", "poweroff", "halt", "ctrlaltdel"];
 
@@ -29,7 +29,8 @@ async fn enter_reboot() -> Result<(), Error> {
         .ok();
 
     stop_all_services(Duration::from_millis(5000)).await;
-
+    airupd().lifetime.reboot();
+
     Ok(())
 }
 
@@ -40,6 +41,7 @@ async fn enter_poweroff() -> Result<(), Error> {
         .ok();
 
     stop_all_services(Duration::from_millis(5000)).await;
+    airupd().lifetime.poweroff();
 
     Ok(())
 }
@@ -51,6 +53,7 @@ async fn enter_halt() -> Result<(), Error> {
         .ok();
 
     stop_all_services(Duration::from_millis(5000)).await;
+    airupd().lifetime.halt();
 
     Ok(())
 }
@@ -103,4 +106,4 @@ fn stop_service_task(service: String) -> JoinHandle<()> {
             }
         };
     })
-}
+}

airupfx/Cargo.toml

@@ -14,7 +14,7 @@ publish = false
 ahash = "0.8"
 airup-sdk = { path = "../airup-sdk", features = ["_internal"] }
 anyhow = "1"
-bytes = "1"
+async-trait = "0.1"
 cfg-if = "1"
 libc = "0.2"
 mini-moka = "0.10"

airupfx/src/fs.rs

@@ -18,11 +18,6 @@ pub async fn read_to_string_limited(
     Ok(buffer)
 }
 
-/// Commits filesystem caches to disk.
-pub async fn sync() {
-    crate::sys::fs::sync().await
-}
-
 /// Sets a file with socket permissions.
 pub async fn set_sock_permission<P: AsRef<Path>>(path: P) -> std::io::Result<()> {
     crate::sys::fs::set_sock_permission(path).await

airupfx/src/power.rs

@@ -3,24 +3,25 @@
 use std::convert::Infallible;
 
 /// Interface of power management.
+#[async_trait::async_trait]
 pub trait PowerManager: Send + Sync {
     /// Immediately powers the device off.
     ///
     /// # Errors
     /// An `Err(_)` is returned if the underlying OS function failed.
-    fn poweroff(&self) -> std::io::Result<Infallible>;
+    async fn poweroff(&self) -> std::io::Result<Infallible>;
 
     /// Immediately reboots the device.
     ///
     /// # Errors
     /// An `Err(_)` is returned if the underlying OS function failed.
-    fn reboot(&self) -> std::io::Result<Infallible>;
+    async fn reboot(&self) -> std::io::Result<Infallible>;
 
     /// Immediately halts the device.
     ///
     /// # Errors
     /// An `Err(_)` is returned if the underlying OS function failed.
-    fn halt(&self) -> std::io::Result<Infallible>;
+    async fn halt(&self) -> std::io::Result<Infallible>;
 }
 
 /// A fallback implementation of `AirupFX` power management.
@@ -29,16 +30,17 @@ pub trait PowerManager: Send + Sync {
 /// to standard error stream and parks the thread if we are `pid == 1`. Otherwise, it directly exits with code `0`.
 #[derive(Default)]
 pub struct Fallback;
+#[async_trait::async_trait]
 impl PowerManager for Fallback {
-    fn poweroff(&self) -> std::io::Result<Infallible> {
+    async fn poweroff(&self) -> std::io::Result<Infallible> {
         Self::halt_process();
     }
 
-    fn reboot(&self) -> std::io::Result<Infallible> {
+    async fn reboot(&self) -> std::io::Result<Infallible> {
         Self::halt_process();
     }
 
-    fn halt(&self) -> std::io::Result<Infallible> {
+    async fn halt(&self) -> std::io::Result<Infallible> {
         Self::halt_process();
     }
 }

airupfx/src/sys/freebsd/power.rs

@@ -6,16 +6,26 @@ use std::convert::Infallible;
 
 #[derive(Default)]
 pub struct FreeBsd;
+impl FreeBsd {
+    async fn prepare(&self) {
+        crate::sys::process::kill_all(Duration::from_millis(5000)).await;
+        crate::sys::fs::sync();
+    }
+}
+#[async_trait::async_trait]
 impl PowerManager for FreeBsd {
-    fn poweroff(&self) -> std::io::Result<Infallible> {
+    async fn poweroff(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         freebsd_reboot(RB_POWEROFF)
     }
 
-    fn reboot(&self) -> std::io::Result<Infallible> {
+    async fn reboot(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         freebsd_reboot(RB_AUTOBOOT)
     }
 
-    fn halt(&self) -> std::io::Result<Infallible> {
+    async fn halt(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         freebsd_reboot(RB_HALT)
     }
 }

airupfx/src/sys/linux/power.rs

@@ -5,22 +5,32 @@ use libc::{
     c_void, syscall, SYS_reboot, LINUX_REBOOT_CMD_HALT, LINUX_REBOOT_CMD_POWER_OFF,
     LINUX_REBOOT_CMD_RESTART, LINUX_REBOOT_MAGIC1, LINUX_REBOOT_MAGIC2,
 };
-use std::{convert::Infallible, ptr::NonNull};
+use std::{convert::Infallible, ptr::NonNull, time::Duration};
 
 #[derive(Default)]
 pub struct Linux;
+impl Linux {
+    async fn prepare(&self) {
+        crate::sys::process::kill_all(Duration::from_millis(5000)).await;
+        crate::sys::fs::sync();
+    }
+}
+#[async_trait::async_trait]
 impl PowerManager for Linux {
-    fn poweroff(&self) -> std::io::Result<Infallible> {
+    async fn poweroff(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         linux_reboot(LINUX_REBOOT_CMD_POWER_OFF, None)?;
         unreachable!()
     }
 
-    fn reboot(&self) -> std::io::Result<Infallible> {
+    async fn reboot(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         linux_reboot(LINUX_REBOOT_CMD_RESTART, None)?;
         unreachable!()
     }
 
-    fn halt(&self) -> std::io::Result<Infallible> {
+    async fn halt(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         linux_reboot(LINUX_REBOOT_CMD_HALT, None)?;
         unreachable!()
     }

airupfx/src/sys/macos/power.rs

@@ -9,23 +9,33 @@ extern "C" {
 }
 
 use crate::power::PowerManager;
-use std::convert::Infallible;
+use std::{convert::Infallible, time::Duration};
 
 const RB_AUTOBOOT: libc::c_int = 0;
 const RB_HALT: libc::c_int = 0x08;
 
 #[derive(Default)]
 pub struct MacOS;
+impl MacOS {
+    async fn prepare(&self) {
+        crate::sys::process::kill_all(Duration::from_millis(5000)).await;
+        crate::sys::fs::sync();
+    }
+}
+#[async_trait::async_trait]
 impl PowerManager for MacOS {
-    fn poweroff(&self) -> std::io::Result<Infallible> {
+    async fn poweroff(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         macos_reboot(RB_HALT)
     }
 
-    fn reboot(&self) -> std::io::Result<Infallible> {
+    async fn reboot(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         macos_reboot(RB_AUTOBOOT)
     }
 
-    fn halt(&self) -> std::io::Result<Infallible> {
+    async fn halt(&self) -> std::io::Result<Infallible> {
+        self.prepare().await;
         macos_reboot(RB_HALT)
     }
 }

airupfx/src/sys/unix/fs.rs

@@ -18,10 +18,8 @@ pub async fn set_sock_permission<P: AsRef<Path>>(path: P) -> std::io::Result<()>
 }
 
 /// Commits filesystem caches to disk.
-pub async fn sync() {
-    tokio::task::spawn_blocking(|| unsafe {
+pub fn sync() {
+    unsafe {
         libc::sync();
-    })
-    .await
-    .unwrap();
+    }
 }

airupfx/src/sys/unix/process.rs

@@ -16,6 +16,7 @@ use std::{
     convert::Infallible,
     os::unix::process::CommandExt,
     sync::{Arc, Mutex, OnceLock, RwLock},
+    time::Duration,
 };
 use sysinfo::UserExt;
 use tokio::{signal::unix::SignalKind, sync::mpsc};
@@ -49,7 +50,7 @@ pub fn reload_image() -> std::io::Result<Infallible> {
 ///
 /// # Errors
 /// An `Err(_)` is returned if the underlying OS function failed.
-pub fn kill(pid: Pid, signum: i32) -> std::io::Result<()> {
+fn kill(pid: Pid, signum: i32) -> std::io::Result<()> {
     let result = unsafe { libc::kill(pid as _, signum) };
     match result {
         0 => Ok(()),
@@ -58,6 +59,29 @@ pub fn kill(pid: Pid, signum: i32) -> std::io::Result<()> {
     }
 }
 
+/// Sends a signal to all running processes, then wait for them to be terminated. If the timeout expired, the processes are
+/// force-killed.
+///
+/// On systems without signal support, forces all running processes to be killed.
+pub(crate) async fn kill_all(timeout: Duration) {
+    eprintln!("Sending SIGTERM to all processes");
+    kill(-1, super::signal::SIGTERM).ok();
+    eprintln!("Waiting for all processes to be terminated");
+    let _lock = child_queue().lock.lock().await;
+    tokio::time::timeout(
+        timeout,
+        tokio::task::spawn_blocking(|| {
+            let mut status = 0;
+            while unsafe { libc::wait(&mut status) > 0 } {}
+        }),
+    )
+    .await
+    .ok();
+    drop(_lock);
+    eprintln!("Sending SIGKILL to all processes");
+    kill(-1, super::signal::SIGKILL).ok();
+}
+
 pub trait ExitStatusExt {
     /// Converts from a `status` returned by [`libc::waitpid`] to [`ExitStatus`].
     fn from_unix(status: libc::c_int) -> Self;
@@ -188,6 +212,7 @@ impl Drop for Child {
 #[derive(Debug, Default)]
 struct ChildQueue {
     queue: RwLock<AHashMap<Pid, mpsc::Sender<Wait>>>,
+    lock: tokio::sync::Mutex<()>,
 }
 impl ChildQueue {
     /// Creates a new [`ChildQueue`] instance.
@@ -206,18 +231,13 @@ impl ChildQueue {
     }
 
     /// Starts the child queue task.
-    ///
-    /// # Panics
-    /// This method would panic if it is called more than once.
-    fn start(&'static self) -> anyhow::Result<()> {
-        static CALLED: OnceLock<()> = OnceLock::new();
-        CALLED.set(()).unwrap();
-
+    fn start(&'static self) -> anyhow::Result<tokio::task::JoinHandle<()>> {
         let mut signal = tokio::signal::unix::signal(SignalKind::child())?;
-        tokio::spawn(async move {
+        Ok(tokio::spawn(async move {
             loop {
                 signal.recv().await;
                 loop {
+                    let _lock = self.lock.lock().await;
                     let wait = match wait_nonblocking(-1) {
                         Ok(Some(x)) => x,
                         Ok(None) => break,
@@ -226,15 +246,15 @@ impl ChildQueue {
                             break;
                         }
                     };
+                    drop(_lock);
 
                     if wait.code().is_some() || wait.signal().is_some() {
                         self.send(wait).await;
                         continue;
                     }
                 }
             }
-        });
-        Ok(())
+        }))
     }
 
     /// Creates a new [`mpsc::Receiver`] handle that will receive [`Wait`] sent after this call to `subscribe`.