1 files changed, 86 insertions, 0 deletions
diff --git a/crates/rebel-runner/src/util/unix.rs b/crates/rebel-runner/src/util/unix.rs
new file mode 100644
index 0000000..a97b1db
--- /dev/null
+++ b/crates/rebel-runner/src/util/unix.rs
@@ -0,0 +1,86 @@
+use std::{fs::File, os::unix::prelude::*, path::Path};
+
+use nix::{
+	fcntl::{self, FcntlArg, FdFlag, Flock, OFlag},
+	sched,
+	unistd::Pid,
+};
+
+use rebel_common::error::*;
+
+use super::fs;
+
+pub fn set_blocking<Fd: AsFd>(fd: &Fd, blocking: bool) -> Result<()> {
+	let raw_fd = fd.as_fd().as_raw_fd();
+
+	let flags =
+		OFlag::from_bits_retain(fcntl::fcntl(raw_fd, FcntlArg::F_GETFL).context("fcntl(F_GETFL)")?);
+
+	let new_flags = if blocking {
+		flags & !OFlag::O_NONBLOCK
+	} else {
+		flags | OFlag::O_NONBLOCK
+	};
+
+	if new_flags != flags {
+		fcntl::fcntl(raw_fd, FcntlArg::F_SETFL(new_flags)).context("fcntl(F_SETFL)")?;
+	}
+
+	Ok(())
+}
+
+pub fn set_cloexec<Fd: AsFd>(fd: &Fd, cloexec: bool) -> Result<()> {
+	let raw_fd = fd.as_fd().as_raw_fd();
+
+	let flags = FdFlag::from_bits_retain(
+		fcntl::fcntl(raw_fd, FcntlArg::F_GETFD).context("fcntl(F_GETFD)")?,
+	);
+
+	let new_flags = if cloexec {
+		flags | FdFlag::FD_CLOEXEC
+	} else {
+		flags & !FdFlag::FD_CLOEXEC
+	};
+
+	if new_flags != flags {
+		fcntl::fcntl(raw_fd, FcntlArg::F_SETFD(new_flags)).context("fcntl(F_SETFD)")?;
+	}
+
+	Ok(())
+}
+
+pub fn nproc() -> Result<usize> {
+	const MAXCPU: usize = sched::CpuSet::count();
+
+	let affinity = sched::sched_getaffinity(Pid::from_raw(0)).context("sched_getaffinity()")?;
+
+	let mut count = 0;
+
+	for cpu in 0..MAXCPU {
+		if affinity.is_set(cpu).unwrap() {
+			count += 1;
+		}
+	}
+
+	Ok(count)
+}
+
+pub fn lock<P: AsRef<Path>>(path: P, exclusive: bool, blocking: bool) -> Result<Flock<File>> {
+	use fcntl::FlockArg::*;
+
+	if let Some(parent) = path.as_ref().parent() {
+		fs::mkdir(parent)?;
+	}
+
+	let arg = match (exclusive, blocking) {
+		(true, true) => LockExclusive,
+		(true, false) => LockExclusiveNonblock,
+		(false, true) => LockShared,
+		(false, false) => LockSharedNonblock,
+	};
+
+	let file = fs::create(path.as_ref())?;
+	fcntl::Flock::lock(file, arg)
+		.map_err(|(_, errno)| errno)
+		.with_context(|| format!("flock failed on {:?}", path.as_ref()))
+}