5 changed files with 271 additions and 0 deletions
@ -0,0 +1,245 @@
@@ -0,0 +1,245 @@
|
||||
use crate::{ |
||||
database::{ |
||||
abstraction::{DatabaseEngine, Tree}, |
||||
Config, |
||||
}, |
||||
Result, |
||||
}; |
||||
use persy::{ByteVec, OpenOptions, Persy, Transaction, TransactionConfig, ValueMode}; |
||||
|
||||
use std::{ |
||||
collections::HashMap, |
||||
future::Future, |
||||
pin::Pin, |
||||
sync::{Arc, RwLock}, |
||||
}; |
||||
|
||||
use tokio::sync::oneshot::Sender; |
||||
use tracing::warn; |
||||
|
||||
pub struct PersyEngine { |
||||
persy: Persy, |
||||
} |
||||
|
||||
impl DatabaseEngine for PersyEngine { |
||||
fn open(config: &Config) -> Result<Arc<Self>> { |
||||
let mut cfg = persy::Config::new(); |
||||
cfg.change_cache_size((config.db_cache_capacity_mb * 1024.0 * 1024.0) as u64); |
||||
|
||||
let persy = OpenOptions::new() |
||||
.create(true) |
||||
.config(cfg) |
||||
.open(&format!("{}/db.persy", config.database_path))?; |
||||
Ok(Arc::new(PersyEngine { persy })) |
||||
} |
||||
|
||||
fn open_tree(self: &Arc<Self>, name: &'static str) -> Result<Arc<dyn Tree>> { |
||||
// Create if it doesn't exist
|
||||
if !self.persy.exists_index(name)? { |
||||
let mut tx = self.persy.begin()?; |
||||
tx.create_index::<ByteVec, ByteVec>(name, ValueMode::Replace)?; |
||||
tx.prepare()?.commit()?; |
||||
} |
||||
|
||||
Ok(Arc::new(PersyTree { |
||||
persy: self.persy.clone(), |
||||
name: name.to_owned(), |
||||
watchers: RwLock::new(HashMap::new()), |
||||
})) |
||||
} |
||||
|
||||
fn flush(self: &Arc<Self>) -> Result<()> { |
||||
Ok(()) |
||||
} |
||||
} |
||||
|
||||
pub struct PersyTree { |
||||
persy: Persy, |
||||
name: String, |
||||
watchers: RwLock<HashMap<Vec<u8>, Vec<Sender<()>>>>, |
||||
} |
||||
|
||||
impl PersyTree { |
||||
fn begin(&self) -> Result<Transaction> { |
||||
Ok(self |
||||
.persy |
||||
.begin_with(TransactionConfig::new().set_background_sync(true))?) |
||||
} |
||||
} |
||||
|
||||
impl Tree for PersyTree { |
||||
#[tracing::instrument(skip(self, key))] |
||||
fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> { |
||||
let result = self |
||||
.persy |
||||
.get::<ByteVec, ByteVec>(&self.name, &ByteVec::from(key))? |
||||
.next() |
||||
.map(|v| (*v).to_owned()); |
||||
Ok(result) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, key, value))] |
||||
fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> { |
||||
self.insert_batch(&mut Some((key.to_owned(), value.to_owned())).into_iter())?; |
||||
let watchers = self.watchers.read().unwrap(); |
||||
let mut triggered = Vec::new(); |
||||
|
||||
for length in 0..=key.len() { |
||||
if watchers.contains_key(&key[..length]) { |
||||
triggered.push(&key[..length]); |
||||
} |
||||
} |
||||
|
||||
drop(watchers); |
||||
|
||||
if !triggered.is_empty() { |
||||
let mut watchers = self.watchers.write().unwrap(); |
||||
for prefix in triggered { |
||||
if let Some(txs) = watchers.remove(prefix) { |
||||
for tx in txs { |
||||
let _ = tx.send(()); |
||||
} |
||||
} |
||||
} |
||||
} |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, iter))] |
||||
fn insert_batch<'a>(&self, iter: &mut dyn Iterator<Item = (Vec<u8>, Vec<u8>)>) -> Result<()> { |
||||
let mut tx = self.begin()?; |
||||
for (key, value) in iter { |
||||
tx.put::<ByteVec, ByteVec>( |
||||
&self.name, |
||||
ByteVec::from(key.clone()), |
||||
ByteVec::from(value), |
||||
)?; |
||||
} |
||||
tx.prepare()?.commit()?; |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, iter))] |
||||
fn increment_batch<'a>(&self, iter: &mut dyn Iterator<Item = Vec<u8>>) -> Result<()> { |
||||
let mut tx = self.begin()?; |
||||
for key in iter { |
||||
let old = tx |
||||
.get::<ByteVec, ByteVec>(&self.name, &ByteVec::from(key.clone()))? |
||||
.next() |
||||
.map(|v| (*v).to_owned()); |
||||
let new = crate::utils::increment(old.as_deref()).unwrap(); |
||||
tx.put::<ByteVec, ByteVec>(&self.name, ByteVec::from(key), ByteVec::from(new))?; |
||||
} |
||||
tx.prepare()?.commit()?; |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, key))] |
||||
fn remove(&self, key: &[u8]) -> Result<()> { |
||||
let mut tx = self.begin()?; |
||||
tx.remove::<ByteVec, ByteVec>(&self.name, ByteVec::from(key), None)?; |
||||
tx.prepare()?.commit()?; |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self))] |
||||
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> { |
||||
let iter = self.persy.range::<ByteVec, ByteVec, _>(&self.name, ..); |
||||
match iter { |
||||
Ok(iter) => Box::new(iter.filter_map(|(k, v)| { |
||||
v.into_iter() |
||||
.map(|val| ((*k).to_owned().into(), (*val).to_owned().into())) |
||||
.next() |
||||
})), |
||||
Err(e) => { |
||||
warn!("error iterating {:?}", e); |
||||
Box::new(std::iter::empty()) |
||||
} |
||||
} |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, from, backwards))] |
||||
fn iter_from<'a>( |
||||
&'a self, |
||||
from: &[u8], |
||||
backwards: bool, |
||||
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> { |
||||
let range = if backwards { |
||||
self.persy |
||||
.range::<ByteVec, ByteVec, _>(&self.name, ..=ByteVec::from(from)) |
||||
} else { |
||||
self.persy |
||||
.range::<ByteVec, ByteVec, _>(&self.name, ByteVec::from(from)..) |
||||
}; |
||||
match range { |
||||
Ok(iter) => { |
||||
let map = iter.filter_map(|(k, v)| { |
||||
v.into_iter() |
||||
.map(|val| ((*k).to_owned().into(), (*val).to_owned().into())) |
||||
.next() |
||||
}); |
||||
if backwards { |
||||
Box::new(map.rev()) |
||||
} else { |
||||
Box::new(map) |
||||
} |
||||
} |
||||
Err(e) => { |
||||
warn!("error iterating with prefix {:?}", e); |
||||
Box::new(std::iter::empty()) |
||||
} |
||||
} |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, key))] |
||||
fn increment(&self, key: &[u8]) -> Result<Vec<u8>> { |
||||
self.increment_batch(&mut Some(key.to_owned()).into_iter())?; |
||||
Ok(self.get(key)?.unwrap()) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, prefix))] |
||||
fn scan_prefix<'a>( |
||||
&'a self, |
||||
prefix: Vec<u8>, |
||||
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> { |
||||
let range_prefix = ByteVec::from(prefix.clone()); |
||||
let range = self |
||||
.persy |
||||
.range::<ByteVec, ByteVec, _>(&self.name, range_prefix..); |
||||
|
||||
match range { |
||||
Ok(iter) => { |
||||
let owned_prefix = prefix.clone(); |
||||
Box::new( |
||||
iter.take_while(move |(k, _)| (*k).starts_with(&owned_prefix)) |
||||
.filter_map(|(k, v)| { |
||||
v.into_iter() |
||||
.map(|val| ((*k).to_owned().into(), (*val).to_owned().into())) |
||||
.next() |
||||
}), |
||||
) |
||||
} |
||||
Err(e) => { |
||||
warn!("error scanning prefix {:?}", e); |
||||
Box::new(std::iter::empty()) |
||||
} |
||||
} |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, prefix))] |
||||
fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> { |
||||
let (tx, rx) = tokio::sync::oneshot::channel(); |
||||
|
||||
self.watchers |
||||
.write() |
||||
.unwrap() |
||||
.entry(prefix.to_vec()) |
||||
.or_default() |
||||
.push(tx); |
||||
|
||||
Box::pin(async move { |
||||
// Tx is never destroyed
|
||||
rx.await.unwrap(); |
||||
}) |
||||
} |
||||
} |
||||
Loading…
Reference in new issue