improvement: allow rocksdb again

Cargo.toml

@@ -78,6 +78,7 @@ crossbeam = { version = "0.8.1", optional = true }
 num_cpus = "1.13.0"
 threadpool = "1.8.1"
 heed = { git = "https://github.com/timokoesters/heed.git", rev = "f6f825da7fb2c758867e05ad973ef800a6fe1d5d", optional = true }
+rocksdb = { version = "0.16.0", features = ["multi-threaded-cf"], optional = true }
 thread_local = "1.1.3"
 # used for TURN server authentication
 hmac = "0.11.0"
@@ -87,7 +88,8 @@ sha-1 = "0.9.8"
 default = ["conduit_bin", "backend_sqlite"]
 backend_sled = ["sled"]
 backend_sqlite = ["sqlite"]
-backend_heed = ["heed", "crossbeam", "parking_lot"]
+backend_heed = ["heed", "crossbeam"]
+backend_rocksdb = ["rocksdb"]
 sqlite = ["rusqlite", "parking_lot", "crossbeam", "tokio/signal"]
 conduit_bin = [] # TODO: add rocket to this when it is optional
 

src/database.rs

@@ -154,6 +154,9 @@ pub type Engine = abstraction::sqlite::Engine;
 #[cfg(feature = "heed")]
 pub type Engine = abstraction::heed::Engine;
 
+#[cfg(feature = "rocksdb")]
+pub type Engine = abstraction::rocksdb::Engine;
+
 pub struct Database {
     _db: Arc<Engine>,
     pub globals: globals::Globals,
@@ -314,10 +317,10 @@ impl Database {
                         .expect("pdu cache capacity fits into usize"),
                 )),
                 auth_chain_cache: Mutex::new(LruCache::new(1_000_000)),
-                shorteventid_cache: Mutex::new(LruCache::new(1_000_000)),
-                eventidshort_cache: Mutex::new(LruCache::new(1_000_000)),
-                shortstatekey_cache: Mutex::new(LruCache::new(1_000_000)),
-                statekeyshort_cache: Mutex::new(LruCache::new(1_000_000)),
+                shorteventid_cache: Mutex::new(LruCache::new(100_000_000)),
+                eventidshort_cache: Mutex::new(LruCache::new(100_000_000)),
+                shortstatekey_cache: Mutex::new(LruCache::new(100_000_000)),
+                statekeyshort_cache: Mutex::new(LruCache::new(100_000_000)),
                 our_real_users_cache: RwLock::new(HashMap::new()),
                 appservice_in_room_cache: RwLock::new(HashMap::new()),
                 stateinfo_cache: Mutex::new(LruCache::new(1000)),

src/database/abstraction.rs

@@ -12,7 +12,10 @@ pub mod sqlite;
 #[cfg(feature = "heed")]
 pub mod heed;
 
-#[cfg(any(feature = "sqlite", feature = "heed"))]
+#[cfg(feature = "rocksdb")]
+pub mod rocksdb;
+
+#[cfg(any(feature = "sqlite", feature = "rocksdb", feature = "heed"))]
 pub mod watchers;
 
 pub trait DatabaseEngine: Sized {

src/database/abstraction/rocksdb.rs

@@ -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)
+    }
+}

src/database/abstraction/sqlite.rs

@@ -132,7 +132,7 @@ type TupleOfBytes = (Vec<u8>, Vec<u8>);
 impl SqliteTable {
     #[tracing::instrument(skip(self, guard, key))]
     fn get_with_guard(&self, guard: &Connection, key: &[u8]) -> Result<Option<Vec<u8>>> {
-        //dbg!(&self.name);
+        dbg!(&self.name);
         Ok(guard
             .prepare(format!("SELECT value FROM {} WHERE key = ?", self.name).as_str())?
             .query_row([key], |row| row.get(0))
@@ -141,7 +141,7 @@ impl SqliteTable {
 
     #[tracing::instrument(skip(self, guard, key, value))]
     fn insert_with_guard(&self, guard: &Connection, key: &[u8], value: &[u8]) -> Result<()> {
-        //dbg!(&self.name);
+        dbg!(&self.name);
         guard.execute(
             format!(
                 "INSERT OR REPLACE INTO {} (key, value) VALUES (?, ?)",
@@ -168,14 +168,14 @@ impl SqliteTable {
 
         let statement_ref = NonAliasingBox(statement);
 
-        //let name = self.name.clone();
+        let name = self.name.clone();
 
         let iterator = Box::new(
             statement
                 .query_map([], |row| Ok((row.get_unwrap(0), row.get_unwrap(1))))
                 .unwrap()
                 .map(move |r| {
-                    //dbg!(&name);
+                    dbg!(&name);
                     r.unwrap()
                 }),
         );
@@ -263,7 +263,7 @@ impl Tree for SqliteTable {
         let guard = self.engine.read_lock_iterator();
         let from = from.to_vec(); // TODO change interface?
 
-        //let name = self.name.clone();
+        let name = self.name.clone();
 
         if backwards {
             let statement = Box::leak(Box::new(
@@ -282,7 +282,7 @@ impl Tree for SqliteTable {
                     .query_map([from], |row| Ok((row.get_unwrap(0), row.get_unwrap(1))))
                     .unwrap()
                     .map(move |r| {
-                        //dbg!(&name);
+                        dbg!(&name);
                         r.unwrap()
                     }),
             );
@@ -307,7 +307,7 @@ impl Tree for SqliteTable {
                     .query_map([from], |row| Ok((row.get_unwrap(0), row.get_unwrap(1))))
                     .unwrap()
                     .map(move |r| {
-                        //dbg!(&name);
+                        dbg!(&name);
                         r.unwrap()
                     }),
             );

src/error.rs

@@ -39,6 +39,12 @@ pub enum Error {
     #[cfg(feature = "heed")]
     #[error("There was a problem with the connection to the heed database: {error}")]
     HeedError { error: String },
+    #[cfg(feature = "rocksdb")]
+    #[error("There was a problem with the connection to the rocksdb database: {source}")]
+    RocksDbError {
+        #[from]
+        source: rocksdb::Error,
+    },
     #[error("Could not generate an image.")]
     ImageError {
         #[from]

src/utils.rs

@@ -29,6 +29,17 @@ pub fn increment(old: Option<&[u8]>) -> Option<Vec<u8>> {
     Some(number.to_be_bytes().to_vec())
 }
 
+#[cfg(feature = "rocksdb")]
+pub fn increment_rocksdb(
+    _new_key: &[u8],
+    old: Option<&[u8]>,
+    _operands: &mut rocksdb::MergeOperands,
+) -> Option<Vec<u8>> {
+    dbg!(_new_key);
+    dbg!(old);
+    increment(old)
+}
+
 pub fn generate_keypair() -> Vec<u8> {
     let mut value = random_string(8).as_bytes().to_vec();
     value.push(0xff);