use crate::{Result, StorageError};
use std::collections::HashMap;
use std::fs::File;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use std::sync::{Arc, RwLock};
struct FileRef {
file: Arc<File>,
ref_count: AtomicUsize,
delete_pending: AtomicBool,
}
#[derive(Clone)]
pub struct FileRefManager {
refs: Arc<RwLock<HashMap<PathBuf, Arc<FileRef>>>>,
}
impl FileRefManager {
pub fn new() -> Self {
Self {
refs: Arc::new(RwLock::new(HashMap::new())),
}
}
pub fn open<P: AsRef<Path>>(&self, path: P) -> Result<FileHandle> {
let path = path.as_ref().to_path_buf();
let mut refs = self.refs.write()
.map_err(|_| StorageError::Lock("FileRefManager lock poisoned".into()))?;
let file_ref = if let Some(existing) = refs.get(&path) {
existing.ref_count.fetch_add(1, Ordering::SeqCst);
existing.clone()
} else {
let file = File::open(&path)?;
let file_ref = Arc::new(FileRef {
file: Arc::new(file),
ref_count: AtomicUsize::new(1),
delete_pending: AtomicBool::new(false),
});
refs.insert(path.clone(), file_ref.clone());
file_ref
};
Ok(FileHandle {
file: file_ref.file.clone(),
path: path.clone(),
file_ref,
manager: self.clone(),
})
}
pub fn acquire<P: AsRef<Path>>(&self, path: P) -> Result<FileHandle> {
self.open(path)
}
pub fn mark_for_deletion<P: AsRef<Path>>(&self, path: P) -> Result<()> {
let path = path.as_ref();
let refs = self.refs.read()
.map_err(|_| StorageError::Lock("FileRefManager lock poisoned".into()))?;
if let Some(file_ref) = refs.get(path) {
file_ref.delete_pending.store(true, Ordering::SeqCst);
}
Ok(())
}
fn close(&self, path: &Path, file_ref: &Arc<FileRef>) {
let count = file_ref.ref_count.fetch_sub(1, Ordering::SeqCst);
if count == 1 && file_ref.delete_pending.load(Ordering::SeqCst) {
if let Ok(mut refs) = self.refs.write() {
refs.remove(path);
}
let _ = std::fs::remove_file(path);
}
}
pub fn ref_count<P: AsRef<Path>>(&self, path: P) -> usize {
let path = path.as_ref();
if let Ok(refs) = self.refs.read() {
if let Some(file_ref) = refs.get(path) {
return file_ref.ref_count.load(Ordering::SeqCst);
}
}
0
}
pub fn is_pending_deletion<P: AsRef<Path>>(&self, path: P) -> bool {
let path = path.as_ref();
if let Ok(refs) = self.refs.read() {
if let Some(file_ref) = refs.get(path) {
return file_ref.delete_pending.load(Ordering::SeqCst);
}
}
false
}
}
impl Default for FileRefManager {
fn default() -> Self {
Self::new()
}
}
pub struct FileHandle {
pub file: Arc<File>,
path: PathBuf,
file_ref: Arc<FileRef>,
manager: FileRefManager,
}
impl Drop for FileHandle {
fn drop(&mut self) {
self.manager.close(&self.path, &self.file_ref);
}
}
impl FileHandle {
pub fn path(&self) -> &Path {
&self.path
}
pub fn file(&self) -> &Arc<File> {
&self.file
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use tempfile::NamedTempFile;
#[test]
fn test_open_and_close() {
let manager = FileRefManager::new();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"test data").unwrap();
temp_file.flush().unwrap();
let path = temp_file.path();
{
let handle = manager.open(path).unwrap();
assert_eq!(manager.ref_count(path), 1);
assert_eq!(handle.path(), path);
}
assert_eq!(manager.ref_count(path), 0);
}
#[test]
fn test_multiple_references() {
let manager = FileRefManager::new();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"test data").unwrap();
temp_file.flush().unwrap();
let path = temp_file.path();
let handle1 = manager.open(path).unwrap();
assert_eq!(manager.ref_count(path), 1);
let handle2 = manager.open(path).unwrap();
assert_eq!(manager.ref_count(path), 2);
let handle3 = manager.open(path).unwrap();
assert_eq!(manager.ref_count(path), 3);
drop(handle1);
assert_eq!(manager.ref_count(path), 2);
drop(handle2);
drop(handle3);
assert_eq!(manager.ref_count(path), 0);
}
#[test]
fn test_deferred_deletion() {
let manager = FileRefManager::new();
let temp_file = NamedTempFile::new().unwrap();
let path = temp_file.into_temp_path();
std::fs::write(&path, b"test data").unwrap();
let handle = manager.open(&path).unwrap();
assert_eq!(manager.ref_count(&path), 1);
assert!(path.exists());
manager.mark_for_deletion(&path).unwrap();
assert!(manager.is_pending_deletion(&path));
assert!(path.exists());
drop(handle);
std::thread::sleep(std::time::Duration::from_millis(100));
assert!(!path.exists()); assert_eq!(manager.ref_count(&path), 0);
}
#[test]
fn test_concurrent_access() {
use std::thread;
let manager = FileRefManager::new();
let mut temp_file = NamedTempFile::new().unwrap();
temp_file.write_all(b"test data").unwrap();
temp_file.flush().unwrap();
let path = temp_file.path().to_path_buf();
let manager_clone = manager.clone();
let path_clone = path.clone();
let handle1 = manager.open(&path).unwrap();
let thread = thread::spawn(move || {
let handle2 = manager_clone.open(&path_clone).unwrap();
assert_eq!(manager_clone.ref_count(&path_clone), 2);
handle2
});
let handle2 = thread.join().unwrap();
assert_eq!(manager.ref_count(&path), 2);
drop(handle1);
drop(handle2);
assert_eq!(manager.ref_count(&path), 0);
}
}