mirror of https://gitlab.com/famedly/conduit
improvement: allow rocksdb again
Timo Kösters
3 years ago
parent
b25354c747
commit
1d647a1a9a
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,183 @@
|
||||
use super::super::Config;
|
||||
use crate::{utils, Result};
|
||||
|
||||
use std::{future::Future, pin::Pin, sync::Arc};
|
||||
|
||||
use super::{DatabaseEngine, Tree};
|
||||
|
||||
use std::{collections::HashMap, sync::RwLock};
|
||||
|
||||
pub struct Engine {
|
||||
rocks: rocksdb::DBWithThreadMode<rocksdb::MultiThreaded>,
|
||||
old_cfs: Vec<String>,
|
||||
}
|
||||
|
||||
pub struct RocksDbEngineTree<'a> {
|
||||
db: Arc<Engine>,
|
||||
name: &'a str,
|
||||
watchers: Watchers,
|
||||
}
|
||||
|
||||
impl DatabaseEngine for Engine {
|
||||
fn open(config: &Config) -> Result<Arc<Self>> {
|
||||
let mut db_opts = rocksdb::Options::default();
|
||||
db_opts.create_if_missing(true);
|
||||
db_opts.set_max_open_files(16);
|
||||
db_opts.set_compaction_style(rocksdb::DBCompactionStyle::Level);
|
||||
db_opts.set_compression_type(rocksdb::DBCompressionType::Snappy);
|
||||
db_opts.set_target_file_size_base(256 << 20);
|
||||
db_opts.set_write_buffer_size(256 << 20);
|
||||
|
||||
let mut block_based_options = rocksdb::BlockBasedOptions::default();
|
||||
block_based_options.set_block_size(512 << 10);
|
||||
db_opts.set_block_based_table_factory(&block_based_options);
|
||||
|
||||
let cfs = rocksdb::DBWithThreadMode::<rocksdb::MultiThreaded>::list_cf(
|
||||
&db_opts,
|
||||
&config.database_path,
|
||||
)
|
||||
.unwrap_or_default();
|
||||
|
||||
let db = rocksdb::DBWithThreadMode::<rocksdb::MultiThreaded>::open_cf_descriptors(
|
||||
&db_opts,
|
||||
&config.database_path,
|
||||
cfs.iter().map(|name| {
|
||||
let mut options = rocksdb::Options::default();
|
||||
let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(1);
|
||||
options.set_prefix_extractor(prefix_extractor);
|
||||
options.set_merge_operator_associative("increment", utils::increment_rocksdb);
|
||||
|
||||
rocksdb::ColumnFamilyDescriptor::new(name, options)
|
||||
}),
|
||||
)?;
|
||||
|
||||
Ok(Arc::new(Engine {
|
||||
rocks: db,
|
||||
old_cfs: cfs,
|
||||
}))
|
||||
}
|
||||
|
||||
fn open_tree(self: &Arc<Self>, name: &'static str) -> Result<Arc<dyn Tree>> {
|
||||
if !self.old_cfs.contains(&name.to_owned()) {
|
||||
// Create if it didn't exist
|
||||
let mut options = rocksdb::Options::default();
|
||||
let prefix_extractor = rocksdb::SliceTransform::create_fixed_prefix(1);
|
||||
options.set_prefix_extractor(prefix_extractor);
|
||||
options.set_merge_operator_associative("increment", utils::increment_rocksdb);
|
||||
|
||||
let _ = self.rocks.create_cf(name, &options);
|
||||
println!("created cf");
|
||||
}
|
||||
|
||||
Ok(Arc::new(RocksDbEngineTree {
|
||||
name,
|
||||
db: Arc::clone(self),
|
||||
watchers: Watchers::default(),
|
||||
}))
|
||||
}
|
||||
|
||||
fn flush(self: &Arc<Self>) -> Result<()> {
|
||||
// TODO?
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl RocksDbEngineTree<'_> {
|
||||
fn cf(&self) -> rocksdb::BoundColumnFamily<'_> {
|
||||
self.db.rocks.cf_handle(self.name).unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl Tree for RocksDbEngineTree<'_> {
|
||||
fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
|
||||
Ok(self.db.rocks.get_cf(self.cf(), key)?)
|
||||
}
|
||||
|
||||
fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> {
|
||||
self.db.rocks.put_cf(self.cf(), key, value)?;
|
||||
self.watchers.wake(key);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn insert_batch<'a>(&self, iter: &mut dyn Iterator<Item = (Vec<u8>, Vec<u8>)>) -> Result<()> {
|
||||
for (key, value) in iter {
|
||||
self.db.rocks.put_cf(self.cf(), key, value)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn remove(&self, key: &[u8]) -> Result<()> {
|
||||
Ok(self.db.rocks.delete_cf(self.cf(), key)?)
|
||||
}
|
||||
|
||||
fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
|
||||
Box::new(
|
||||
self.db
|
||||
.rocks
|
||||
.iterator_cf(self.cf(), rocksdb::IteratorMode::Start)
|
||||
.map(|(k, v)| (Vec::from(k), Vec::from(v))),
|
||||
)
|
||||
}
|
||||
|
||||
fn iter_from<'a>(
|
||||
&'a self,
|
||||
from: &[u8],
|
||||
backwards: bool,
|
||||
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
|
||||
Box::new(
|
||||
self.db
|
||||
.rocks
|
||||
.iterator_cf(
|
||||
self.cf(),
|
||||
rocksdb::IteratorMode::From(
|
||||
from,
|
||||
if backwards {
|
||||
rocksdb::Direction::Reverse
|
||||
} else {
|
||||
rocksdb::Direction::Forward
|
||||
},
|
||||
),
|
||||
)
|
||||
.map(|(k, v)| (Vec::from(k), Vec::from(v))),
|
||||
)
|
||||
}
|
||||
|
||||
fn increment(&self, key: &[u8]) -> Result<Vec<u8>> {
|
||||
// TODO: make atomic
|
||||
let old = self.db.rocks.get_cf(self.cf(), &key)?;
|
||||
let new = utils::increment(old.as_deref()).unwrap();
|
||||
self.db.rocks.put_cf(self.cf(), key, &new)?;
|
||||
Ok(new)
|
||||
}
|
||||
|
||||
fn increment_batch<'a>(&self, iter: &mut dyn Iterator<Item = Vec<u8>>) -> Result<()> {
|
||||
for key in iter {
|
||||
let old = self.db.rocks.get_cf(self.cf(), &key)?;
|
||||
let new = utils::increment(old.as_deref()).unwrap();
|
||||
self.db.rocks.put_cf(self.cf(), key, new)?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn scan_prefix<'a>(
|
||||
&'a self,
|
||||
prefix: Vec<u8>,
|
||||
) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
|
||||
Box::new(
|
||||
self.db
|
||||
.rocks
|
||||
.iterator_cf(
|
||||
self.cf(),
|
||||
rocksdb::IteratorMode::From(&prefix, rocksdb::Direction::Forward),
|
||||
)
|
||||
.map(|(k, v)| (Vec::from(k), Vec::from(v)))
|
||||
.take_while(move |(k, _)| k.starts_with(&prefix)),
|
||||
)
|
||||
}
|
||||
|
||||
fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> {
|
||||
self.watchers.watch(prefix)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in New Issue