1use std::{
2 env,
3 io::{PipeWriter, Write},
4 os::fd::{FromRawFd, RawFd},
5};
6
7#[derive(Debug, thiserror::Error)]
8pub enum Error {
9 #[error("Not running under procdog supervisor")]
10 Unsupervised,
11 #[error("procdog file descriptor number invalid")]
12 FdInvalid,
13 #[error("I/O error on procdog file handle ({0})")]
14 IoError(nix::Error),
15}
16
17pub struct ProcdogAgent {
18 writer: Option<PipeWriter>,
19}
20
21pub fn init(allow_unsupervised: bool) -> Result<ProcdogAgent, Error> {
22 if let Ok(s_fd) = env::var("PROCDOG_FD") {
23 let fd = RawFd::from_str_radix(&s_fd, 10).map_err(|_| Error::FdInvalid)?;
24 let writer = unsafe { PipeWriter::from_raw_fd(fd) };
25 nix::fcntl::fcntl(
26 &writer,
27 nix::fcntl::FcntlArg::F_SETFL(nix::fcntl::OFlag::O_NONBLOCK),
28 )
29 .map_err(|errno| Error::IoError(errno))?;
30 Ok(ProcdogAgent {
31 writer: Some(writer),
32 })
33 } else if allow_unsupervised {
34 Ok(ProcdogAgent { writer: None })
35 } else {
36 Err(Error::Unsupervised)
37 }
38}
39
40impl ProcdogAgent {
41 pub fn heartbeat(&mut self) {
42 if let Some(writer) = self.writer.as_mut() {
43 let buf = [b'H'];
44 writer
45 .write(&buf)
46 .expect("Error writing to procdog supervisor");
47 }
48 }
49}