From 32435d8a4c3a5ac84f8f6e3e46da292e5e907849 Mon Sep 17 00:00:00 2001
From: Luke Parker <lukeparker5132@gmail.com>
Date: Tue, 25 Jul 2023 21:39:29 -0400
Subject: [PATCH] Consolidate RockDB code

Moves explicitly to zstd. RockDB recommends zstd, or at least lz4 over snappy,
and this minimizes which dependencies we pull in.
---
 common/db/Cargo.toml      | 2 +-
 common/db/src/lib.rs      | 2 ++
 common/db/src/rocks.rs    | 9 ++++++++-
 coordinator/src/main.rs   | 7 +------
 message-queue/src/main.rs | 9 ++-------
 processor/src/main.rs     | 8 +-------
 6 files changed, 15 insertions(+), 22 deletions(-)

diff --git a/common/db/Cargo.toml b/common/db/Cargo.toml
index 9606bda8..83c536bb 100644
--- a/common/db/Cargo.toml
+++ b/common/db/Cargo.toml
@@ -13,7 +13,7 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs"]
 
 [dependencies]
-rocksdb = { version = "0.21", optional = true }
+rocksdb = { version = "0.21", default-features = false, features = ["zstd"], optional = true }
 
 [features]
 rocksdb = ["dep:rocksdb"]
diff --git a/common/db/src/lib.rs b/common/db/src/lib.rs
index 5baf35f1..0bcf9810 100644
--- a/common/db/src/lib.rs
+++ b/common/db/src/lib.rs
@@ -3,6 +3,8 @@ pub use mem::*;
 
 #[cfg(feature = "rocksdb")]
 mod rocks;
+#[cfg(feature = "rocksdb")]
+pub use rocks::{RocksDB, new_rocksdb};
 
 /// An object implementing get.
 pub trait Get {
diff --git a/common/db/src/rocks.rs b/common/db/src/rocks.rs
index 0866bb7e..bf8d2574 100644
--- a/common/db/src/rocks.rs
+++ b/common/db/src/rocks.rs
@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use rocksdb::{ThreadMode, Transaction, TransactionDB};
+use rocksdb::{DBCompressionType, ThreadMode, SingleThreaded, Options, Transaction, TransactionDB};
 
 use crate::*;
 
@@ -32,3 +32,10 @@ impl<T: ThreadMode + 'static> Db for Arc<TransactionDB<T>> {
     self.transaction()
   }
 }
+
+pub type RocksDB = Arc<TransactionDB<SingleThreaded>>;
+pub fn new_rocksdb(path: &str) -> RocksDB {
+  let mut options = Options::default();
+  options.set_compression_type(DBCompressionType::Zstd);
+  Arc::new(TransactionDB::open(&options, &Default::default(), path).unwrap())
+}
diff --git a/coordinator/src/main.rs b/coordinator/src/main.rs
index 56aed81a..a0d09cb1 100644
--- a/coordinator/src/main.rs
+++ b/coordinator/src/main.rs
@@ -669,12 +669,7 @@ pub async fn run<D: Db, Pro: Processors, P: P2p>(
 
 #[tokio::main]
 async fn main() {
-  let db = Arc::new(
-    rocksdb::TransactionDB::<rocksdb::SingleThreaded>::open_default(
-      env::var("DB_PATH").expect("path to DB wasn't specified"),
-    )
-    .unwrap(),
-  );
+  let db = serai_db::new_rocksdb(&env::var("DB_PATH").expect("path to DB wasn't specified"));
 
   let key = Zeroizing::new(<Ristretto as Ciphersuite>::F::ZERO); // TODO
   let p2p = LocalP2p::new(1).swap_remove(0); // TODO
diff --git a/message-queue/src/main.rs b/message-queue/src/main.rs
index 8827bab0..4dc5ad2f 100644
--- a/message-queue/src/main.rs
+++ b/message-queue/src/main.rs
@@ -16,7 +16,7 @@ use messages::*;
 mod queue;
 use queue::Queue;
 
-type Db = Arc<rocksdb::TransactionDB>;
+type Db = serai_db::RocksDB;
 
 lazy_static::lazy_static! {
   static ref KEYS: Arc<RwLock<HashMap<Service, <Ristretto as Ciphersuite>::G>>> =
@@ -117,12 +117,7 @@ async fn main() {
   log::info!("Starting message-queue service...");
 
   // Open the DB
-  let db = Arc::new(
-    rocksdb::TransactionDB::open_default(
-      serai_env::var("DB_PATH").expect("path to DB wasn't specified"),
-    )
-    .unwrap(),
-  );
+  let db = serai_db::new_rocksdb(&serai_env::var("DB_PATH").expect("path to DB wasn't specified"));
 
   let read_key = |str| {
     let key = serai_env::var(str)?;
diff --git a/processor/src/main.rs b/processor/src/main.rs
index c196efae..a156b1d5 100644
--- a/processor/src/main.rs
+++ b/processor/src/main.rs
@@ -1,6 +1,5 @@
 use std::{
   time::Duration,
-  sync::Arc,
   collections::{VecDeque, HashMap},
 };
 
@@ -730,12 +729,7 @@ async fn main() {
   }
   env_logger::init();
 
-  let db = Arc::new(
-    rocksdb::TransactionDB::<rocksdb::SingleThreaded>::open_default(
-      env::var("DB_PATH").expect("path to DB wasn't specified"),
-    )
-    .unwrap(),
-  );
+  let db = serai_db::new_rocksdb(&env::var("DB_PATH").expect("path to DB wasn't specified"));
 
   // Network configuration
   let url = {