refactor(streaming): use LRU cache for group top-n (#6491)

* temp commit

* top_n cache use lru manager

* ready for review

* call evict manually

* fix

* fix

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: wcy-fdu

src/stream/src/executor/mview/materialize.rs

@@ -210,7 +210,7 @@ impl<S: StateStore> MaterializeExecutor<S> {
                     if let Some(vnode_bitmap) = b.as_update_vnode_bitmap(self.actor_context.id) {
                         let _ = self.state_table.update_vnode_bitmap(vnode_bitmap);
                     }
-
+                    self.materialize_cache.evict();
                     Message::Barrier(b)
                 }
             }
@@ -501,6 +501,10 @@ impl MaterializeCache {
     pub fn put(&mut self, key: Vec<u8>, value: Option<CompactedRow>) {
         self.data.push(key, value);
     }
+
+    fn evict(&mut self) {
+        self.data.evict()
+    }
 }
 #[cfg(test)]
 mod tests {

src/stream/src/executor/top_n/group_top_n.rs

@@ -12,8 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::collections::hash_map::Entry::Vacant;
-use std::collections::{HashMap, HashSet};
+use std::collections::HashSet;
 use std::sync::Arc;
 
 use async_trait::async_trait;
@@ -30,7 +29,7 @@ use risingwave_storage::StateStore;
 use super::top_n_cache::TopNCacheTrait;
 use super::utils::*;
 use super::TopNCache;
-use crate::cache::cache_may_stale;
+use crate::cache::{cache_may_stale, EvictableHashMap, ExecutorCache, LruManagerRef};
 use crate::common::table::state_table::StateTable;
 use crate::error::StreamResult;
 use crate::executor::error::StreamExecutorResult;
@@ -52,6 +51,8 @@ impl<S: StateStore> GroupTopNExecutor<S, false> {
         executor_id: u64,
         group_by: Vec<usize>,
         state_table: StateTable<S>,
+        lru_manager: Option<LruManagerRef>,
+        cache_size: usize,
     ) -> StreamResult<Self> {
         let info = input.info();
         let schema = input.schema().clone();
@@ -68,6 +69,8 @@ impl<S: StateStore> GroupTopNExecutor<S, false> {
                 executor_id,
                 group_by,
                 state_table,
+                lru_manager,
+                cache_size,
             )?,
         })
     }
@@ -85,6 +88,8 @@ impl<S: StateStore> GroupTopNExecutor<S, true> {
         executor_id: u64,
         group_by: Vec<usize>,
         state_table: StateTable<S>,
+        lru_manager: Option<LruManagerRef>,
+        cache_size: usize,
     ) -> StreamResult<Self> {
         let info = input.info();
         let schema = input.schema().clone();
@@ -102,6 +107,8 @@ impl<S: StateStore> GroupTopNExecutor<S, true> {
                 executor_id,
                 group_by,
                 state_table,
+                lru_manager,
+                cache_size,
             )?,
         })
     }
@@ -132,7 +139,7 @@ pub struct InnerGroupTopNExecutorNew<S: StateStore, const WITH_TIES: bool> {
     group_by: Vec<usize>,
 
     /// group key -> cache for this group
-    caches: HashMap<Vec<Datum>, TopNCache<WITH_TIES>>,
+    caches: GroupTopNCache<WITH_TIES>,
 
     /// The number of fields of the ORDER BY clause, and will be used to split key into `CacheKey`.
     order_by_len: usize,
@@ -153,6 +160,8 @@ impl<S: StateStore, const WITH_TIES: bool> InnerGroupTopNExecutorNew<S, WITH_TIE
         executor_id: u64,
         group_by: Vec<usize>,
         state_table: StateTable<S>,
+        lru_manager: Option<LruManagerRef>,
+        cache_size: usize,
     ) -> StreamResult<Self> {
         // order_pairs is superset of pk
         assert!(order_pairs
@@ -196,13 +205,47 @@ impl<S: StateStore, const WITH_TIES: bool> InnerGroupTopNExecutorNew<S, WITH_TIE
             pk_indices,
             internal_key_indices,
             group_by,
-            caches: HashMap::new(),
+            caches: GroupTopNCache::new(lru_manager, cache_size),
             order_by_len,
             cache_key_serde,
         })
     }
 }
 
+pub struct GroupTopNCache<const WITH_TIES: bool> {
+    data: ExecutorCache<Vec<Datum>, TopNCache<WITH_TIES>>,
+}
+
+impl<const WITH_TIES: bool> GroupTopNCache<WITH_TIES> {
+    pub fn new(lru_manager: Option<LruManagerRef>, cache_size: usize) -> Self {
+        let cache = if let Some(lru_manager) = lru_manager {
+            ExecutorCache::Managed(lru_manager.create_cache())
+        } else {
+            ExecutorCache::Local(EvictableHashMap::new(cache_size))
+        };
+        Self { data: cache }
+    }
+
+    fn clear(&mut self) {
+        self.data.clear()
+    }
+
+    fn get_mut(&mut self, key: &[Datum]) -> Option<&mut TopNCache<WITH_TIES>> {
+        self.data.get_mut(key)
+    }
+
+    fn contains(&mut self, key: &[Datum]) -> bool {
+        self.data.contains(key)
+    }
+
+    fn insert(&mut self, key: Vec<Datum>, value: TopNCache<WITH_TIES>) {
+        self.data.push(key, value);
+    }
+
+    fn evict(&mut self) {
+        self.data.evict()
+    }
+}
 #[async_trait]
 impl<S: StateStore, const WITH_TIES: bool> TopNExecutorBase
     for InnerGroupTopNExecutorNew<S, WITH_TIES>
@@ -229,12 +272,12 @@ where
 
             // If 'self.caches' does not already have a cache for the current group, create a new
             // cache for it and insert it into `self.caches`
-            if let Vacant(entry) = self.caches.entry(group_key) {
+            if !self.caches.contains(&group_key) {
                 let mut topn_cache = TopNCache::new(self.offset, self.limit, self.order_by_len);
                 self.managed_state
                     .init_topn_cache(Some(&pk_prefix), &mut topn_cache, self.order_by_len)
                     .await?;
-                entry.insert(topn_cache);
+                self.caches.insert(group_key, topn_cache);
             }
             let cache = self.caches.get_mut(&pk_prefix.0).unwrap();
 
@@ -291,6 +334,10 @@ where
         }
     }
 
+    fn evict(&mut self) {
+        self.caches.evict()
+    }
+
     async fn init(&mut self, epoch: EpochPair) -> StreamExecutorResult<()> {
         self.managed_state.state_table.init_epoch(epoch);
         Ok(())
@@ -406,6 +453,8 @@ mod tests {
                 1,
                 vec![1],
                 state_table,
+                None,
+                0,
             )
             .unwrap(),
         );
@@ -504,6 +553,8 @@ mod tests {
                 1,
                 vec![1],
                 state_table,
+                None,
+                0,
             )
             .unwrap(),
         );
@@ -594,6 +645,8 @@ mod tests {
                 1,
                 vec![1, 2],
                 state_table,
+                None,
+                0,
             )
             .unwrap(),
         );

src/stream/src/executor/top_n/utils.rs

@@ -58,6 +58,7 @@ pub trait TopNExecutorBase: Send + 'static {
         unreachable!()
     }
 
+    fn evict(&mut self) {}
     async fn init(&mut self, epoch: EpochPair) -> StreamExecutorResult<()>;
 }
 
@@ -120,7 +121,7 @@ where
                     if let Some(vnode_bitmap) = barrier.as_update_vnode_bitmap(self.ctx.id) {
                         self.inner.update_vnode_bitmap(vnode_bitmap);
                     }
-
+                    self.inner.evict();
                     yield Message::Barrier(barrier)
                 }
             };

src/stream/src/from_proto/group_top_n.rs

@@ -31,7 +31,7 @@ impl ExecutorBuilder for GroupTopNExecutorBuilder {
         mut params: ExecutorParams,
         node: &Self::Node,
         store: impl StateStore,
-        _stream: &mut LocalStreamManagerCore,
+        stream: &mut LocalStreamManagerCore,
     ) -> StreamResult<BoxedExecutor> {
         let group_by = node
             .get_group_key()
@@ -54,6 +54,8 @@ impl ExecutorBuilder for GroupTopNExecutorBuilder {
                 params.executor_id,
                 group_by,
                 state_table,
+                stream.context.lru_manager.clone(),
+                1 << 16,
             )?
             .boxed())
         } else {
@@ -67,6 +69,8 @@ impl ExecutorBuilder for GroupTopNExecutorBuilder {
                 params.executor_id,
                 group_by,
                 state_table,
+                stream.context.lru_manager.clone(),
+                1 << 16,
             )?
             .boxed())
         }