use crate::error::{GunError, GunResult};
use crate::state::Node;
use async_trait::async_trait;
use parking_lot::RwLock;
use std::collections::{HashMap, HashSet};
use std::fs;
use std::io::{Read, Write};
use std::path::PathBuf;
#[async_trait]
pub trait Storage: Send + Sync {
async fn get(&self, soul: &str) -> GunResult<Option<Node>>;
async fn put(&self, soul: &str, node: &Node) -> GunResult<()>;
async fn has(&self, soul: &str) -> GunResult<bool>;
}
pub struct MemoryStorage {
data: RwLock<HashMap<String, Node>>,
}
impl MemoryStorage {
pub fn new() -> Self {
Self {
data: RwLock::new(HashMap::new()),
}
}
}
#[async_trait]
impl Storage for MemoryStorage {
async fn get(&self, soul: &str) -> GunResult<Option<Node>> {
let data = self.data.read();
Ok(data.get(soul).cloned())
}
async fn put(&self, soul: &str, node: &Node) -> GunResult<()> {
let mut data = self.data.write();
data.insert(soul.to_string(), node.clone());
Ok(())
}
async fn has(&self, soul: &str) -> GunResult<bool> {
let data = self.data.read();
Ok(data.contains_key(soul))
}
}
impl Default for MemoryStorage {
fn default() -> Self {
Self::new()
}
}
pub struct SledStorage {
db: sled::Db,
}
impl SledStorage {
pub fn new(path: &str) -> GunResult<Self> {
let db = sled::open(path)?;
Ok(Self { db })
}
}
#[async_trait]
impl Storage for SledStorage {
async fn get(&self, soul: &str) -> GunResult<Option<Node>> {
match self.db.get(soul)? {
Some(ivec) => {
let json_str = String::from_utf8(ivec.to_vec())
.map_err(|e| GunError::InvalidData(format!("Invalid UTF-8: {}", e)))?;
let node: Node = serde_json::from_str(&json_str)?;
Ok(Some(node))
}
None => Ok(None),
}
}
async fn put(&self, soul: &str, node: &Node) -> GunResult<()> {
let json_str = serde_json::to_string(node)?;
self.db.insert(soul, json_str.as_bytes())?;
self.db.flush_async().await?;
Ok(())
}
async fn has(&self, soul: &str) -> GunResult<bool> {
Ok(self.db.contains_key(soul)?)
}
}
pub struct LocalStorage {
data_dir: PathBuf,
cache: RwLock<HashMap<String, Node>>, dirty: RwLock<HashSet<String>>, }
impl LocalStorage {
pub fn new(data_dir: &str) -> GunResult<Self> {
let path = PathBuf::from(data_dir);
fs::create_dir_all(&path).map_err(|e| {
GunError::Io(std::io::Error::other(format!(
"Failed to create storage directory: {}",
e
)))
})?;
let cache = Self::load_all(&path)?;
Ok(Self {
data_dir: path,
cache: RwLock::new(cache),
dirty: RwLock::new(HashSet::new()),
})
}
fn load_all(path: &PathBuf) -> GunResult<HashMap<String, Node>> {
let mut data = HashMap::new();
if let Ok(entries) = fs::read_dir(path) {
for entry in entries.flatten() {
let file_path = entry.path();
if file_path.is_file() {
if let Some(file_name) = file_path.file_name() {
if let Some(soul) = file_name.to_str() {
let soul = urlencoding::decode(soul)
.unwrap_or(std::borrow::Cow::Borrowed(soul))
.into_owned();
if let Ok(node) = Self::load_file(&file_path) {
data.insert(soul, node);
}
}
}
}
}
}
Ok(data)
}
fn load_file(path: &PathBuf) -> GunResult<Node> {
let mut file = fs::File::open(path)?;
let mut contents = String::new();
file.read_to_string(&mut contents)?;
let node: Node = serde_json::from_str(&contents)?;
Ok(node)
}
fn save_file(&self, soul: &str, node: &Node) -> GunResult<()> {
let encoded_soul = urlencoding::encode(soul);
let file_path = self.data_dir.join(encoded_soul.as_ref());
let json_str = serde_json::to_string_pretty(node).map_err(GunError::Serialization)?;
let temp_path = file_path.with_extension("tmp");
let mut file = fs::File::create(&temp_path)?;
file.write_all(json_str.as_bytes())?;
file.sync_all()?;
drop(file);
fs::rename(&temp_path, &file_path)?;
Ok(())
}
pub async fn flush(&self) -> GunResult<()> {
let dirty_keys: Vec<String> = {
let dirty = self.dirty.read();
dirty.iter().cloned().collect()
};
let cache = self.cache.read();
for soul in dirty_keys {
if let Some(node) = cache.get(&soul) {
if let Err(e) = self.save_file(&soul, node) {
eprintln!("Error saving {} to disk: {}", soul, e);
}
}
}
let mut dirty = self.dirty.write();
dirty.clear();
Ok(())
}
}
#[async_trait]
impl Storage for LocalStorage {
async fn get(&self, soul: &str) -> GunResult<Option<Node>> {
let cache = self.cache.read();
Ok(cache.get(soul).cloned())
}
async fn put(&self, soul: &str, node: &Node) -> GunResult<()> {
{
let mut cache = self.cache.write();
cache.insert(soul.to_string(), node.clone());
}
{
let mut dirty = self.dirty.write();
dirty.insert(soul.to_string());
}
self.save_file(soul, node)?;
let mut dirty = self.dirty.write();
dirty.remove(soul);
Ok(())
}
async fn has(&self, soul: &str) -> GunResult<bool> {
let cache = self.cache.read();
Ok(cache.contains_key(soul))
}
}
impl Drop for LocalStorage {
fn drop(&mut self) {
let dirty_keys: Vec<String> = {
let dirty = self.dirty.read();
dirty.iter().cloned().collect()
};
if !dirty_keys.is_empty() {
let cache = self.cache.read();
for soul in dirty_keys {
if let Some(node) = cache.get(&soul) {
let _ = self.save_file(&soul, node);
}
}
}
}
}