Timo Kösters
4 years ago
9 changed files with 456 additions and 97 deletions
@ -0,0 +1,241 @@
@@ -0,0 +1,241 @@
|
||||
use super::super::Config; |
||||
use crossbeam::channel::{bounded, Sender as ChannelSender}; |
||||
use threadpool::ThreadPool; |
||||
|
||||
use crate::{Error, Result}; |
||||
use std::{ |
||||
collections::HashMap, |
||||
future::Future, |
||||
pin::Pin, |
||||
sync::{Arc, Mutex, RwLock}, |
||||
}; |
||||
use tokio::sync::oneshot::Sender; |
||||
|
||||
use super::{DatabaseEngine, Tree}; |
||||
|
||||
type TupleOfBytes = (Vec<u8>, Vec<u8>); |
||||
|
||||
pub struct Engine { |
||||
env: heed::Env, |
||||
iter_pool: Mutex<ThreadPool>, |
||||
} |
||||
|
||||
pub struct EngineTree { |
||||
engine: Arc<Engine>, |
||||
tree: Arc<heed::UntypedDatabase>, |
||||
watchers: RwLock<HashMap<Vec<u8>, Vec<Sender<()>>>>, |
||||
} |
||||
|
||||
fn convert_error(error: heed::Error) -> Error { |
||||
panic!(error.to_string()); |
||||
Error::HeedError { |
||||
error: error.to_string(), |
||||
} |
||||
} |
||||
|
||||
impl DatabaseEngine for Engine { |
||||
fn open(config: &Config) -> Result<Arc<Self>> { |
||||
let mut env_builder = heed::EnvOpenOptions::new(); |
||||
env_builder.map_size(1024 * 1024 * 1024 * 1024); // 1 Terabyte
|
||||
env_builder.max_readers(126); |
||||
env_builder.max_dbs(128); |
||||
unsafe { |
||||
env_builder.flag(heed::flags::Flags::MdbNoSync); |
||||
env_builder.flag(heed::flags::Flags::MdbNoMetaSync); |
||||
} |
||||
|
||||
Ok(Arc::new(Engine { |
||||
env: env_builder |
||||
.open(&config.database_path) |
||||
.map_err(convert_error)?, |
||||
iter_pool: Mutex::new(ThreadPool::new(10)), |
||||
})) |
||||
} |
||||
|
||||
fn open_tree(self: &Arc<Self>, name: &'static str) -> Result<Arc<dyn Tree>> { |
||||
// Creates the db if it doesn't exist already
|
||||
Ok(Arc::new(EngineTree { |
||||
engine: Arc::clone(self), |
||||
tree: Arc::new( |
||||
self.env |
||||
.create_database(Some(name)) |
||||
.map_err(convert_error)?, |
||||
), |
||||
watchers: RwLock::new(HashMap::new()), |
||||
})) |
||||
} |
||||
|
||||
fn flush(self: &Arc<Self>) -> Result<()> { |
||||
self.env.force_sync().map_err(convert_error)?; |
||||
Ok(()) |
||||
} |
||||
} |
||||
|
||||
impl EngineTree { |
||||
#[tracing::instrument(skip(self, tree, from, backwards))] |
||||
fn iter_from_thread( |
||||
&self, |
||||
tree: Arc<heed::UntypedDatabase>, |
||||
from: Vec<u8>, |
||||
backwards: bool, |
||||
) -> Box<dyn Iterator<Item = TupleOfBytes> + Send + Sync> { |
||||
let (s, r) = bounded::<TupleOfBytes>(5); |
||||
let engine = Arc::clone(&self.engine); |
||||
|
||||
let lock = self.engine.iter_pool.lock().unwrap(); |
||||
if lock.active_count() < lock.max_count() { |
||||
lock.execute(move || { |
||||
iter_from_thread_work(tree, &engine.env.read_txn().unwrap(), from, backwards, &s); |
||||
}); |
||||
} else { |
||||
std::thread::spawn(move || { |
||||
iter_from_thread_work(tree, &engine.env.read_txn().unwrap(), from, backwards, &s); |
||||
}); |
||||
} |
||||
|
||||
Box::new(r.into_iter()) |
||||
} |
||||
} |
||||
|
||||
#[tracing::instrument(skip(tree, txn, from, backwards))] |
||||
fn iter_from_thread_work( |
||||
tree: Arc<heed::UntypedDatabase>, |
||||
txn: &heed::RoTxn<'_>, |
||||
from: Vec<u8>, |
||||
backwards: bool, |
||||
s: &ChannelSender<(Vec<u8>, Vec<u8>)>, |
||||
) { |
||||
if backwards { |
||||
for (k, v) in tree.rev_range(txn, ..=&*from).unwrap().map(|r| r.unwrap()) { |
||||
if s.send((k.to_vec(), v.to_vec())).is_err() { |
||||
return; |
||||
} |
||||
} |
||||
} else { |
||||
if from.is_empty() { |
||||
for (k, v) in tree.iter(txn).unwrap().map(|r| r.unwrap()) { |
||||
if s.send((k.to_vec(), v.to_vec())).is_err() { |
||||
return; |
||||
} |
||||
} |
||||
} else { |
||||
for (k, v) in tree.range(txn, &*from..).unwrap().map(|r| r.unwrap()) { |
||||
if s.send((k.to_vec(), v.to_vec())).is_err() { |
||||
return; |
||||
} |
||||
} |
||||
} |
||||
} |
||||
} |
||||
|
||||
impl Tree for EngineTree { |
||||
#[tracing::instrument(skip(self, key))] |
||||
fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> { |
||||
let txn = self.engine.env.read_txn().map_err(convert_error)?; |
||||
Ok(self |
||||
.tree |
||||
.get(&txn, &key) |
||||
.map_err(convert_error)? |
||||
.map(|s| s.to_vec())) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, key, value))] |
||||
fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> { |
||||
let mut txn = self.engine.env.write_txn().map_err(convert_error)?; |
||||
self.tree |
||||
.put(&mut txn, &key, &value) |
||||
.map_err(convert_error)?; |
||||
txn.commit().map_err(convert_error)?; |
||||
|
||||
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, key))] |
||||
fn remove(&self, key: &[u8]) -> Result<()> { |
||||
let mut txn = self.engine.env.write_txn().map_err(convert_error)?; |
||||
self.tree.delete(&mut txn, &key).map_err(convert_error)?; |
||||
txn.commit().map_err(convert_error)?; |
||||
Ok(()) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self))] |
||||
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + Send + 'a> { |
||||
self.iter_from(&[], false) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, from, backwards))] |
||||
fn iter_from( |
||||
&self, |
||||
from: &[u8], |
||||
backwards: bool, |
||||
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + Send> { |
||||
self.iter_from_thread(Arc::clone(&self.tree), from.to_vec(), backwards) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, key))] |
||||
fn increment(&self, key: &[u8]) -> Result<Vec<u8>> { |
||||
let mut txn = self.engine.env.write_txn().map_err(convert_error)?; |
||||
|
||||
let old = self.tree.get(&txn, &key).map_err(convert_error)?; |
||||
let new = |
||||
crate::utils::increment(old.as_deref()).expect("utils::increment always returns Some"); |
||||
|
||||
self.tree |
||||
.put(&mut txn, &key, &&*new) |
||||
.map_err(convert_error)?; |
||||
|
||||
txn.commit().map_err(convert_error)?; |
||||
|
||||
Ok(new) |
||||
} |
||||
|
||||
#[tracing::instrument(skip(self, prefix))] |
||||
fn scan_prefix<'a>( |
||||
&'a self, |
||||
prefix: Vec<u8>, |
||||
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + Send + 'a> { |
||||
Box::new( |
||||
self.iter_from(&prefix, false) |
||||
.take_while(move |(key, _)| key.starts_with(&prefix)), |
||||
) |
||||
} |
||||
|
||||
#[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