feat(cubestore): Introduce limit for cache_max_entry_size (#7363)

cube-js · Oct 31, 2023 · 2295fe0 · 2295fe0
1 parent 5dbe0ab
commit 2295fe0
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 7 deletions.
diff --git a/rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs b/rust/cubestore/cubestore/src/cachestore/cache_rocksstore.rs
@@ -715,6 +715,15 @@ impl CacheStore for RocksCacheStore {
         item: CacheItem,
         update_if_not_exists: bool,
     ) -> Result<bool, CubeError> {
+        if item.get_value().len() >= self.store.config.cachestore_cache_max_entry_size() {
+            return Err(CubeError::user(format!(
+                "Unable to SET cache with '{}' key, exceeds maximum allowed size for payload: {}, max allowed: {}",
+                item.key,
+                humansize::format_size(item.get_value().len(), humansize::DECIMAL),
+                humansize::format_size(self.store.config.cachestore_cache_max_entry_size(), humansize::DECIMAL),
+            )));
+        }
+
         self.cache_eviction_manager
             .before_insert(item.get_value().len() as u64)
             .await?;
@@ -1521,6 +1530,38 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_cache_set_max_entry_size() -> Result<(), CubeError> {
+        init_test_logger().await;
+
+        let config = Config::test("cachestore_set").update_config(|mut c| {
+            c.cachestore_cache_max_entry_size = 1 << 20;
+            c
+        });
+
+        let (_, cachestore) =
+            RocksCacheStore::prepare_test_cachestore("cache_set_max_entry_size", config);
+
+        let err = cachestore
+            .cache_set(
+                CacheItem::new(
+                    "prefix:key-with-wrong-size".to_string(),
+                    Some(60),
+                    "a".repeat(2 << 20),
+                ),
+                false,
+            )
+            .await;
+
+        assert_eq!(err, Err(CubeError::user(
+            "Unable to SET cache with 'key-with-wrong-size' key, exceeds maximum allowed size for payload: 2.10 MB, max allowed: 1.05 MB".to_string()
+        )));
+
+        RocksCacheStore::cleanup_test_cachestore("cache_set_max_entry_size");
+
+        Ok(())
+    }
+
     async fn test_cachestore_force_eviction(
         name: &str,
         update_config: impl FnOnce(ConfigObjImpl) -> ConfigObjImpl,

diff --git a/rust/cubestore/cubestore/src/config/mod.rs b/rust/cubestore/cubestore/src/config/mod.rs
@@ -442,6 +442,8 @@ pub trait ConfigObj: DIService {
 
     fn cachestore_cache_max_size(&self) -> u64;
 
+    fn cachestore_cache_max_entry_size(&self) -> usize;
+
     fn cachestore_cache_compaction_trigger_size(&self) -> u64;
 
     fn cachestore_cache_max_keys(&self) -> u32;
@@ -582,6 +584,7 @@ pub struct ConfigObjImpl {
     pub cachestore_cache_eviction_loop_interval: u64,
     pub cachestore_cache_ttl_persist_loop_interval: u64,
     pub cachestore_cache_max_size: u64,
+    pub cachestore_cache_max_entry_size: usize,
     pub cachestore_cache_compaction_trigger_size: u64,
     pub cachestore_cache_threshold_to_force_eviction: u8,
     pub cachestore_queue_results_expire: u64,
@@ -798,6 +801,10 @@ impl ConfigObj for ConfigObjImpl {
         self.cachestore_cache_max_size
     }
 
+    fn cachestore_cache_max_entry_size(&self) -> usize {
+        self.cachestore_cache_max_entry_size
+    }
+
     fn cachestore_cache_compaction_trigger_size(&self) -> u64 {
         self.cachestore_cache_compaction_trigger_size
     }
@@ -1151,6 +1158,20 @@ impl Config {
             Some(32 << 20),
         );
 
+        let transport_max_message_size = env_parse_size(
+            "CUBESTORE_TRANSPORT_MAX_MESSAGE_SIZE",
+            64 << 20,
+            Some(256 << 20),
+            Some(16 << 20),
+        );
+
+        let cachestore_cache_max_entry_size = env_parse_size(
+            "CUBESTORE_CACHE_MAX_ENTRY_SIZE",
+            (64 << 20) - (1024 << 10),
+            Some(transport_max_message_size - (1024 << 10)),
+            None,
+        );
+
         let cachestore_cache_compaction_trigger_size = env_parse_size(
             "CUBESTORE_CACHE_COMPACTION_TRIGGER_SIZE",
             Self::calculate_cache_compaction_trigger_size(cachestore_cache_max_size),
@@ -1322,6 +1343,7 @@ impl Config {
                     Some(0),
                 ),
                 cachestore_cache_max_size: cachestore_cache_max_size as u64,
+                cachestore_cache_max_entry_size,
                 cachestore_cache_compaction_trigger_size,
                 cachestore_cache_threshold_to_force_eviction: 25,
                 cachestore_queue_results_expire: env_parse_duration(
@@ -1441,15 +1463,10 @@ impl Config {
                     * 1024
                     * 1024,
                 disk_space_cache_duration_secs: 300,
-                transport_max_message_size: env_parse_size(
-                    "CUBESTORE_TRANSPORT_MAX_MESSAGE_SIZE",
-                    64 << 20,
-                    Some(256 << 20),
-                    Some(16 << 20),
-                ),
+                transport_max_message_size,
                 transport_max_frame_size: env_parse_size(
                     "CUBESTORE_TRANSPORT_MAX_FRAME_SIZE",
-                    32 << 20,
+                    64 << 20,
                     Some(256 << 20),
                     Some(4 << 20),
                 ),
@@ -1535,6 +1552,7 @@ impl Config {
                 cachestore_cache_eviction_loop_interval: 60,
                 cachestore_cache_ttl_persist_loop_interval: 15,
                 cachestore_cache_max_size: 4096 << 20,
+                cachestore_cache_max_entry_size: 64 << 20,
                 cachestore_cache_compaction_trigger_size: 4096 * 2 << 20,
                 cachestore_cache_threshold_to_force_eviction: 25,
                 cachestore_queue_results_expire: 90,