use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs;
use std::io::{Read, Write};
use std::path::{Path, PathBuf};
use std::time::SystemTime;
use torsh_core::error::{Result, TorshError};
pub trait StorageBackend: Send + Sync {
fn put(&mut self, key: &str, data: &[u8]) -> Result<()>;
fn get(&self, key: &str) -> Result<Vec<u8>>;
fn delete(&mut self, key: &str) -> Result<()>;
fn exists(&self, key: &str) -> Result<bool>;
fn list(&self, prefix: &str) -> Result<Vec<StorageObject>>;
fn get_metadata(&self, key: &str) -> Result<StorageObject>;
fn copy(&mut self, from_key: &str, to_key: &str) -> Result<()> {
let data = self.get(from_key)?;
self.put(to_key, &data)
}
fn backend_type(&self) -> &str;
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StorageObject {
pub key: String,
pub size: u64,
pub last_modified: SystemTime,
pub content_type: Option<String>,
pub etag: Option<String>,
pub metadata: HashMap<String, String>,
}
pub struct LocalStorage {
base_path: PathBuf,
}
impl LocalStorage {
pub fn new(base_path: PathBuf) -> Result<Self> {
if !base_path.exists() {
fs::create_dir_all(&base_path).map_err(|e| {
TorshError::IoError(format!(
"Failed to create storage directory {}: {}",
base_path.display(),
e
))
})?;
}
Ok(Self { base_path })
}
fn get_path(&self, key: &str) -> PathBuf {
self.base_path.join(key)
}
fn ensure_parent_dir(&self, key: &str) -> Result<()> {
let path = self.get_path(key);
if let Some(parent) = path.parent() {
if !parent.exists() {
fs::create_dir_all(parent).map_err(|e| {
TorshError::IoError(format!("Failed to create parent directory: {}", e))
})?;
}
}
Ok(())
}
}
impl StorageBackend for LocalStorage {
fn put(&mut self, key: &str, data: &[u8]) -> Result<()> {
self.ensure_parent_dir(key)?;
let path = self.get_path(key);
let mut file = fs::File::create(&path).map_err(|e| {
TorshError::IoError(format!("Failed to create file {}: {}", path.display(), e))
})?;
file.write_all(data).map_err(|e| {
TorshError::IoError(format!("Failed to write to file {}: {}", path.display(), e))
})?;
Ok(())
}
fn get(&self, key: &str) -> Result<Vec<u8>> {
let path = self.get_path(key);
if !path.exists() {
return Err(TorshError::InvalidArgument(format!(
"Storage key not found: {}",
key
)));
}
let mut file = fs::File::open(&path).map_err(|e| {
TorshError::IoError(format!("Failed to open file {}: {}", path.display(), e))
})?;
let mut data = Vec::new();
file.read_to_end(&mut data).map_err(|e| {
TorshError::IoError(format!("Failed to read file {}: {}", path.display(), e))
})?;
Ok(data)
}
fn delete(&mut self, key: &str) -> Result<()> {
let path = self.get_path(key);
if path.exists() {
fs::remove_file(&path).map_err(|e| {
TorshError::IoError(format!("Failed to delete file {}: {}", path.display(), e))
})?;
}
Ok(())
}
fn exists(&self, key: &str) -> Result<bool> {
let path = self.get_path(key);
Ok(path.exists())
}
fn list(&self, prefix: &str) -> Result<Vec<StorageObject>> {
let prefix_path = self.get_path(prefix);
if !prefix_path.exists() {
return Ok(Vec::new());
}
let mut objects = Vec::new();
fn walk_dir(dir: &Path, base: &Path, objects: &mut Vec<StorageObject>) -> Result<()> {
if dir.is_dir() {
for entry in fs::read_dir(dir)
.map_err(|e| TorshError::IoError(format!("Failed to read directory: {}", e)))?
{
let entry = entry.map_err(|e| {
TorshError::IoError(format!("Failed to read directory entry: {}", e))
})?;
let path = entry.path();
if path.is_file() {
let metadata = fs::metadata(&path).map_err(|e| {
TorshError::IoError(format!("Failed to get metadata: {}", e))
})?;
let relative_path = path
.strip_prefix(base)
.map_err(|e| {
TorshError::InvalidArgument(format!("Invalid path: {}", e))
})?
.to_string_lossy()
.to_string();
objects.push(StorageObject {
key: relative_path,
size: metadata.len(),
last_modified: metadata
.modified()
.unwrap_or_else(|_| SystemTime::now()),
content_type: None,
etag: None,
metadata: HashMap::new(),
});
} else if path.is_dir() {
walk_dir(&path, base, objects)?;
}
}
}
Ok(())
}
walk_dir(&prefix_path, &self.base_path, &mut objects)?;
Ok(objects)
}
fn get_metadata(&self, key: &str) -> Result<StorageObject> {
let path = self.get_path(key);
if !path.exists() {
return Err(TorshError::InvalidArgument(format!(
"Storage key not found: {}",
key
)));
}
let metadata = fs::metadata(&path).map_err(|e| {
TorshError::IoError(format!("Failed to get metadata for {}: {}", key, e))
})?;
Ok(StorageObject {
key: key.to_string(),
size: metadata.len(),
last_modified: metadata.modified().unwrap_or_else(|_| SystemTime::now()),
content_type: None,
etag: None,
metadata: HashMap::new(),
})
}
fn backend_type(&self) -> &str {
"local"
}
}
pub struct StorageManager {
backend: Box<dyn StorageBackend>,
cache: HashMap<String, CachedObject>,
cache_size_limit: usize,
current_cache_size: usize,
retry_count: u32,
stats: StorageStats,
}
#[derive(Clone)]
struct CachedObject {
data: Vec<u8>,
accessed_at: SystemTime,
access_count: u64,
}
#[derive(Debug, Default, Clone)]
pub struct StorageStats {
pub gets: u64,
pub puts: u64,
pub deletes: u64,
pub bytes_read: u64,
pub bytes_written: u64,
pub cache_hits: u64,
pub cache_misses: u64,
}
impl StorageManager {
pub fn new(backend: Box<dyn StorageBackend>) -> Self {
Self {
backend,
cache: HashMap::new(),
cache_size_limit: 100 * 1024 * 1024, current_cache_size: 0,
retry_count: 3,
stats: StorageStats::default(),
}
}
pub fn with_cache_size(mut self, size_bytes: usize) -> Self {
self.cache_size_limit = size_bytes;
self
}
pub fn with_retry_count(mut self, count: u32) -> Self {
self.retry_count = count;
self
}
pub fn put(&mut self, key: &str, data: &[u8]) -> Result<()> {
let mut last_error = None;
for attempt in 0..=self.retry_count {
match self.backend.put(key, data) {
Ok(()) => {
self.stats.puts += 1;
self.stats.bytes_written += data.len() as u64;
if self.cache.contains_key(key) {
self.put_in_cache(key, data);
}
return Ok(());
}
Err(e) => {
last_error = Some(e);
if attempt < self.retry_count {
let backoff_ms = 100 * 2u64.pow(attempt);
std::thread::sleep(std::time::Duration::from_millis(backoff_ms));
}
}
}
}
Err(last_error.expect("last_error is set when retries exhausted"))
}
pub fn get(&mut self, key: &str) -> Result<Vec<u8>> {
if let Some(cached) = self.cache.get_mut(key) {
cached.accessed_at = SystemTime::now();
cached.access_count += 1;
self.stats.cache_hits += 1;
self.stats.gets += 1;
return Ok(cached.data.clone());
}
self.stats.cache_misses += 1;
let mut last_error = None;
for attempt in 0..=self.retry_count {
match self.backend.get(key) {
Ok(data) => {
self.stats.gets += 1;
self.stats.bytes_read += data.len() as u64;
self.put_in_cache(key, &data);
return Ok(data);
}
Err(e) => {
last_error = Some(e);
if attempt < self.retry_count {
let backoff_ms = 100 * 2u64.pow(attempt);
std::thread::sleep(std::time::Duration::from_millis(backoff_ms));
}
}
}
}
Err(last_error.expect("last_error is set when retries exhausted"))
}
pub fn delete(&mut self, key: &str) -> Result<()> {
if let Some(cached) = self.cache.remove(key) {
self.current_cache_size -= cached.data.len();
}
let mut last_error = None;
for attempt in 0..=self.retry_count {
match self.backend.delete(key) {
Ok(()) => {
self.stats.deletes += 1;
return Ok(());
}
Err(e) => {
last_error = Some(e);
if attempt < self.retry_count {
let backoff_ms = 100 * 2u64.pow(attempt);
std::thread::sleep(std::time::Duration::from_millis(backoff_ms));
}
}
}
}
Err(last_error.expect("last_error is set when retries exhausted"))
}
pub fn exists(&self, key: &str) -> Result<bool> {
if self.cache.contains_key(key) {
return Ok(true);
}
self.backend.exists(key)
}
pub fn list(&self, prefix: &str) -> Result<Vec<StorageObject>> {
self.backend.list(prefix)
}
pub fn get_metadata(&self, key: &str) -> Result<StorageObject> {
self.backend.get_metadata(key)
}
pub fn copy(&mut self, from_key: &str, to_key: &str) -> Result<()> {
self.backend.copy(from_key, to_key)
}
pub fn clear_cache(&mut self) {
self.cache.clear();
self.current_cache_size = 0;
}
pub fn stats(&self) -> &StorageStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats = StorageStats::default();
}
fn put_in_cache(&mut self, key: &str, data: &[u8]) {
while self.current_cache_size + data.len() > self.cache_size_limit && !self.cache.is_empty()
{
if let Some(lru_key) = self.find_lru_key() {
if let Some(removed) = self.cache.remove(&lru_key) {
self.current_cache_size -= removed.data.len();
}
} else {
break;
}
}
if data.len() <= self.cache_size_limit {
self.current_cache_size += data.len();
self.cache.insert(
key.to_string(),
CachedObject {
data: data.to_vec(),
accessed_at: SystemTime::now(),
access_count: 1,
},
);
}
}
fn find_lru_key(&self) -> Option<String> {
self.cache
.iter()
.min_by_key(|(_, obj)| obj.accessed_at)
.map(|(key, _)| key.clone())
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
#[test]
fn test_local_storage_creation() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
assert_eq!(storage.backend_type(), "local");
}
#[test]
fn test_local_storage_put_get() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let mut storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
let data = b"test package data";
storage.put("test/package.bin", data).unwrap();
let retrieved = storage.get("test/package.bin").unwrap();
assert_eq!(retrieved, data);
}
#[test]
fn test_local_storage_exists() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let mut storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
assert!(!storage.exists("nonexistent").unwrap());
storage.put("exists", b"data").unwrap();
assert!(storage.exists("exists").unwrap());
}
#[test]
fn test_local_storage_delete() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let mut storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
storage.put("to_delete", b"data").unwrap();
assert!(storage.exists("to_delete").unwrap());
storage.delete("to_delete").unwrap();
assert!(!storage.exists("to_delete").unwrap());
}
#[test]
fn test_local_storage_list() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let mut storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
storage.put("models/model1.bin", b"data1").unwrap();
storage.put("models/model2.bin", b"data2").unwrap();
storage.put("other/file.txt", b"data3").unwrap();
let models = storage.list("models/").unwrap();
assert_eq!(models.len(), 2);
let all = storage.list("").unwrap();
assert_eq!(all.len(), 3);
}
#[test]
fn test_local_storage_metadata() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let mut storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
let data = b"test data";
storage.put("metadata_test", data).unwrap();
let metadata = storage.get_metadata("metadata_test").unwrap();
assert_eq!(metadata.size, data.len() as u64);
assert_eq!(metadata.key, "metadata_test");
}
#[test]
fn test_storage_manager_caching() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
let mut manager = StorageManager::new(Box::new(storage)).with_cache_size(1024 * 1024);
let data = b"cached data";
manager.put("cache_test", data).unwrap();
let retrieved1 = manager.get("cache_test").unwrap();
assert_eq!(retrieved1, data);
assert_eq!(manager.stats().cache_misses, 1);
assert_eq!(manager.stats().cache_hits, 0);
let retrieved2 = manager.get("cache_test").unwrap();
assert_eq!(retrieved2, data);
assert_eq!(manager.stats().cache_misses, 1);
assert_eq!(manager.stats().cache_hits, 1);
}
#[test]
fn test_storage_manager_cache_eviction() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
let mut manager = StorageManager::new(Box::new(storage)).with_cache_size(100);
manager.put("large1", &vec![1u8; 60]).unwrap();
manager.put("large2", &vec![2u8; 60]).unwrap();
manager.get("large1").unwrap();
manager.get("large2").unwrap();
assert!(manager.current_cache_size <= 100);
}
#[test]
fn test_storage_manager_stats() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
let mut manager = StorageManager::new(Box::new(storage));
let data = b"test data";
manager.put("stats_test", data).unwrap();
manager.get("stats_test").unwrap();
manager.delete("stats_test").unwrap();
let stats = manager.stats();
assert_eq!(stats.puts, 1);
assert_eq!(stats.gets, 1);
assert_eq!(stats.deletes, 1);
assert_eq!(stats.bytes_written, data.len() as u64);
assert_eq!(stats.bytes_read, data.len() as u64);
}
#[test]
fn test_storage_manager_copy() {
let temp_dir = TempDir::new().expect("Failed to create temp directory for test");
let storage = LocalStorage::new(temp_dir.path().to_path_buf())
.expect("Failed to create storage path");
let mut manager = StorageManager::new(Box::new(storage));
let data = b"copy test data";
manager.put("source", data).unwrap();
manager.copy("source", "destination").unwrap();
let copied = manager.get("destination").unwrap();
assert_eq!(copied, data);
}
}