313 lines
8.8 KiB
Rust
313 lines
8.8 KiB
Rust
//! Local chunk storage
|
|
|
|
use dashmap::DashMap;
|
|
use std::path::PathBuf;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use thiserror::Error;
|
|
use tokio::fs;
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tracing::debug;
|
|
|
|
/// Errors from chunk storage operations
|
|
#[derive(Debug, Error)]
|
|
pub enum StorageError {
|
|
#[error("IO error: {0}")]
|
|
Io(#[from] std::io::Error),
|
|
|
|
#[error("Chunk not found: {0}")]
|
|
NotFound(String),
|
|
|
|
#[error("Storage capacity exceeded")]
|
|
CapacityExceeded,
|
|
}
|
|
|
|
pub type StorageResult<T> = Result<T, StorageError>;
|
|
|
|
/// Local filesystem-based chunk storage
|
|
pub struct LocalChunkStore {
|
|
/// Data directory
|
|
data_dir: PathBuf,
|
|
|
|
/// In-memory index of chunk sizes for fast lookups
|
|
chunk_sizes: DashMap<String, u64>,
|
|
|
|
/// Total bytes stored
|
|
total_bytes: AtomicU64,
|
|
|
|
/// Maximum capacity (0 = unlimited)
|
|
max_capacity: u64,
|
|
|
|
/// Number of chunks stored
|
|
chunk_count: AtomicU64,
|
|
}
|
|
|
|
impl LocalChunkStore {
|
|
/// Create a new local chunk store
|
|
pub async fn new(data_dir: PathBuf, max_capacity: u64) -> StorageResult<Self> {
|
|
// Ensure data directory exists
|
|
fs::create_dir_all(&data_dir).await?;
|
|
|
|
let store = Self {
|
|
data_dir,
|
|
chunk_sizes: DashMap::new(),
|
|
total_bytes: AtomicU64::new(0),
|
|
max_capacity,
|
|
chunk_count: AtomicU64::new(0),
|
|
};
|
|
|
|
// Scan existing chunks
|
|
store.scan_existing_chunks().await?;
|
|
|
|
Ok(store)
|
|
}
|
|
|
|
/// Scan existing chunks in the data directory
|
|
async fn scan_existing_chunks(&self) -> StorageResult<()> {
|
|
let mut entries = fs::read_dir(&self.data_dir).await?;
|
|
let mut total_bytes = 0u64;
|
|
let mut chunk_count = 0u64;
|
|
|
|
while let Some(entry) = entries.next_entry().await? {
|
|
let path = entry.path();
|
|
if path.is_file() {
|
|
if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
|
|
if let Ok(metadata) = entry.metadata().await {
|
|
let size = metadata.len();
|
|
self.chunk_sizes.insert(name.to_string(), size);
|
|
total_bytes += size;
|
|
chunk_count += 1;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
self.total_bytes.store(total_bytes, Ordering::SeqCst);
|
|
self.chunk_count.store(chunk_count, Ordering::SeqCst);
|
|
|
|
debug!(
|
|
total_bytes,
|
|
chunk_count,
|
|
"Scanned existing chunks"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Get the path for a chunk
|
|
fn chunk_path(&self, chunk_id: &str) -> PathBuf {
|
|
// Sanitize chunk_id to be a valid filename
|
|
let safe_id = chunk_id.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|'], "_");
|
|
self.data_dir.join(safe_id)
|
|
}
|
|
|
|
/// Store a chunk
|
|
pub async fn put(&self, chunk_id: &str, data: &[u8]) -> StorageResult<u64> {
|
|
let size = data.len() as u64;
|
|
|
|
// Check capacity
|
|
if self.max_capacity > 0 {
|
|
let current = self.total_bytes.load(Ordering::SeqCst);
|
|
if current + size > self.max_capacity {
|
|
return Err(StorageError::CapacityExceeded);
|
|
}
|
|
}
|
|
|
|
let path = self.chunk_path(chunk_id);
|
|
|
|
// Check if replacing existing chunk
|
|
let old_size = self.chunk_sizes.get(chunk_id).map(|v| *v).unwrap_or(0);
|
|
|
|
// Write data
|
|
let mut file = fs::File::create(&path).await?;
|
|
file.write_all(data).await?;
|
|
file.sync_all().await?;
|
|
|
|
// Update index
|
|
self.chunk_sizes.insert(chunk_id.to_string(), size);
|
|
|
|
// Update totals
|
|
if old_size > 0 {
|
|
// Replacing existing chunk
|
|
self.total_bytes.fetch_sub(old_size, Ordering::SeqCst);
|
|
} else {
|
|
// New chunk
|
|
self.chunk_count.fetch_add(1, Ordering::SeqCst);
|
|
}
|
|
self.total_bytes.fetch_add(size, Ordering::SeqCst);
|
|
|
|
debug!(chunk_id, size, "Stored chunk");
|
|
|
|
Ok(size)
|
|
}
|
|
|
|
/// Retrieve a chunk
|
|
pub async fn get(&self, chunk_id: &str) -> StorageResult<Vec<u8>> {
|
|
let path = self.chunk_path(chunk_id);
|
|
|
|
if !path.exists() {
|
|
return Err(StorageError::NotFound(chunk_id.to_string()));
|
|
}
|
|
|
|
let mut file = fs::File::open(&path).await?;
|
|
let mut data = Vec::new();
|
|
file.read_to_end(&mut data).await?;
|
|
|
|
debug!(chunk_id, size = data.len(), "Retrieved chunk");
|
|
|
|
Ok(data)
|
|
}
|
|
|
|
/// Delete a chunk
|
|
pub async fn delete(&self, chunk_id: &str) -> StorageResult<()> {
|
|
let path = self.chunk_path(chunk_id);
|
|
|
|
if let Some((_, size)) = self.chunk_sizes.remove(chunk_id) {
|
|
if path.exists() {
|
|
fs::remove_file(&path).await?;
|
|
}
|
|
self.total_bytes.fetch_sub(size, Ordering::SeqCst);
|
|
self.chunk_count.fetch_sub(1, Ordering::SeqCst);
|
|
debug!(chunk_id, "Deleted chunk");
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Check if a chunk exists
|
|
pub fn exists(&self, chunk_id: &str) -> bool {
|
|
self.chunk_sizes.contains_key(chunk_id)
|
|
}
|
|
|
|
/// Get the size of a chunk
|
|
pub fn size(&self, chunk_id: &str) -> Option<u64> {
|
|
self.chunk_sizes.get(chunk_id).map(|v| *v)
|
|
}
|
|
|
|
/// Get total bytes stored
|
|
pub fn total_bytes(&self) -> u64 {
|
|
self.total_bytes.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Get chunk count
|
|
pub fn chunk_count(&self) -> u64 {
|
|
self.chunk_count.load(Ordering::SeqCst)
|
|
}
|
|
|
|
/// Get maximum capacity
|
|
pub fn max_capacity(&self) -> u64 {
|
|
self.max_capacity
|
|
}
|
|
|
|
/// Get available capacity
|
|
pub fn available_bytes(&self) -> u64 {
|
|
if self.max_capacity == 0 {
|
|
u64::MAX
|
|
} else {
|
|
self.max_capacity.saturating_sub(self.total_bytes())
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use tempfile::TempDir;
|
|
|
|
async fn create_test_store() -> (LocalChunkStore, TempDir) {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let store = LocalChunkStore::new(temp_dir.path().to_path_buf(), 0)
|
|
.await
|
|
.unwrap();
|
|
(store, temp_dir)
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_put_get() {
|
|
let (store, _temp) = create_test_store().await;
|
|
|
|
let chunk_id = "test-chunk-1";
|
|
let data = vec![42u8; 1024];
|
|
|
|
let size = store.put(chunk_id, &data).await.unwrap();
|
|
assert_eq!(size, 1024);
|
|
|
|
let retrieved = store.get(chunk_id).await.unwrap();
|
|
assert_eq!(retrieved, data);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_delete() {
|
|
let (store, _temp) = create_test_store().await;
|
|
|
|
let chunk_id = "test-chunk-2";
|
|
let data = vec![42u8; 512];
|
|
|
|
store.put(chunk_id, &data).await.unwrap();
|
|
assert!(store.exists(chunk_id));
|
|
|
|
store.delete(chunk_id).await.unwrap();
|
|
assert!(!store.exists(chunk_id));
|
|
|
|
let result = store.get(chunk_id).await;
|
|
assert!(result.is_err());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_size_tracking() {
|
|
let (store, _temp) = create_test_store().await;
|
|
|
|
assert_eq!(store.total_bytes(), 0);
|
|
assert_eq!(store.chunk_count(), 0);
|
|
|
|
store.put("chunk1", &vec![0u8; 100]).await.unwrap();
|
|
assert_eq!(store.total_bytes(), 100);
|
|
assert_eq!(store.chunk_count(), 1);
|
|
|
|
store.put("chunk2", &vec![0u8; 200]).await.unwrap();
|
|
assert_eq!(store.total_bytes(), 300);
|
|
assert_eq!(store.chunk_count(), 2);
|
|
|
|
store.delete("chunk1").await.unwrap();
|
|
assert_eq!(store.total_bytes(), 200);
|
|
assert_eq!(store.chunk_count(), 1);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_capacity_limit() {
|
|
let temp_dir = TempDir::new().unwrap();
|
|
let store = LocalChunkStore::new(temp_dir.path().to_path_buf(), 1000)
|
|
.await
|
|
.unwrap();
|
|
|
|
// Should succeed
|
|
store.put("chunk1", &vec![0u8; 500]).await.unwrap();
|
|
|
|
// Should fail - would exceed capacity
|
|
let result = store.put("chunk2", &vec![0u8; 600]).await;
|
|
assert!(matches!(result, Err(StorageError::CapacityExceeded)));
|
|
|
|
// Should succeed - within remaining capacity
|
|
store.put("chunk2", &vec![0u8; 400]).await.unwrap();
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn test_replace_chunk() {
|
|
let (store, _temp) = create_test_store().await;
|
|
|
|
let chunk_id = "test-chunk";
|
|
|
|
store.put(chunk_id, &vec![0u8; 100]).await.unwrap();
|
|
assert_eq!(store.total_bytes(), 100);
|
|
assert_eq!(store.chunk_count(), 1);
|
|
|
|
// Replace with larger data
|
|
store.put(chunk_id, &vec![0u8; 200]).await.unwrap();
|
|
assert_eq!(store.total_bytes(), 200);
|
|
assert_eq!(store.chunk_count(), 1); // Still 1 chunk
|
|
|
|
// Replace with smaller data
|
|
store.put(chunk_id, &vec![0u8; 50]).await.unwrap();
|
|
assert_eq!(store.total_bytes(), 50);
|
|
assert_eq!(store.chunk_count(), 1);
|
|
}
|
|
}
|