use std::fs::{File, OpenOptions};
use std::io;
use std::path::{Path, PathBuf};
#[derive(Debug)]
pub struct FileLock {
lock_path: PathBuf,
}
impl FileLock {
pub fn new(path: impl AsRef<Path>) -> io::Result<Self> {
let path = path.as_ref();
let lock_path = path.with_extension(format!(
"{}.lock",
path.extension()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_default()
));
if let Some(parent) = lock_path.parent() {
std::fs::create_dir_all(parent)?;
}
Ok(Self { lock_path })
}
pub fn shared(&self) -> io::Result<LockGuard> {
self.acquire(LockType::Shared)
}
pub fn try_shared(&self) -> io::Result<Option<LockGuard>> {
self.try_acquire(LockType::Shared)
}
pub fn exclusive(&self) -> io::Result<LockGuard> {
self.acquire(LockType::Exclusive)
}
pub fn try_exclusive(&self) -> io::Result<Option<LockGuard>> {
self.try_acquire(LockType::Exclusive)
}
fn acquire(&self, lock_type: LockType) -> io::Result<LockGuard> {
let file = self.open_lock_file()?;
#[cfg(unix)]
{
use nix::fcntl::{Flock, FlockArg};
let arg = match lock_type {
LockType::Shared => FlockArg::LockShared,
LockType::Exclusive => FlockArg::LockExclusive,
};
match Flock::lock(file, arg) {
Ok(flock) => Ok(LockGuard {
_flock: flock,
_lock_type: lock_type,
}),
Err((_, errno)) => Err(io::Error::new(
io::ErrorKind::Other,
format!("flock failed: {}", errno),
)),
}
}
#[cfg(not(unix))]
{
let _ = (file, lock_type);
Err(io::Error::new(
io::ErrorKind::Unsupported,
"File locking not supported on this platform",
))
}
}
fn try_acquire(&self, lock_type: LockType) -> io::Result<Option<LockGuard>> {
let file = self.open_lock_file()?;
#[cfg(unix)]
{
use nix::errno::Errno;
use nix::fcntl::{Flock, FlockArg};
let arg = match lock_type {
LockType::Shared => FlockArg::LockSharedNonblock,
LockType::Exclusive => FlockArg::LockExclusiveNonblock,
};
match Flock::lock(file, arg) {
Ok(flock) => Ok(Some(LockGuard {
_flock: flock,
_lock_type: lock_type,
})),
Err((_, errno)) if errno == Errno::EWOULDBLOCK || errno == Errno::EAGAIN => {
Ok(None)
}
Err((_, errno)) => Err(io::Error::new(
io::ErrorKind::Other,
format!("flock failed: {}", errno),
)),
}
}
#[cfg(not(unix))]
{
let _ = (file, lock_type);
Err(io::Error::new(
io::ErrorKind::Unsupported,
"File locking not supported on this platform",
))
}
}
fn open_lock_file(&self) -> io::Result<File> {
OpenOptions::new()
.read(true)
.write(true)
.create(true)
.truncate(false)
.open(&self.lock_path)
}
pub fn lock_path(&self) -> &Path {
&self.lock_path
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum LockType {
Shared,
Exclusive,
}
#[derive(Debug)]
pub struct LockGuard {
#[cfg(unix)]
_flock: nix::fcntl::Flock<File>,
_lock_type: LockType,
}
pub struct LockedFile {
lock: FileLock,
}
impl LockedFile {
pub fn new(path: impl AsRef<Path>) -> io::Result<Self> {
Ok(Self {
lock: FileLock::new(path)?,
})
}
pub fn read(&self, path: &Path) -> io::Result<String> {
let _guard = self.lock.shared()?;
if path.exists() {
std::fs::read_to_string(path)
} else {
Ok(String::new())
}
}
pub fn write(&self, path: &Path, content: &str) -> io::Result<()> {
let _guard = self.lock.exclusive()?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
std::fs::write(path, content)
}
pub fn with_shared_lock<T, F>(&self, f: F) -> io::Result<T>
where
F: FnOnce() -> io::Result<T>,
{
let _guard = self.lock.shared()?;
f()
}
pub fn with_exclusive_lock<T, F>(&self, f: F) -> io::Result<T>
where
F: FnOnce() -> io::Result<T>,
{
let _guard = self.lock.exclusive()?;
f()
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::{Arc, Barrier};
use std::thread;
use std::time::{Duration, Instant};
use tempfile::TempDir;
#[test]
fn test_lock_file_path() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.jsonl");
let lock = FileLock::new(&file_path).unwrap();
assert_eq!(lock.lock_path(), temp_dir.path().join("test.jsonl.lock"));
}
#[test]
fn test_lock_file_path_no_extension() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("tasks");
let lock = FileLock::new(&file_path).unwrap();
assert_eq!(lock.lock_path(), temp_dir.path().join("tasks..lock"));
}
#[test]
fn test_shared_lock_acquired() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let lock = FileLock::new(&file_path).unwrap();
let guard = lock.shared();
assert!(guard.is_ok());
}
#[test]
fn test_exclusive_lock_acquired() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let lock = FileLock::new(&file_path).unwrap();
let guard = lock.exclusive();
assert!(guard.is_ok());
}
#[test]
fn test_multiple_shared_locks() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let lock1 = FileLock::new(&file_path).unwrap();
let lock2 = FileLock::new(&file_path).unwrap();
let _guard1 = lock1.shared().unwrap();
let guard2 = lock2.try_shared();
assert!(guard2.is_ok());
assert!(guard2.unwrap().is_some());
}
#[test]
fn test_exclusive_blocks_shared() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let lock1 = FileLock::new(&file_path).unwrap();
let lock2 = FileLock::new(&file_path).unwrap();
let _guard1 = lock1.exclusive().unwrap();
let guard2 = lock2.try_shared();
assert!(guard2.is_ok());
assert!(guard2.unwrap().is_none());
}
#[test]
fn test_shared_blocks_exclusive() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let lock1 = FileLock::new(&file_path).unwrap();
let lock2 = FileLock::new(&file_path).unwrap();
let _guard1 = lock1.shared().unwrap();
let guard2 = lock2.try_exclusive();
assert!(guard2.is_ok());
assert!(guard2.unwrap().is_none());
}
#[test]
fn test_lock_released_on_drop() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let lock1 = FileLock::new(&file_path).unwrap();
let lock2 = FileLock::new(&file_path).unwrap();
{
let _guard1 = lock1.exclusive().unwrap();
}
let guard2 = lock2.try_exclusive();
assert!(guard2.is_ok());
assert!(guard2.unwrap().is_some());
}
#[test]
fn test_locked_file_read_write() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("test.txt");
let locked = LockedFile::new(&file_path).unwrap();
locked.write(&file_path, "Hello, World!").unwrap();
let content = locked.read(&file_path).unwrap();
assert_eq!(content, "Hello, World!");
}
#[test]
fn test_locked_file_read_nonexistent() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("nonexistent.txt");
let locked = LockedFile::new(&file_path).unwrap();
let content = locked.read(&file_path).unwrap();
assert!(content.is_empty());
}
#[test]
fn test_concurrent_writes_serialized() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("counter.txt");
let file_path_clone = file_path.clone();
std::fs::write(&file_path, "0").unwrap();
let barrier = Arc::new(Barrier::new(2));
let barrier_clone = barrier.clone();
let handle1 = thread::spawn(move || {
let locked = LockedFile::new(&file_path).unwrap();
barrier.wait();
locked
.with_exclusive_lock(|| {
let content = std::fs::read_to_string(&file_path)?;
let n: i32 = content.trim().parse().unwrap_or(0);
thread::sleep(Duration::from_millis(10));
std::fs::write(&file_path, format!("{}", n + 1))
})
.unwrap();
});
let handle2 = thread::spawn(move || {
let locked = LockedFile::new(&file_path_clone).unwrap();
barrier_clone.wait();
locked
.with_exclusive_lock(|| {
let content = std::fs::read_to_string(&file_path_clone)?;
let n: i32 = content.trim().parse().unwrap_or(0);
thread::sleep(Duration::from_millis(10));
std::fs::write(&file_path_clone, format!("{}", n + 1))
})
.unwrap();
});
handle1.join().unwrap();
handle2.join().unwrap();
let final_content = std::fs::read_to_string(temp_dir.path().join("counter.txt")).unwrap();
assert_eq!(final_content.trim(), "2");
}
#[test]
fn test_blocking_lock_waits() {
let temp_dir = TempDir::new().unwrap();
let file_path = temp_dir.path().join("wait.txt");
let file_path_clone = file_path.clone();
let barrier = Arc::new(Barrier::new(2));
let barrier_clone = barrier.clone();
let handle1 = thread::spawn(move || {
let lock = FileLock::new(&file_path).unwrap();
let _guard = lock.exclusive().unwrap();
barrier.wait();
thread::sleep(Duration::from_millis(50));
});
let start = Instant::now();
let handle2 = thread::spawn(move || {
let lock = FileLock::new(&file_path_clone).unwrap();
barrier_clone.wait();
let _guard = lock.exclusive().unwrap();
});
handle1.join().unwrap();
handle2.join().unwrap();
assert!(start.elapsed() >= Duration::from_millis(40));
}
}