use std::{collections::HashMap, path::Path, sync::Arc, time::Duration}; use async_trait::async_trait; use tokio::sync::RwLock; use uuid::Uuid; use crate::{ drwmutex::{DRWMutex, Options}, lrwmutex::LRWMutex, LockApi, }; use common::error::Result; pub type RWLockerImpl = Box; #[async_trait] pub trait RWLocker { async fn get_lock(&mut self, opts: &Options) -> Result; async fn un_lock(&mut self) -> Result<()>; async fn get_u_lock(&mut self, opts: &Options) -> Result; async fn un_r_lock(&mut self) -> Result<()>; } #[derive(Debug)] struct NsLock { reference: usize, lock: LRWMutex, } #[derive(Debug, Default)] pub struct NsLockMap { is_dist_erasure: bool, lock_map: RwLock>, } impl NsLockMap { pub fn new(is_dist_erasure: bool) -> Self { Self { is_dist_erasure, ..Default::default() } } async fn lock( &mut self, volume: &String, path: &String, lock_source: &str, ops_id: &str, read_lock: bool, timeout: Duration, ) -> bool { let resource = Path::new(volume).join(path).to_str().unwrap().to_string(); let mut w_lock_map = self.lock_map.write().await; let nslk = w_lock_map.entry(resource.clone()).or_insert(NsLock { reference: 0, lock: LRWMutex::default(), }); nslk.reference += 1; let locked = if read_lock { nslk.lock.get_r_lock(ops_id, lock_source, &timeout).await } else { nslk.lock.get_lock(ops_id, lock_source, &timeout).await }; if !locked { nslk.reference -= 1; if nslk.reference == 0 { w_lock_map.remove(&resource); } } locked } async fn un_lock(&mut self, volume: &String, path: &String, read_lock: bool) { let resource = Path::new(volume).join(path).to_str().unwrap().to_string(); let mut w_lock_map = self.lock_map.write().await; if let Some(nslk) = w_lock_map.get_mut(&resource) { if read_lock { nslk.lock.un_r_lock().await; } else { nslk.lock.un_lock().await; } nslk.reference -= 0; if nslk.reference == 0 { w_lock_map.remove(&resource); } } } } pub struct WrapperLocker(pub Arc>); impl Drop for WrapperLocker { fn drop(&mut self) { let inner = self.0.clone(); tokio::spawn(async move { let _ = inner.write().await.un_lock().await; }); } } pub async fn new_nslock( ns: Arc>, owner: String, volume: String, paths: Vec, lockers: Vec, ) -> WrapperLocker { if ns.read().await.is_dist_erasure { let names = paths .iter() .map(|path| Path::new(&volume).join(path).to_str().unwrap().to_string()) .collect(); return WrapperLocker(Arc::new(RwLock::new(Box::new(DistLockInstance::new(owner, names, lockers))))); } WrapperLocker(Arc::new(RwLock::new(Box::new(LocalLockInstance::new(ns, volume, paths))))) } struct DistLockInstance { lock: Box, ops_id: String, } impl DistLockInstance { fn new(owner: String, names: Vec, lockers: Vec) -> Self { let ops_id = Uuid::new_v4().to_string(); Self { lock: Box::new(DRWMutex::new(owner, names, lockers)), ops_id, } } } #[async_trait] impl RWLocker for DistLockInstance { async fn get_lock(&mut self, opts: &Options) -> Result { let source = "".to_string(); Ok(self.lock.get_lock(&self.ops_id, &source, opts).await) } async fn un_lock(&mut self) -> Result<()> { self.lock.un_lock().await; Ok(()) } async fn get_u_lock(&mut self, opts: &Options) -> Result { let source = "".to_string(); Ok(self.lock.get_r_lock(&self.ops_id, &source, opts).await) } async fn un_r_lock(&mut self) -> Result<()> { self.lock.un_r_lock().await; Ok(()) } } struct LocalLockInstance { ns: Arc>, volume: String, paths: Vec, ops_id: String, } impl LocalLockInstance { fn new(ns: Arc>, volume: String, paths: Vec) -> Self { let ops_id = Uuid::new_v4().to_string(); Self { ns, volume, paths, ops_id, } } } #[async_trait] impl RWLocker for LocalLockInstance { async fn get_lock(&mut self, opts: &Options) -> Result { let source = "".to_string(); let read_lock = false; let mut success = vec![false; self.paths.len()]; for (idx, path) in self.paths.iter().enumerate() { if !self .ns .write() .await .lock(&self.volume, path, &source, &self.ops_id, read_lock, opts.timeout) .await { for (i, x) in success.iter().enumerate() { if *x { self.ns.write().await.un_lock(&self.volume, &self.paths[i], read_lock).await; } } return Ok(false); } success[idx] = true; } Ok(true) } async fn un_lock(&mut self) -> Result<()> { let read_lock = false; for path in self.paths.iter() { self.ns.write().await.un_lock(&self.volume, path, read_lock).await; } Ok(()) } async fn get_u_lock(&mut self, opts: &Options) -> Result { let source = "".to_string(); let read_lock = true; let mut success = Vec::with_capacity(self.paths.len()); for (idx, path) in self.paths.iter().enumerate() { if !self .ns .write() .await .lock(&self.volume, path, &source, &self.ops_id, read_lock, opts.timeout) .await { for (i, x) in success.iter().enumerate() { if *x { self.ns.write().await.un_lock(&self.volume, &self.paths[i], read_lock).await; } } return Ok(false); } success[idx] = true; } Ok(true) } async fn un_r_lock(&mut self) -> Result<()> { let read_lock = true; for path in self.paths.iter() { self.ns.write().await.un_lock(&self.volume, path, read_lock).await; } Ok(()) } } #[cfg(test)] mod test { use std::{sync::Arc, time::Duration}; use common::error::Result; use tokio::sync::RwLock; use crate::{ drwmutex::Options, namespace_lock::{new_nslock, NsLockMap}, }; #[tokio::test] async fn test_local_instance() -> Result<()> { let ns_lock_map = Arc::new(RwLock::new(NsLockMap::default())); let ns = new_nslock( Arc::clone(&ns_lock_map), "local".to_string(), "test".to_string(), vec!["foo".to_string()], Vec::new(), ) .await; let result = ns.0.write() .await .get_lock(&Options { timeout: Duration::from_secs(5), retry_interval: Duration::from_secs(1), }) .await?; assert!(result); Ok(()) } }