use anyhow::{Context, Result};
use std::fs::{File, OpenOptions};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum LockType {
Read,
Write,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LockInfo {
pub lock_type: String, pub process_id: u32,
pub acquired_at: u64,
pub expires_at: u64,
pub holder_info: String, }
impl LockInfo {
pub fn new(lock_type: LockType, duration_secs: u64, holder_info: String) -> Self {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
Self {
lock_type: match lock_type {
LockType::Read => "read".to_string(),
LockType::Write => "write".to_string(),
},
process_id: std::process::id(),
acquired_at: now,
expires_at: now + duration_secs,
holder_info,
}
}
pub fn is_expired(&self) -> bool {
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs();
now >= self.expires_at
}
pub fn is_write_lock(&self) -> bool {
self.lock_type == "write"
}
pub fn is_read_lock(&self) -> bool {
self.lock_type == "read"
}
}
pub struct SWMRLockManager {
collection_path: PathBuf,
locks_dir: PathBuf,
write_lock_path: PathBuf,
}
impl SWMRLockManager {
pub fn new(collection_path: &Path) -> Self {
let locks_dir = collection_path.join("locks");
let write_lock_path = locks_dir.join("write.lock");
Self {
collection_path: collection_path.to_path_buf(),
locks_dir,
write_lock_path,
}
}
pub fn init(&self) -> Result<()> {
std::fs::create_dir_all(&self.locks_dir)
.context("Failed to create locks directory")?;
Ok(())
}
pub fn acquire_read_lock(&self, timeout_secs: u64, holder_info: String) -> Result<ReadLock> {
self.init()?;
let start_time = SystemTime::now();
let timeout = Duration::from_secs(timeout_secs);
loop {
if let Ok(write_lock_info) = self.read_write_lock() {
if !write_lock_info.is_expired() {
if start_time.elapsed().unwrap() >= timeout {
anyhow::bail!("Timeout waiting for read lock - write lock held by process {}",
write_lock_info.process_id);
}
std::thread::sleep(Duration::from_millis(100));
continue;
}
}
let lock_info = LockInfo::new(LockType::Read, 3600, holder_info); let read_lock_path = self.locks_dir.join(format!("read_{}.lock", std::process::id()));
self.write_lock_file(&read_lock_path, &lock_info)?;
return Ok(ReadLock {
manager: self,
lock_path: read_lock_path,
lock_info,
});
}
}
pub fn acquire_write_lock(&self, timeout_secs: u64, holder_info: String) -> Result<WriteLock> {
self.init()?;
let start_time = SystemTime::now();
let timeout = Duration::from_secs(timeout_secs);
loop {
if let Ok(existing_write) = self.read_write_lock() {
if !existing_write.is_expired() {
if start_time.elapsed().unwrap() >= timeout {
anyhow::bail!("Timeout waiting for write lock - held by process {}",
existing_write.process_id);
}
std::thread::sleep(Duration::from_millis(100));
continue;
}
}
let active_read_locks = self.get_active_read_locks()?;
if !active_read_locks.is_empty() {
if start_time.elapsed().unwrap() >= timeout {
anyhow::bail!("Timeout waiting for write lock - {} read locks active",
active_read_locks.len());
}
std::thread::sleep(Duration::from_millis(100));
continue;
}
let lock_info = LockInfo::new(LockType::Write, 1800, holder_info);
self.write_lock_file(&self.write_lock_path, &lock_info)?;
return Ok(WriteLock {
manager: self,
lock_info,
});
}
}
pub fn is_write_locked(&self) -> Result<bool> {
match self.read_write_lock() {
Ok(lock_info) => Ok(!lock_info.is_expired()),
Err(_) => Ok(false),
}
}
pub fn active_read_lock_count(&self) -> Result<usize> {
Ok(self.get_active_read_locks()?.len())
}
fn read_write_lock(&self) -> Result<LockInfo> {
let content = std::fs::read_to_string(&self.write_lock_path)
.context("Failed to read write lock file")?;
let lock_info: LockInfo = serde_json::from_str(&content)
.context("Failed to parse write lock info")?;
Ok(lock_info)
}
fn get_active_read_locks(&self) -> Result<Vec<LockInfo>> {
let mut active_locks = Vec::new();
if !self.locks_dir.exists() {
return Ok(active_locks);
}
for entry in std::fs::read_dir(&self.locks_dir)? {
let entry = entry?;
let path = entry.path();
if let Some(filename) = path.file_name() {
if let Some(filename_str) = filename.to_str() {
if filename_str.starts_with("read_") && filename_str.ends_with(".lock") {
if let Ok(content) = std::fs::read_to_string(&path) {
if let Ok(lock_info) = serde_json::from_str::<LockInfo>(&content) {
if !lock_info.is_expired() {
active_locks.push(lock_info);
} else {
let _ = std::fs::remove_file(&path);
}
}
}
}
}
}
}
Ok(active_locks)
}
fn write_lock_file(&self, path: &Path, lock_info: &LockInfo) -> Result<()> {
let json = serde_json::to_string_pretty(lock_info)
.context("Failed to serialize lock info")?;
std::fs::write(path, json)
.context("Failed to write lock file")?;
Ok(())
}
fn release_read_lock(&self, lock_path: &Path) -> Result<()> {
if lock_path.exists() {
std::fs::remove_file(lock_path)
.context("Failed to remove read lock file")?;
}
Ok(())
}
fn release_write_lock(&self) -> Result<()> {
if self.write_lock_path.exists() {
std::fs::remove_file(&self.write_lock_path)
.context("Failed to remove write lock file")?;
}
Ok(())
}
}
pub struct ReadLock<'a> {
manager: &'a SWMRLockManager,
lock_path: PathBuf,
lock_info: LockInfo,
}
impl<'a> ReadLock<'a> {
pub fn lock_info(&self) -> &LockInfo {
&self.lock_info
}
pub fn extend(&mut self, additional_secs: u64) -> Result<()> {
self.lock_info.expires_at += additional_secs;
self.manager.write_lock_file(&self.lock_path, &self.lock_info)?;
Ok(())
}
}
impl<'a> Drop for ReadLock<'a> {
fn drop(&mut self) {
let _ = self.manager.release_read_lock(&self.lock_path);
}
}
pub struct WriteLock<'a> {
manager: &'a SWMRLockManager,
lock_info: LockInfo,
}
impl<'a> WriteLock<'a> {
pub fn lock_info(&self) -> &LockInfo {
&self.lock_info
}
pub fn extend(&mut self, additional_secs: u64) -> Result<()> {
self.lock_info.expires_at += additional_secs;
self.manager.write_lock_file(&self.manager.write_lock_path, &self.lock_info)?;
Ok(())
}
}
impl<'a> Drop for WriteLock<'a> {
fn drop(&mut self) {
let _ = self.manager.release_write_lock();
}
}
pub trait LockAware {
fn with_read_lock<F, R>(&self, timeout_secs: u64, operation: F) -> Result<R>
where
F: FnOnce() -> Result<R>;
fn with_write_lock<F, R>(&self, timeout_secs: u64, operation: F) -> Result<R>
where
F: FnOnce() -> Result<R>;
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_read_lock_acquisition() {
let temp_dir = TempDir::new().unwrap();
let manager = SWMRLockManager::new(temp_dir.path());
let _lock = manager.acquire_read_lock(5, "test".to_string()).unwrap();
assert_eq!(manager.active_read_lock_count().unwrap(), 1);
}
#[test]
fn test_write_lock_exclusivity() {
let temp_dir = TempDir::new().unwrap();
let manager = SWMRLockManager::new(temp_dir.path());
let _write_lock = manager.acquire_write_lock(5, "test".to_string()).unwrap();
let result = manager.acquire_read_lock(1, "test2".to_string());
assert!(result.is_err());
}
}