serai_db/
rocks.rs

1use alloc::sync::Arc;
2
3use rocksdb::{
4  DBCompressionType, SingleThreaded, LogLevel, WriteOptions, Transaction as RocksTransaction,
5  Options, OptimisticTransactionDB,
6};
7
8use crate::{Get, Db};
9
10/// A transaction for a RocksDB database.
11pub struct Transaction<'db> {
12  db: &'db OptimisticTransactionDB<SingleThreaded>,
13  txn: RocksTransaction<'db, OptimisticTransactionDB<SingleThreaded>>,
14}
15
16impl Get for Transaction<'_> {
17  fn get(&self, key: impl AsRef<[u8]>) -> Option<impl AsRef<[u8]>> {
18    self.txn.get(key).expect("couldn't read from RocksDB via transaction")
19  }
20}
21impl crate::Transaction for Transaction<'_> {
22  fn set(&mut self, key: impl AsRef<[u8]>, value: impl AsRef<[u8]>) {
23    self.txn.put(key, value).expect("couldn't write to RocksDB via transaction");
24  }
25  fn del(&mut self, key: impl AsRef<[u8]>) {
26    self.txn.delete(key).expect("couldn't delete from RocksDB via transaction");
27  }
28  fn commit(self) {
29    self.txn.commit().expect("couldn't commit to RocksDB via transaction");
30    /*
31      TODO: This is incredibly aggressive flushing due to an observed fault where the WAL was
32      _never_ flushed, despite trusting RocksDB to do so as soon as the size exceeded the limit.
33      This is the overkill solution to the above problem. Proper tests for this should be added and
34      a solution which does bound the WAL, but doesn't remove the performance benefit of the WAL,
35      should be implemented.
36    */
37    self.db.flush_wal(true).expect("couldn't flush RocksDB WAL");
38    self.db.flush().expect("couldn't flush RocksDB");
39  }
40}
41
42impl Get for Arc<OptimisticTransactionDB<SingleThreaded>> {
43  fn get(&self, key: impl AsRef<[u8]>) -> Option<impl AsRef<[u8]>> {
44    OptimisticTransactionDB::get(self, key).expect("couldn't read from RocksDB")
45  }
46}
47impl Db for Arc<OptimisticTransactionDB<SingleThreaded>> {
48  type Transaction<'db> = Transaction<'db>;
49  fn txn(&mut self) -> Self::Transaction<'_> {
50    let mut opts = WriteOptions::default();
51    opts.set_sync(true);
52    Transaction { db: self, txn: self.transaction_opt(&opts, &Default::default()) }
53  }
54}
55
56/// The RocksDB database type.
57pub type RocksDB = Arc<OptimisticTransactionDB<SingleThreaded>>;
58
59/// Open a RocksDB database.
60///
61/// This will create the database if it does not already exist.
62///
63/// This will panic if opening/creating the database errors.
64pub fn new_rocksdb(path: &str) -> RocksDB {
65  let mut options = Options::default();
66  options.create_if_missing(true);
67  options.set_compression_type(DBCompressionType::Zstd);
68
69  options.set_wal_compression_type(DBCompressionType::Zstd);
70  // 10 MB
71  options.set_max_total_wal_size(10 * 1024 * 1024);
72  options.set_wal_size_limit_mb(10);
73
74  options.set_log_level(LogLevel::Warn);
75  // 1 MB
76  options.set_max_log_file_size(1024 * 1024);
77  options.set_recycle_log_file_num(1);
78
79  Arc::new(OptimisticTransactionDB::open(&options, path).expect("failed to create/open RocksDB DB"))
80}
serai_db/rocks.rs

serai_db/
rocks.rs
