Initial commit

Author: kdrag0n

LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2023 Orbital Labs, LLC <[email protected]>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,6 @@
+# pidfd-rs
+
+Fast, polling-free approach to kill and wait for all processes to exit in an init implementation.
+
+- Rust + Tokio
+- pidfd + epoll + timer

pidfd.rs

@@ -0,0 +1,59 @@
+/*
+Copyright (c) 2023 Orbital Labs, LLC <[email protected]>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+use std::os::fd::{OwnedFd, FromRawFd, AsRawFd, RawFd};
+
+use nix::{libc::{syscall, SYS_pidfd_open, PIDFD_NONBLOCK, SYS_pidfd_send_signal, siginfo_t}, sys::signal::Signal};
+use tokio::io::unix::{AsyncFd, AsyncFdReadyGuard};
+
+pub struct PidFd(AsyncFd<OwnedFd>);
+
+impl PidFd {
+    pub fn open(pid: i32) -> std::io::Result<Self> {
+        let fd = unsafe { syscall(SYS_pidfd_open, pid, PIDFD_NONBLOCK) };
+        if fd < 0 {
+            return Err(std::io::Error::last_os_error());
+        }
+        let fd = unsafe { OwnedFd::from_raw_fd(fd as _) };
+        let fd = AsyncFd::new(fd)?;
+        Ok(Self(fd))
+    }
+
+    pub fn kill(&self, signal: Signal) -> nix::Result<()> {
+        let res = unsafe { syscall(SYS_pidfd_send_signal, self.as_raw_fd(), signal, std::ptr::null::<*const siginfo_t>(), 0) };
+        if res < 0 {
+            return Err(nix::Error::last());
+        }
+
+        Ok(())
+    }
+
+    pub async fn wait(&self) -> tokio::io::Result<AsyncFdReadyGuard<OwnedFd>> {
+        self.0.readable().await
+    }
+}
+
+impl AsRawFd for PidFd {
+    fn as_raw_fd(&self) -> RawFd {
+        self.0.as_raw_fd()
+    }
+}

shutdown.rs

@@ -0,0 +1,94 @@
+/*
+Copyright (c) 2023 Orbital Labs, LLC <[email protected]>
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+fn kill_one_entry(entry: Result<DirEntry, io::Error>, signal: Signal) -> Result<Option<PidFd>, Box<dyn Error>> {
+    let filename = entry?.file_name();
+    if let Ok(pid) =  filename.to_str().unwrap().parse::<i32>() {
+        // skip pid 1
+        if pid == 1 {
+            return Ok(None);
+        }
+        
+        // skip kthreads (they won't exit)
+        if is_process_kthread(pid)? {
+            return Ok(None);
+        }
+
+        // open a pidfd before killing, then kill via pidfd for safety
+        let pidfd = PidFd::open(pid)?;
+        pidfd.kill(signal)?;
+        Ok(Some(pidfd))
+    } else {
+        Ok(None)
+    }
+}
+
+fn broadcast_signal(signal: Signal) -> nix::Result<Vec<PidFd>> {
+    // freeze to get consistent snapshot and avoid thrashing
+    kill(Pid::from_raw(-1), Signal::SIGSTOP)?;
+
+    // can't use kill(-1) because we need to know which PIDs to wait for exit
+    // otherwise unmount returns EBUSY
+    let mut pids = Vec::new();
+    match fs::read_dir("/proc") {
+        Ok(entries) => {
+            for entry in entries {
+                match kill_one_entry(entry, signal) {
+                    Ok(Some(pid)) => {
+                        pids.push(pid);
+                    },
+                    Err(e) => {
+                        println!(" !!! Failed to read /proc entry: {}", e);
+                    },
+                    _ => {},
+                }
+            }
+        },
+        Err(e) => {
+            println!(" !!! Failed to read /proc: {}", e);
+        },
+    }
+
+    // always make sure to unfreeze
+    kill(Pid::from_raw(-1), Signal::SIGCONT)?;
+    Ok(pids)
+}
+
+async fn wait_for_pidfds_exit(pidfds: Vec<PidFd>, timeout: Duration) -> Result<(), Box<dyn Error>> {
+    let futures = pidfds.into_iter()
+        .map(|pidfd| {
+            async move {
+                let _guard = pidfd.wait().await?;
+                Ok::<(), tokio::io::Error>(())
+            }
+        })
+        .collect::<Vec<_>>();
+
+    let results = tokio::time::timeout(timeout, futures::future::join_all(futures)).await?;
+    for result in results {
+        if let Err(err) = result {
+            return Err(InitError::PollPidFd(err).into());
+        }
+    }
+
+    Ok(())
+}