From f87e7fcfd5ff82b184dd2950c3b4750d32fc1f87 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Karolis=20Gudi=C5=A1kis?= <karolisgud@gmail.com>
Date: Thu, 28 Sep 2023 09:49:02 +0300
Subject: [PATCH] feat: Add block cache size to rocksdb config

---
 dozer-recordstore/src/lib.rs         | 11 ++++++---
 dozer-recordstore/src/rocksdb.rs     |  5 ++--
 dozer-storage/src/rocksdb_map.rs     | 18 +++++++++++---
 dozer-types/src/models/app_config.rs |  8 +++++-
 json_schemas/dozer.json              | 37 +++++++++++++++++++++++++---
 5 files changed, 66 insertions(+), 13 deletions(-)

diff --git a/dozer-recordstore/src/lib.rs b/dozer-recordstore/src/lib.rs
index ef971b2b23..3bab157542 100644
--- a/dozer-recordstore/src/lib.rs
+++ b/dozer-recordstore/src/lib.rs
@@ -109,7 +109,9 @@ impl ProcessorRecordStore {
     pub fn new(record_store: RecordStore) -> Result<Self, RecordStoreError> {
         match record_store {
             RecordStore::InMemory => Ok(Self::InMemory(in_memory::ProcessorRecordStore::new()?)),
-            RecordStore::Rocksdb => Ok(Self::Rocksdb(rocksdb::ProcessorRecordStore::new()?)),
+            RecordStore::Rocksdb(config) => {
+                Ok(Self::Rocksdb(rocksdb::ProcessorRecordStore::new(config)?))
+            }
         }
     }
 
@@ -178,7 +180,9 @@ impl ProcessorRecordStoreDeserializer {
             RecordStore::InMemory => Ok(Self::InMemory(
                 in_memory::ProcessorRecordStoreDeserializer::new()?,
             )),
-            RecordStore::Rocksdb => Ok(Self::Rocksdb(rocksdb::ProcessorRecordStore::new()?)),
+            RecordStore::Rocksdb(config) => {
+                Ok(Self::Rocksdb(rocksdb::ProcessorRecordStore::new(config)?))
+            }
         }
     }
 
@@ -253,6 +257,7 @@ struct ProcessorRecordForSerialization {
 
 #[cfg(test)]
 mod tests {
+    use dozer_types::models::app_config::RocksdbConfig;
     use std::time::Duration;
 
     use dozer_types::types::Timestamp;
@@ -283,7 +288,7 @@ mod tests {
     #[test]
     fn test_record_roundtrip() {
         test_record_roundtrip_impl(RecordStore::InMemory);
-        test_record_roundtrip_impl(RecordStore::Rocksdb);
+        test_record_roundtrip_impl(RecordStore::Rocksdb(RocksdbConfig::default()));
     }
 
     fn test_record_serialization_roundtrip_impl(record_store_kind: RecordStore) {
diff --git a/dozer-recordstore/src/rocksdb.rs b/dozer-recordstore/src/rocksdb.rs
index 09a82d2d0b..f54f28d86f 100644
--- a/dozer-recordstore/src/rocksdb.rs
+++ b/dozer-recordstore/src/rocksdb.rs
@@ -1,6 +1,7 @@
 use std::sync::atomic::AtomicU64;
 
 use dozer_storage::RocksdbMap;
+use dozer_types::models::app_config::RocksdbConfig;
 use dozer_types::types::Field;
 use tempdir::TempDir;
 
@@ -14,10 +15,10 @@ pub struct ProcessorRecordStore {
 }
 
 impl ProcessorRecordStore {
-    pub fn new() -> Result<Self, RecordStoreError> {
+    pub fn new(config: RocksdbConfig) -> Result<Self, RecordStoreError> {
         let temp_dir = TempDir::new("rocksdb_processor_record_store")
             .map_err(RecordStoreError::FailedToCreateTempDir)?;
-        let records = RocksdbMap::<u64, Vec<Field>>::create(temp_dir.path())?;
+        let records = RocksdbMap::<u64, Vec<Field>>::create(temp_dir.path(), config)?;
 
         Ok(Self {
             _temp_dir: temp_dir,
diff --git a/dozer-storage/src/rocksdb_map.rs b/dozer-storage/src/rocksdb_map.rs
index f2b59e73d6..2075644ff9 100644
--- a/dozer-storage/src/rocksdb_map.rs
+++ b/dozer-storage/src/rocksdb_map.rs
@@ -1,8 +1,9 @@
 use std::path::Path;
 
-use rocksdb::DB;
+use rocksdb::{BlockBasedOptions, Cache, Options, DB};
 
 use dozer_types::borrow::IntoOwned;
+use dozer_types::models::app_config::RocksdbConfig;
 
 use crate::{errors::StorageError, BorrowEncode, Encode, LmdbVal};
 
@@ -17,8 +18,19 @@ impl<K: BorrowEncode, V: LmdbVal> RocksdbMap<K, V>
 where
     for<'a> V::Borrowed<'a>: IntoOwned<V>,
 {
-    pub fn create(path: &Path) -> Result<Self, StorageError> {
-        let db = DB::open_default(path)?;
+    pub fn create(path: &Path, config: RocksdbConfig) -> Result<Self, StorageError> {
+        let mut options = Options::default();
+        options.create_if_missing(true);
+
+        if let Some(block_cache_size) = config.block_cache_size {
+            let mut block_options = BlockBasedOptions::default();
+            let cache = Cache::new_lru_cache(block_cache_size);
+            block_options.set_block_cache(&cache);
+
+            options.set_block_based_table_factory(&block_options);
+        }
+
+        let db = DB::open(&options, path)?;
         Ok(Self {
             db,
             _key: std::marker::PhantomData,
diff --git a/dozer-types/src/models/app_config.rs b/dozer-types/src/models/app_config.rs
index 13647565d2..799229d419 100644
--- a/dozer-types/src/models/app_config.rs
+++ b/dozer-types/src/models/app_config.rs
@@ -63,7 +63,13 @@ pub struct S3Storage {
 pub enum RecordStore {
     #[default]
     InMemory,
-    Rocksdb,
+    Rocksdb(RocksdbConfig),
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema, Default)]
+#[serde(deny_unknown_fields)]
+pub struct RocksdbConfig {
+    pub block_cache_size: Option<usize>,
 }
 
 pub fn default_persist_queue_capacity() -> u32 {
diff --git a/json_schemas/dozer.json b/json_schemas/dozer.json
index 84a3a24d2d..429b75ee03 100644
--- a/json_schemas/dozer.json
+++ b/json_schemas/dozer.json
@@ -1307,10 +1307,25 @@
       "additionalProperties": false
     },
     "RecordStore": {
-      "type": "string",
-      "enum": [
-        "InMemory",
-        "Rocksdb"
+      "oneOf": [
+        {
+          "type": "string",
+          "enum": [
+            "InMemory"
+          ]
+        },
+        {
+          "type": "object",
+          "required": [
+            "Rocksdb"
+          ],
+          "properties": {
+            "Rocksdb": {
+              "$ref": "#/definitions/RocksdbConfig"
+            }
+          },
+          "additionalProperties": false
+        }
       ]
     },
     "RefreshConfig": {
@@ -1351,6 +1366,20 @@
       },
       "additionalProperties": false
     },
+    "RocksdbConfig": {
+      "type": "object",
+      "properties": {
+        "block_cache_size": {
+          "type": [
+            "integer",
+            "null"
+          ],
+          "format": "uint",
+          "minimum": 0.0
+        }
+      },
+      "additionalProperties": false
+    },
     "S3Details": {
       "type": "object",
       "required": [