fix(cubestore): Channel closed in query queue cache

rust/cubestore/cubestore/src/config/mod.rs

@@ -494,6 +494,8 @@ pub trait ConfigObj: DIService {
 
     fn query_cache_max_capacity_bytes(&self) -> u64;
 
+    fn query_queue_cache_max_capacity(&self) -> u64;
+
     fn query_cache_time_to_idle_secs(&self) -> Option<u64>;
 
     fn metadata_cache_max_capacity_bytes(&self) -> u64;
@@ -616,6 +618,7 @@ pub struct ConfigObjImpl {
     pub enable_startup_warmup: bool,
     pub malloc_trim_every_secs: u64,
     pub query_cache_max_capacity_bytes: u64,
+    pub query_queue_cache_max_capacity: u64,
     pub query_cache_time_to_idle_secs: Option<u64>,
     pub metadata_cache_max_capacity_bytes: u64,
     pub metadata_cache_time_to_idle_secs: u64,
@@ -908,6 +911,9 @@ impl ConfigObj for ConfigObjImpl {
     fn query_cache_max_capacity_bytes(&self) -> u64 {
         self.query_cache_max_capacity_bytes
     }
+    fn query_queue_cache_max_capacity(&self) -> u64 {
+        self.query_queue_cache_max_capacity
+    }
     fn query_cache_time_to_idle_secs(&self) -> Option<u64> {
         self.query_cache_time_to_idle_secs
     }
@@ -1443,6 +1449,10 @@ impl Config {
                     Some(16384 << 20),
                     Some(0),
                 ) as u64,
+                query_queue_cache_max_capacity: env_parse(
+                    "CUBESTORE_QUEUE_CACHE_MAX_CAPACITY",
+                    10000,
+                ),
                 query_cache_time_to_idle_secs: if query_cache_time_to_idle_secs == 0 {
                     None
                 } else {
@@ -1612,6 +1622,7 @@ impl Config {
                 enable_startup_warmup: true,
                 malloc_trim_every_secs: 0,
                 query_cache_max_capacity_bytes: 512 << 20,
+                query_queue_cache_max_capacity: 10000,
                 query_cache_time_to_idle_secs: Some(600),
                 metadata_cache_max_capacity_bytes: 0,
                 metadata_cache_time_to_idle_secs: 1_000,
@@ -2143,6 +2154,7 @@ impl Config {
         let query_cache = Arc::new(SqlResultCache::new(
             self.config_obj.query_cache_max_capacity_bytes(),
             self.config_obj.query_cache_time_to_idle_secs(),
+            self.config_obj.query_queue_cache_max_capacity(),
         ));
 
         let query_cache_to_move = query_cache.clone();

rust/cubestore/cubestore/src/queryplanner/mod.rs

@@ -790,7 +790,7 @@ pub mod tests {
             Arc::new(test_utils::MetaStoreMock {}),
             Arc::new(test_utils::CacheStoreMock {}),
             &vec![],
-            Arc::new(SqlResultCache::new(1 << 20, None)),
+            Arc::new(SqlResultCache::new(1 << 20, None, 10000)),
         )
     }
 

rust/cubestore/cubestore/src/sql/cache.rs

@@ -81,15 +81,19 @@ pub fn sql_result_cache_sizeof(key: &SqlResultCacheKey, df: &Arc<DataFrame>) ->
 }
 
 impl SqlResultCache {
-    pub fn new(capacity_bytes: u64, time_to_idle_secs: Option<u64>) -> Self {
+    pub fn new(
+        capacity_bytes: u64,
+        time_to_idle_secs: Option<u64>,
+        queue_cache_max_capacity: u64,
+    ) -> Self {
         let cache_builder = if let Some(time_to_idle_secs) = time_to_idle_secs {
             Cache::builder().time_to_idle(Duration::from_secs(time_to_idle_secs))
         } else {
             Cache::builder()
         };
 
         Self {
-            queue_cache: Mutex::new(lru::LruCache::new(1000)),
+            queue_cache: Mutex::new(lru::LruCache::new(queue_cache_max_capacity as usize)),
             result_cache: cache_builder
                 .max_capacity(capacity_bytes)
                 .weigher(sql_result_cache_sizeof)
@@ -139,6 +143,19 @@ impl SqlResultCache {
         let (sender, receiver) = {
             let key = queue_key.clone();
             let mut cache = self.queue_cache.lock().await;
+
+            if cache.contains(&key) {
+                if let Some(receiver) = cache.get(&key) {
+                    if receiver.has_changed().is_err() {
+                        log::error!("Queue cache contains closed channel");
+                        cache.pop(&key);
+                    }
+                } else {
+                    log::error!("Queue cache doesn't contains channel");
+                    cache.pop(&key);
+                }
+            }
+
             if !cache.contains(&key) {
                 let (tx, rx) = watch::channel(None);
                 cache.put(key, rx);
@@ -290,7 +307,7 @@ mod tests {
 
     #[tokio::test]
     async fn simple() -> Result<(), CubeError> {
-        let cache = SqlResultCache::new(1 << 20, Some(120));
+        let cache = SqlResultCache::new(1 << 20, Some(120), 1000);
         let schema = Arc::new(DFSchema::new(Vec::new())?);
         let plan = SerializedPlan::try_new(
             LogicalPlan::EmptyRelation {

rust/cubestore/cubestore/src/sql/mod.rs

@@ -1743,6 +1743,7 @@ mod tests {
                 Arc::new(SqlResultCache::new(
                     config.config_obj().query_cache_max_capacity_bytes(),
                     config.config_obj().query_cache_time_to_idle_secs(),
+                    1000,
                 )),
                 BasicProcessRateLimiter::new(),
             );
@@ -1819,6 +1820,7 @@ mod tests {
                 Arc::new(SqlResultCache::new(
                     config.config_obj().query_cache_max_capacity_bytes(),
                     config.config_obj().query_cache_time_to_idle_secs(),
+                    1000,
                 )),
                 BasicProcessRateLimiter::new(),
             );
@@ -1925,6 +1927,7 @@ mod tests {
                 Arc::new(SqlResultCache::new(
                     config.config_obj().query_cache_max_capacity_bytes(),
                     config.config_obj().query_cache_time_to_idle_secs(),
+                    1000,
                 )),
                 BasicProcessRateLimiter::new(),
             );