runner: make lock() function more generic

1 files changed, 24 insertions, 1 deletions

diff --git a/crates/runner/src/util/unix.rs b/crates/runner/src/util/unix.rs
index 48db764..710138c 100644
--- a/crates/runner/src/util/unix.rs
+++ b/crates/runner/src/util/unix.rs
@@ -1,4 +1,4 @@
-use std::os::unix::prelude::*;
+use std::{fs::File, os::unix::prelude::*, path::Path};
 
 use nix::{
 	fcntl::{self, FcntlArg, FdFlag, OFlag},
@@ -8,6 +8,8 @@ use nix::{
 
 use common::error::*;
 
+use super::fs;
+
 pub fn set_blocking(fd: RawFd, blocking: bool) -> Result<()> {
 	let flags = unsafe {
 		OFlag::from_bits_unchecked(fcntl::fcntl(fd, FcntlArg::F_GETFL).context("fcntl(F_GETFL)")?)
@@ -59,3 +61,24 @@ pub fn nproc() -> Result<usize> {
 
 	Ok(count)
 }
+
+pub fn lock<P: AsRef<Path>>(path: P, exclusive: bool, blocking: bool) -> Result<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(file.as_raw_fd(), arg)
+		.with_context(|| format!("flock failed on {:?}", path.as_ref()))?;
+
+	Ok(file)
+}