137 lines
3.0 KiB
Rust
137 lines
3.0 KiB
Rust
use std::{
|
|
fmt::Debug,
|
|
marker::PhantomData,
|
|
sync::atomic::{AtomicU64, Ordering},
|
|
};
|
|
|
|
use eyre::{eyre, WrapErr};
|
|
use redb::ReadableTable;
|
|
use serde::{de::DeserializeOwned, Serialize};
|
|
use snitch_lib::LogMsg;
|
|
|
|
use crate::Opt;
|
|
|
|
pub struct Database {
|
|
inner: redb::Database,
|
|
next_log_id: AtomicU64,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
struct Serialized<T> {
|
|
packed: Vec<u8>,
|
|
_type: PhantomData<T>,
|
|
}
|
|
|
|
pub(crate) fn open(opt: &Opt) -> eyre::Result<Database> {
|
|
let inner = redb::Database::create(&opt.db)
|
|
.context(eyre!("Failed to open database at {}", opt.db.display()))?;
|
|
|
|
let txn = inner.begin_write()?;
|
|
let next_log_id = txn
|
|
.open_table(table::LOG_MESSAGE)?
|
|
.range::<u64>(..)?
|
|
.next_back()
|
|
.transpose()
|
|
.ok()
|
|
.flatten()
|
|
.map(|(id, _)| id.value() + 1)
|
|
.unwrap_or_default();
|
|
txn.commit()?;
|
|
|
|
Ok(Database {
|
|
inner,
|
|
next_log_id: next_log_id.into(),
|
|
})
|
|
}
|
|
|
|
mod table {
|
|
use super::Serialized;
|
|
use redb::TableDefinition;
|
|
use snitch_lib::LogMsg;
|
|
|
|
pub const LOG_MESSAGE: TableDefinition<u64, Serialized<LogMsg>> =
|
|
TableDefinition::new("log_message");
|
|
}
|
|
|
|
impl Database {
|
|
pub fn write_log(&self, log: &LogMsg) -> eyre::Result<u64> {
|
|
let txn = self.inner.begin_write()?;
|
|
let id = self.next_log_id.fetch_add(1, Ordering::SeqCst);
|
|
txn.open_table(table::LOG_MESSAGE)?
|
|
.insert(id, Serialized::serialize(log))?;
|
|
txn.commit()?;
|
|
|
|
Ok(id)
|
|
}
|
|
|
|
pub fn get_all_log_messages(&self) -> eyre::Result<Vec<(u64, LogMsg)>> {
|
|
let txn = self.inner.begin_read()?;
|
|
txn.open_table(table::LOG_MESSAGE)?
|
|
.range::<u64>(..)?
|
|
.map(|result| {
|
|
let (id, msg) = result?;
|
|
let id = id.value();
|
|
let msg = msg.value();
|
|
|
|
msg.deserialize().map(|msg| (id, msg))
|
|
})
|
|
.collect()
|
|
}
|
|
}
|
|
|
|
impl<T> Serialized<T>
|
|
where
|
|
T: Debug + Serialize + DeserializeOwned,
|
|
{
|
|
pub fn serialize(value: &T) -> Self {
|
|
let packed = rmp_serde::to_vec(&value).unwrap();
|
|
Self {
|
|
packed,
|
|
_type: PhantomData,
|
|
}
|
|
}
|
|
|
|
pub fn deserialize(&self) -> eyre::Result<T> {
|
|
rmp_serde::from_slice(&self.packed).wrap_err("Failed to deserialize")
|
|
}
|
|
}
|
|
|
|
impl<T> redb::Value for Serialized<T>
|
|
where
|
|
T: Debug + Serialize + DeserializeOwned,
|
|
{
|
|
type SelfType<'a> = Self
|
|
where
|
|
Self: 'a;
|
|
|
|
type AsBytes<'a> = &'a [u8]
|
|
where
|
|
Self: 'a;
|
|
|
|
fn fixed_width() -> Option<usize> {
|
|
None
|
|
}
|
|
|
|
fn from_bytes<'a>(data: &'a [u8]) -> Self::SelfType<'a>
|
|
where
|
|
Self: 'a,
|
|
{
|
|
Serialized {
|
|
packed: data.to_vec(),
|
|
_type: PhantomData,
|
|
}
|
|
}
|
|
|
|
fn as_bytes<'a, 'b: 'a>(value: &'a Self::SelfType<'b>) -> Self::AsBytes<'a>
|
|
where
|
|
Self: 'a,
|
|
Self: 'b,
|
|
{
|
|
&value.packed
|
|
}
|
|
|
|
fn type_name() -> redb::TypeName {
|
|
redb::TypeName::new(std::any::type_name::<Serialized<T>>())
|
|
}
|
|
}
|