use crate::Result;
use std::{
collections::BTreeMap,
path::PathBuf,
sync::{Arc, Mutex},
};
use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};
pub struct FileRwLocks {
locks: Mutex<BTreeMap<PathBuf, Arc<RwLock<()>>>>,
}
#[derive(Debug, Default)]
#[allow(unused)]
pub struct Flocks {
read_locks: Vec<OwnedRwLockReadGuard<()>>,
write_locks: Vec<OwnedRwLockWriteGuard<()>>,
}
impl FileRwLocks {
pub fn new(files: impl IntoIterator<Item = PathBuf>) -> Self {
let locks = files
.into_iter()
.map(|f| (f, Arc::new(RwLock::new(()))))
.collect();
Self {
locks: Mutex::new(locks),
}
}
pub fn files(&self) -> Vec<PathBuf> {
self.locks.lock().unwrap().keys().cloned().collect()
}
pub fn add_files(&self, files: &[PathBuf]) {
let mut locks = self.locks.lock().unwrap();
for file in files {
locks
.entry(file.clone())
.or_insert_with(|| Arc::new(RwLock::new(())));
}
}
fn try_read_locks(&self, files: &[PathBuf]) -> Result<Flocks> {
let mut locks = self.locks.lock().unwrap();
let mut read_locks = Vec::new();
for file in files {
let lock = self.get_or_create_lock(&mut locks, file);
let lock = lock
.try_read_owned()
.map_err(|_| eyre::eyre!("failed to get read lock {}", file.display()))?;
read_locks.push(lock);
}
Ok(Flocks {
read_locks,
write_locks: vec![],
})
}
pub async fn read_locks(&self, files: &[PathBuf]) -> Flocks {
match self.try_read_locks(files) {
Ok(flocks) => return flocks,
Err(e) => {
debug!("failed to get read locks: {e:?}");
}
}
let mut read_locks = Vec::new();
for file in files {
let lock = self.get_or_create_lock(&mut self.locks.lock().unwrap(), file);
read_locks.push(lock.read_owned().await);
}
Flocks {
read_locks,
write_locks: vec![],
}
}
fn try_write_locks(&self, files: &[PathBuf]) -> Result<Flocks> {
let mut locks = self.locks.lock().unwrap();
let mut write_locks = Vec::new();
for file in files {
let lock = self.get_or_create_lock(&mut locks, file);
let lock = lock
.try_write_owned()
.map_err(|_| eyre::eyre!("failed to get write lock {}", file.display()))?;
write_locks.push(lock);
}
Ok(Flocks {
read_locks: vec![],
write_locks,
})
}
pub async fn write_locks(&self, files: &[PathBuf]) -> Flocks {
match self.try_write_locks(files) {
Ok(flocks) => return flocks,
Err(e) => {
debug!("failed to get write locks: {e:?}");
}
}
let mut write_locks = Vec::new();
for file in files {
let lock = self.get_or_create_lock(&mut self.locks.lock().unwrap(), file);
write_locks.push(lock.write_owned().await);
}
Flocks {
read_locks: vec![],
write_locks,
}
}
fn get_or_create_lock(
&self,
locks: &mut BTreeMap<PathBuf, Arc<RwLock<()>>>,
file: &PathBuf,
) -> Arc<RwLock<()>> {
if let Some(lock) = locks.get(file) {
lock
} else {
locks
.entry(file.clone())
.or_insert_with(|| Arc::new(RwLock::new(())))
}
.clone()
}
}