use serde::{de::DeserializeOwned, Serialize};
use std::{
ffi::{OsStr, OsString},
fmt,
fs::{self, File},
io::{self, BufWriter},
ops::{Deref, DerefMut},
path::{Path, PathBuf},
sync::{RwLock, RwLockReadGuard, RwLockWriteGuard},
};
use tracing::{error, info};
pub trait DataStore: Default + Serialize {
fn open<P>(db: P) -> AtomicDatabase<Self>
where
P: AsRef<Path>,
Self: DeserializeOwned,
{
let db_path = db.as_ref();
if db_path.exists() {
AtomicDatabase::load(db_path).unwrap()
} else {
AtomicDatabase::create(db_path).unwrap()
}
}
fn open_in_memory() -> AtomicDatabase<Self>
where
Self: DeserializeOwned,
{
AtomicDatabase::load_in_memory().unwrap()
}
fn load(file: impl io::Read) -> std::io::Result<Self>
where
Self: Sized,
Self: DeserializeOwned,
{
Ok(serde_json::from_reader(file)?)
}
fn save(&self, file: impl io::Write) -> std::io::Result<()> {
let writer = BufWriter::new(file);
serde_json::to_writer_pretty(writer, self)?;
Ok(())
}
}
pub struct AtomicDatabase<T: DataStore> {
path: Option<PathBuf>,
tmp: Option<PathBuf>,
data: RwLock<T>,
}
impl<T: DataStore + DeserializeOwned> AtomicDatabase<T> {
pub fn load_in_memory() -> Result<Self, std::io::Error> {
Ok(Self {
path: None,
tmp: None,
data: RwLock::new(T::default()),
})
}
pub fn load(path: &Path) -> Result<Self, std::io::Error> {
let new_path = path.with_extension("json");
let tmp = Self::tmp_path(&new_path)?;
let file = File::open(path)?;
let data = T::load(file)?;
atomic_write(&tmp, &new_path, &data)?;
Ok(Self {
path: Some(new_path),
tmp: Some(tmp),
data: RwLock::new(data),
})
}
pub fn create(path: &Path) -> Result<Self, std::io::Error> {
let tmp = Self::tmp_path(path)?;
let data = Default::default();
atomic_write(&tmp, path, &data)?;
Ok(Self {
path: Some(path.into()),
tmp: Some(tmp),
data: RwLock::new(data),
})
}
pub fn read(&self) -> AtomicDatabaseRead<'_, T> {
AtomicDatabaseRead {
data: self.data.read().unwrap(),
}
}
pub fn write(&self) -> AtomicDatabaseWrite<'_, T> {
AtomicDatabaseWrite {
path: self.path.as_deref(),
tmp: self.tmp.as_deref(),
data: self.data.write().unwrap(),
}
}
fn tmp_path(path: &Path) -> Result<PathBuf, std::io::Error> {
let mut tmp_name = OsString::from(".");
tmp_name.push(path.file_name().unwrap_or(OsStr::new("db")));
tmp_name.push("~");
let tmp = path.with_file_name(tmp_name);
if tmp.exists() {
error!(
"Found orphaned database temporary file '{tmp:?}'. The server has recently crashed or is already running. Delete this before continuing!"
);
return Err(std::io::Error::last_os_error());
}
Ok(tmp)
}
}
fn atomic_write<T: DataStore>(tmp: &Path, path: &Path, data: &T) -> Result<(), std::io::Error> {
{
let mut tmpfile = File::create(tmp)?;
data.save(&mut tmpfile)?;
tmpfile.sync_all()?; }
fs::rename(tmp, path)?;
Ok(())
}
impl<T: DataStore> fmt::Debug for AtomicDatabase<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("AtomicDatabase")
.field("file", &self.path)
.finish()
}
}
impl<T: DataStore> Drop for AtomicDatabase<T> {
fn drop(&mut self) {
if let Some(tmp) = &self.tmp {
if let Some(path) = &self.path {
info!("Saving database");
let guard = self.data.read().unwrap();
atomic_write(tmp, path, &*guard).unwrap();
}
}
}
}
pub struct AtomicDatabaseRead<'a, T: DataStore> {
data: RwLockReadGuard<'a, T>,
}
impl<'a, T: DataStore> Deref for AtomicDatabaseRead<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.data
}
}
pub struct AtomicDatabaseWrite<'a, T: DataStore> {
tmp: Option<&'a Path>,
path: Option<&'a Path>,
data: RwLockWriteGuard<'a, T>,
}
impl<'a, T: DataStore> Deref for AtomicDatabaseWrite<'a, T> {
type Target = T;
fn deref(&self) -> &Self::Target {
&self.data
}
}
impl<'a, T: DataStore> DerefMut for AtomicDatabaseWrite<'a, T> {
fn deref_mut(&mut self) -> &mut Self::Target {
&mut self.data
}
}
impl<'a, T: DataStore> Drop for AtomicDatabaseWrite<'a, T> {
fn drop(&mut self) {
if let Some(tmp) = self.tmp {
if let Some(path) = self.path {
info!("Saving database");
atomic_write(tmp, path, &*self.data).unwrap();
}
}
}
}