refactor rocksdb opts; split kvtree

Signed-off-by: Jason Volk <jason@zemos.net>

src/database/rocksdb/kvtree.rs

@@ -0,0 +1,197 @@
+use std::{future::Future, pin::Pin, sync::Arc};
+
+use rust_rocksdb::WriteBatchWithTransaction;
+
+use super::{watchers::Watchers, Engine, KeyValueDatabaseEngine, KvTree};
+use crate::{utils, Result};
+
+pub(super) struct RocksDbEngineTree<'a> {
+	pub db: Arc<Engine>,
+	pub name: &'a str,
+	pub watchers: Watchers,
+}
+
+impl RocksDbEngineTree<'_> {
+	fn cf(&self) -> Arc<rust_rocksdb::BoundColumnFamily<'_>> { self.db.rocks.cf_handle(self.name).unwrap() }
+}
+
+impl KvTree for RocksDbEngineTree<'_> {
+	fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+
+		Ok(self.db.rocks.get_cf_opt(&self.cf(), key, &readoptions)?)
+	}
+
+	fn multi_get(
+		&self, iter: Vec<(&Arc<rust_rocksdb::BoundColumnFamily<'_>>, Vec<u8>)>,
+	) -> Vec<Result<Option<Vec<u8>>, rust_rocksdb::Error>> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+
+		self.db.rocks.multi_get_cf_opt(iter, &readoptions)
+	}
+
+	fn insert(&self, key: &[u8], value: &[u8]) -> Result<()> {
+		let writeoptions = rust_rocksdb::WriteOptions::default();
+
+		self.db
+			.rocks
+			.put_cf_opt(&self.cf(), key, value, &writeoptions)?;
+
+		if !self.db.corked() {
+			self.db.flush()?;
+		}
+
+		self.watchers.wake(key);
+
+		Ok(())
+	}
+
+	fn insert_batch(&self, iter: &mut dyn Iterator<Item = (Vec<u8>, Vec<u8>)>) -> Result<()> {
+		let writeoptions = rust_rocksdb::WriteOptions::default();
+
+		let mut batch = WriteBatchWithTransaction::<false>::default();
+
+		for (key, value) in iter {
+			batch.put_cf(&self.cf(), key, value);
+		}
+
+		let result = self.db.rocks.write_opt(batch, &writeoptions);
+
+		if !self.db.corked() {
+			self.db.flush()?;
+		}
+
+		Ok(result?)
+	}
+
+	fn remove(&self, key: &[u8]) -> Result<()> {
+		let writeoptions = rust_rocksdb::WriteOptions::default();
+
+		let result = self.db.rocks.delete_cf_opt(&self.cf(), key, &writeoptions);
+
+		if !self.db.corked() {
+			self.db.flush()?;
+		}
+
+		Ok(result?)
+	}
+
+	fn remove_batch(&self, iter: &mut dyn Iterator<Item = Vec<u8>>) -> Result<()> {
+		let writeoptions = rust_rocksdb::WriteOptions::default();
+
+		let mut batch = WriteBatchWithTransaction::<false>::default();
+
+		for key in iter {
+			batch.delete_cf(&self.cf(), key);
+		}
+
+		let result = self.db.rocks.write_opt(batch, &writeoptions);
+
+		if !self.db.corked() {
+			self.db.flush()?;
+		}
+
+		Ok(result?)
+	}
+
+	fn iter<'a>(&'a self) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+
+		Box::new(
+			self.db
+				.rocks
+				.iterator_cf_opt(&self.cf(), readoptions, rust_rocksdb::IteratorMode::Start)
+				.map(Result::unwrap)
+				.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> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+
+		Box::new(
+			self.db
+				.rocks
+				.iterator_cf_opt(
+					&self.cf(),
+					readoptions,
+					rust_rocksdb::IteratorMode::From(
+						from,
+						if backwards {
+							rust_rocksdb::Direction::Reverse
+						} else {
+							rust_rocksdb::Direction::Forward
+						},
+					),
+				)
+				.map(Result::unwrap)
+				.map(|(k, v)| (Vec::from(k), Vec::from(v))),
+		)
+	}
+
+	fn increment(&self, key: &[u8]) -> Result<Vec<u8>> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+		let writeoptions = rust_rocksdb::WriteOptions::default();
+
+		let old = self.db.rocks.get_cf_opt(&self.cf(), key, &readoptions)?;
+		let new = utils::increment(old.as_deref());
+		self.db
+			.rocks
+			.put_cf_opt(&self.cf(), key, &new, &writeoptions)?;
+
+		if !self.db.corked() {
+			self.db.flush()?;
+		}
+
+		Ok(new)
+	}
+
+	fn increment_batch(&self, iter: &mut dyn Iterator<Item = Vec<u8>>) -> Result<()> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+		let writeoptions = rust_rocksdb::WriteOptions::default();
+
+		let mut batch = WriteBatchWithTransaction::<false>::default();
+
+		for key in iter {
+			let old = self.db.rocks.get_cf_opt(&self.cf(), &key, &readoptions)?;
+			let new = utils::increment(old.as_deref());
+			batch.put_cf(&self.cf(), key, new);
+		}
+
+		self.db.rocks.write_opt(batch, &writeoptions)?;
+
+		if !self.db.corked() {
+			self.db.flush()?;
+		}
+
+		Ok(())
+	}
+
+	fn scan_prefix<'a>(&'a self, prefix: Vec<u8>) -> Box<dyn Iterator<Item = (Vec<u8>, Vec<u8>)> + 'a> {
+		let mut readoptions = rust_rocksdb::ReadOptions::default();
+		readoptions.set_total_order_seek(true);
+
+		Box::new(
+			self.db
+				.rocks
+				.iterator_cf_opt(
+					&self.cf(),
+					readoptions,
+					rust_rocksdb::IteratorMode::From(&prefix, rust_rocksdb::Direction::Forward),
+				)
+				.map(Result::unwrap)
+				.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)
+	}
+}