use std::os::fd::{AsRawFd, FromRawFd, OwnedFd, RawFd};
use std::time::{Duration, Instant};
pub const ENV_READINESS_FD: &str = "OCTL_READINESS_FD";
const MAX_MSG: usize = 4096;
const TAG_READY: u8 = b'R';
const TAG_ERROR: u8 = b'E';
#[derive(Debug, PartialEq, Eq)]
pub enum Readiness {
Ready { pid: u32 },
Died,
Error { code: String, message: String },
Malformed(String),
Timeout,
}
pub fn parse_readiness(bytes: &[u8]) -> Readiness {
if bytes.is_empty() {
return Readiness::Died;
}
let Some(body) = bytes.strip_suffix(b"\n") else {
return Readiness::Malformed(String::from_utf8_lossy(bytes).into_owned());
};
match body.first().copied() {
Some(TAG_READY) => {
let digits = &body[1..];
let pid = if !digits.is_empty() && digits.iter().all(u8::is_ascii_digit) {
std::str::from_utf8(digits)
.ok()
.and_then(|s| s.parse::<u32>().ok())
} else {
None
};
match pid {
Some(pid) if pid != 0 && pid <= libc::pid_t::MAX as u32 => Readiness::Ready { pid },
_ => Readiness::Malformed(String::from_utf8_lossy(bytes).into_owned()),
}
}
Some(TAG_ERROR) => {
let rest = &body[1..];
let (code, message) = match rest.iter().position(|&b| b == b'\t') {
Some(i) => (
String::from_utf8_lossy(&rest[..i]).into_owned(),
String::from_utf8_lossy(&rest[i + 1..]).into_owned(),
),
None => (String::from_utf8_lossy(rest).into_owned(), String::new()),
};
Readiness::Error { code, message }
}
_ => Readiness::Malformed(String::from_utf8_lossy(bytes).into_owned()),
}
}
fn is_writable_pipe(fd: RawFd) -> bool {
let mut st: libc::stat = unsafe { std::mem::zeroed() };
if unsafe { libc::fstat(fd, std::ptr::from_mut(&mut st)) } != 0 {
return false;
}
(st.st_mode & libc::S_IFMT) == libc::S_IFIFO
}
fn set_cloexec(fd: RawFd, on: bool) -> std::io::Result<()> {
let flags = unsafe { libc::fcntl(fd, libc::F_GETFD) };
if flags < 0 {
return Err(std::io::Error::last_os_error());
}
let new = if on {
flags | libc::FD_CLOEXEC
} else {
flags & !libc::FD_CLOEXEC
};
if unsafe { libc::fcntl(fd, libc::F_SETFD, new) } < 0 {
return Err(std::io::Error::last_os_error());
}
Ok(())
}
pub struct ReadinessPipe {
read: OwnedFd,
write: Option<OwnedFd>,
}
impl ReadinessPipe {
pub fn new() -> std::io::Result<Self> {
let mut fds = [0 as RawFd; 2];
if unsafe { libc::pipe(fds.as_mut_ptr()) } != 0 {
return Err(std::io::Error::last_os_error());
}
let read = unsafe { OwnedFd::from_raw_fd(fds[0]) };
let write = unsafe { OwnedFd::from_raw_fd(fds[1]) };
set_cloexec(read.as_raw_fd(), true)?;
set_cloexec(write.as_raw_fd(), true)?;
Ok(Self {
read,
write: Some(write),
})
}
pub fn write_fd(&self) -> RawFd {
self.write
.as_ref()
.expect("write end still held")
.as_raw_fd()
}
pub fn close_write(&mut self) {
self.write = None;
}
pub fn await_ready(self, deadline: Duration) -> Readiness {
debug_assert!(self.write.is_none(), "close_write() before await_ready()");
let fd = self.read.as_raw_fd();
let start = Instant::now();
let mut buf: Vec<u8> = Vec::with_capacity(64);
loop {
let elapsed = start.elapsed();
if elapsed >= deadline {
return Readiness::Timeout;
}
let remaining_ms =
i32::try_from(deadline.saturating_sub(elapsed).as_millis()).unwrap_or(i32::MAX);
let mut pfd = libc::pollfd {
fd,
events: libc::POLLIN,
revents: 0,
};
let rc = unsafe { libc::poll(std::ptr::from_mut(&mut pfd), 1, remaining_ms) };
if rc < 0 {
let err = std::io::Error::last_os_error();
if err.raw_os_error() == Some(libc::EINTR) {
continue;
}
return Readiness::Malformed(format!("readiness poll failed: {err}"));
}
if rc == 0 {
return Readiness::Timeout;
}
let mut chunk = [0u8; 128];
let n =
unsafe { libc::read(fd, chunk.as_mut_ptr().cast::<libc::c_void>(), chunk.len()) };
if n < 0 {
let err = std::io::Error::last_os_error();
if err.raw_os_error() == Some(libc::EINTR) {
continue;
}
return Readiness::Malformed(format!("readiness read failed: {err}"));
}
if n == 0 {
return parse_readiness(&buf);
}
buf.extend_from_slice(&chunk[..n as usize]);
if let Some(pos) = buf.iter().position(|&b| b == b'\n') {
return parse_readiness(&buf[..=pos]);
}
if buf.len() > MAX_MSG {
return Readiness::Malformed(
"readiness frame exceeded size limit before a newline".to_string(),
);
}
}
}
}
pub struct ReadinessReporter {
fd: Option<OwnedFd>,
}
impl ReadinessReporter {
pub fn from_env() -> Self {
let fd = std::env::var(ENV_READINESS_FD)
.ok()
.and_then(|v| v.trim().parse::<RawFd>().ok())
.filter(|&fd| fd > libc::STDERR_FILENO)
.filter(|&fd| is_writable_pipe(fd))
.filter(|&fd| set_cloexec(fd, true).is_ok())
.map(|raw| unsafe { OwnedFd::from_raw_fd(raw) });
Self { fd }
}
pub fn ready(&mut self, pid: u32) {
self.emit(format!("{}{pid}\n", TAG_READY as char));
}
pub fn error(&mut self, code: &str, message: &str) {
let flat = message.replace(['\n', '\t'], " ");
self.emit(format!("{}{code}\t{flat}\n", TAG_ERROR as char));
}
fn emit(&mut self, msg: String) {
if let Some(fd) = self.fd.take() {
write_all(fd.as_raw_fd(), msg.as_bytes());
}
}
}
fn write_all(fd: RawFd, mut buf: &[u8]) {
while !buf.is_empty() {
let n = unsafe { libc::write(fd, buf.as_ptr().cast::<libc::c_void>(), buf.len()) };
if n < 0 {
let err = std::io::Error::last_os_error();
if err.raw_os_error() == Some(libc::EINTR) {
continue;
}
return;
}
if n == 0 {
return;
}
buf = &buf[n as usize..];
}
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_DEADLINE: Duration = Duration::from_secs(5);
#[test]
fn parse_ready_line() {
assert_eq!(
parse_readiness(b"R12345\n"),
Readiness::Ready { pid: 12345 }
);
}
#[test]
fn parse_ready_without_trailing_newline_is_malformed() {
match parse_readiness(b"R777") {
Readiness::Malformed(_) => {}
other => panic!("expected Malformed, got {other:?}"),
}
}
#[test]
fn parse_ready_trailing_garbage_is_malformed() {
match parse_readiness(b"R123xyz\n") {
Readiness::Malformed(_) => {}
other => panic!("expected Malformed, got {other:?}"),
}
}
#[test]
fn parse_ready_overflow_pid_is_malformed() {
match parse_readiness(b"R99999999999\n") {
Readiness::Malformed(_) => {}
other => panic!("expected Malformed, got {other:?}"),
}
}
#[test]
fn parse_empty_is_died() {
assert_eq!(parse_readiness(b""), Readiness::Died);
}
#[test]
fn parse_error_line() {
assert_eq!(
parse_readiness(b"Esupervisor_already_running\tpid 42 is alive\n"),
Readiness::Error {
code: "supervisor_already_running".into(),
message: "pid 42 is alive".into(),
}
);
}
#[test]
fn parse_error_without_message() {
assert_eq!(
parse_readiness(b"Elock_error\n"),
Readiness::Error {
code: "lock_error".into(),
message: String::new(),
}
);
}
#[test]
fn parse_partial_ready_tag_only_is_malformed() {
match parse_readiness(b"R\n") {
Readiness::Malformed(_) => {}
other => panic!("expected Malformed, got {other:?}"),
}
}
#[test]
fn parse_ready_zero_pid_is_malformed() {
match parse_readiness(b"R0\n") {
Readiness::Malformed(_) => {}
other => panic!("expected Malformed, got {other:?}"),
}
}
#[test]
fn parse_garbage_is_malformed() {
match parse_readiness(b"xyzzy\n") {
Readiness::Malformed(_) => {}
other => panic!("expected Malformed, got {other:?}"),
}
}
#[test]
fn pipe_roundtrip_ready() {
let mut pipe = ReadinessPipe::new().unwrap();
let mut reporter = reporter_for(&pipe);
pipe.close_write();
reporter.ready(4242);
assert_eq!(
pipe.await_ready(TEST_DEADLINE),
Readiness::Ready { pid: 4242 }
);
}
#[test]
fn pipe_roundtrip_died_on_drop() {
let mut pipe = ReadinessPipe::new().unwrap();
let reporter = reporter_for(&pipe);
pipe.close_write();
drop(reporter); assert_eq!(pipe.await_ready(TEST_DEADLINE), Readiness::Died);
}
#[test]
fn pipe_roundtrip_error() {
let mut pipe = ReadinessPipe::new().unwrap();
let mut reporter = reporter_for(&pipe);
pipe.close_write();
reporter.error("supervisor_already_running", "pid 9 is alive");
assert_eq!(
pipe.await_ready(TEST_DEADLINE),
Readiness::Error {
code: "supervisor_already_running".into(),
message: "pid 9 is alive".into(),
}
);
}
#[test]
fn reporter_closes_after_ready() {
let mut pipe = ReadinessPipe::new().unwrap();
let mut reporter = reporter_for(&pipe);
pipe.close_write();
reporter.ready(5);
reporter.ready(6);
assert_eq!(pipe.await_ready(TEST_DEADLINE), Readiness::Ready { pid: 5 });
}
#[test]
fn pipe_roundtrip_timeout_on_wedge() {
let mut pipe = ReadinessPipe::new().unwrap();
let _reporter = reporter_for(&pipe);
pipe.close_write();
assert_eq!(
pipe.await_ready(Duration::from_millis(150)),
Readiness::Timeout
);
}
#[test]
#[serial_test::serial(octl_readiness_env)]
fn from_env_absent_is_noop() {
std::env::remove_var(ENV_READINESS_FD);
let mut reporter = ReadinessReporter::from_env();
assert!(reporter.fd.is_none());
reporter.ready(1); reporter.error("x", "y"); }
#[test]
#[serial_test::serial(octl_readiness_env)]
fn from_env_rejects_stdio_fd() {
std::env::set_var(ENV_READINESS_FD, "1");
let reporter = ReadinessReporter::from_env();
std::env::remove_var(ENV_READINESS_FD);
assert!(reporter.fd.is_none(), "must not adopt fd 1 (stdout)");
}
#[test]
#[serial_test::serial(octl_readiness_env)]
fn from_env_rejects_non_pipe_fd() {
let tmp = tempfile::NamedTempFile::new().unwrap();
let raw = tmp.as_file().as_raw_fd();
std::env::set_var(ENV_READINESS_FD, raw.to_string());
let reporter = ReadinessReporter::from_env();
std::env::remove_var(ENV_READINESS_FD);
assert!(reporter.fd.is_none(), "must not adopt a non-pipe fd");
}
fn reporter_for(pipe: &ReadinessPipe) -> ReadinessReporter {
let d = unsafe { libc::dup(pipe.write_fd()) };
assert!(d >= 0, "dup failed");
ReadinessReporter {
fd: Some(unsafe { OwnedFd::from_raw_fd(d) }),
}
}
}