use std::{future::Future, pin::Pin, sync::Arc}; use conduit::{utils, Result}; use rocksdb::{BoundColumnFamily, Direction, IteratorMode, ReadOptions, WriteBatchWithTransaction, WriteOptions}; use super::{or_else, result, watchers::Watchers, Engine}; pub struct Map { db: Arc, name: String, watchers: Watchers, write_options: WriteOptions, read_options: ReadOptions, } type Key = Vec; type Val = Vec; type KeyVal = (Key, Val); impl Map { pub(crate) fn open(db: &Arc, name: &str) -> Result> { db.open_cf(name)?; Ok(Arc::new(Self { db: db.clone(), name: name.to_owned(), watchers: Watchers::default(), write_options: write_options_default(), read_options: read_options_default(), })) } pub fn get(&self, key: &[u8]) -> Result>> { let read_options = &self.read_options; let res = self.db.db.get_cf_opt(&self.cf(), key, read_options); result(res) } pub fn multi_get(&self, keys: &[&[u8]]) -> Result>>> { // Optimization can be `true` if key vector is pre-sorted **by the column // comparator**. const SORTED: bool = false; let mut ret: Vec>> = Vec::with_capacity(keys.len()); let read_options = &self.read_options; for res in self .db .db .batched_multi_get_cf_opt(&self.cf(), keys, SORTED, read_options) { match res { Ok(Some(res)) => ret.push(Some((*res).to_vec())), Ok(None) => ret.push(None), Err(e) => return or_else(e), } } Ok(ret) } pub fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> { let write_options = &self.write_options; self.db .db .put_cf_opt(&self.cf(), key, value, write_options) .or_else(or_else)?; if !self.db.corked() { self.db.flush()?; } self.watchers.wake(key); Ok(()) } pub fn insert_batch(&self, iter: &mut dyn Iterator) -> Result<()> { let mut batch = WriteBatchWithTransaction::::default(); for (key, value) in iter { batch.put_cf(&self.cf(), key, value); } let write_options = &self.write_options; let res = self.db.db.write_opt(batch, write_options); if !self.db.corked() { self.db.flush()?; } result(res) } pub fn remove(&self, key: &[u8]) -> Result<()> { let write_options = &self.write_options; let res = self.db.db.delete_cf_opt(&self.cf(), key, write_options); if !self.db.corked() { self.db.flush()?; } result(res) } pub fn remove_batch(&self, iter: &mut dyn Iterator) -> Result<()> { let mut batch = WriteBatchWithTransaction::::default(); for key in iter { batch.delete_cf(&self.cf(), key); } let write_options = &self.write_options; let res = self.db.db.write_opt(batch, write_options); if !self.db.corked() { self.db.flush()?; } result(res) } pub fn iter<'a>(&'a self) -> Box + 'a> { let read_options = read_options_default(); let it = self .db .db .iterator_cf_opt(&self.cf(), read_options, IteratorMode::Start) .map(Result::unwrap) .map(|(k, v)| (Vec::from(k), Vec::from(v))); Box::new(it) } pub fn iter_from<'a>(&'a self, from: &[u8], backwards: bool) -> Box + 'a> { let direction = if backwards { Direction::Reverse } else { Direction::Forward }; let mode = IteratorMode::From(from, direction); let read_options = read_options_default(); let it = self .db .db .iterator_cf_opt(&self.cf(), read_options, mode) .map(Result::unwrap) .map(|(k, v)| (Vec::from(k), Vec::from(v))); Box::new(it) } pub fn scan_prefix<'a>(&'a self, prefix: Vec) -> Box + 'a> { let mode = IteratorMode::From(&prefix, Direction::Forward); let read_options = read_options_default(); let it = self .db .db .iterator_cf_opt(&self.cf(), read_options, mode) .map(Result::unwrap) .map(|(k, v)| (Vec::from(k), Vec::from(v))) .take_while(move |(k, _)| k.starts_with(&prefix)); Box::new(it) } pub fn increment(&self, key: &[u8]) -> Result> { let read_options = &self.read_options; let old = self .db .db .get_cf_opt(&self.cf(), key, read_options) .or_else(or_else)?; let new = utils::increment(old.as_deref()); let write_options = &self.write_options; self.db .db .put_cf_opt(&self.cf(), key, new, write_options) .or_else(or_else)?; if !self.db.corked() { self.db.flush()?; } Ok(new.to_vec()) } pub fn increment_batch(&self, iter: &mut dyn Iterator) -> Result<()> { let mut batch = WriteBatchWithTransaction::::default(); let read_options = &self.read_options; for key in iter { let old = self .db .db .get_cf_opt(&self.cf(), &key, read_options) .or_else(or_else)?; let new = utils::increment(old.as_deref()); batch.put_cf(&self.cf(), key, new); } let write_options = &self.write_options; self.db .db .write_opt(batch, write_options) .or_else(or_else)?; if !self.db.corked() { self.db.flush()?; } Ok(()) } pub fn watch_prefix<'a>(&'a self, prefix: &[u8]) -> Pin + Send + 'a>> { self.watchers.watch(prefix) } fn cf(&self) -> Arc> { self.db.cf(&self.name) } } impl<'a> IntoIterator for &'a Map { type IntoIter = Box + 'a>; type Item = KeyVal; fn into_iter(self) -> Self::IntoIter { self.iter() } } #[inline] fn read_options_default() -> ReadOptions { let mut read_options = ReadOptions::default(); read_options.set_total_order_seek(true); read_options } #[inline] fn write_options_default() -> WriteOptions { WriteOptions::default() }