feat: introduce TableRouteCache to PartitionRuleManager (#3922)

* chore: add `CompositeTableRouteCacheRef` to `PartitionRuleManager`

* chore: update comments

* fix: add metrics for `get`

* chore: apply suggestions from CR

* chore: correct cache name

* feat: implement `LayeredCacheRegistry`

* fix: invalidate logical tables by physical table id

* refactor: replace `CacheRegistry` with `LayeredCacheRegistry`

* chore: update comments

* chore: apply suggestions from CR

* chore: fix fmt

* refactor: use `TableRouteCache` instead

* chore: apply suggestions from CR

* chore: fix clippy

src/common/meta/src/cache.rs

@@ -19,10 +19,12 @@ mod table;
 
 pub use container::{CacheContainer, Initializer, Invalidator, TokenFilter};
 pub use flow::{new_table_flownode_set_cache, TableFlownodeSetCache, TableFlownodeSetCacheRef};
-pub use registry::{CacheRegistry, CacheRegistryBuilder, CacheRegistryRef};
-pub use table::{
-    new_composite_table_route_cache, new_table_info_cache, new_table_name_cache,
-    new_table_route_cache, CompositeTableRoute, CompositeTableRouteCache,
-    CompositeTableRouteCacheRef, TableInfoCache, TableInfoCacheRef, TableNameCache,
-    TableNameCacheRef, TableRouteCache, TableRouteCacheRef,
+pub use registry::{
+    CacheRegistry, CacheRegistryBuilder, CacheRegistryRef, LayeredCacheRegistry,
+    LayeredCacheRegistryBuilder, LayeredCacheRegistryRef,
+};
+pub use table::{
+    new_table_info_cache, new_table_name_cache, new_table_route_cache, TableInfoCache,
+    TableInfoCacheRef, TableNameCache, TableNameCacheRef, TableRoute, TableRouteCache,
+    TableRouteCacheRef,
 };

src/common/meta/src/cache/container.rs

@@ -101,9 +101,18 @@ where
 {
     /// Returns a _clone_ of the value corresponding to the key.
     pub async fn get(&self, key: K) -> Result<Option<V>> {
+        metrics::CACHE_CONTAINER_CACHE_GET
+            .with_label_values(&[&self.name])
+            .inc();
         let moved_init = self.initializer.clone();
         let moved_key = key;
         let init = async move {
+            metrics::CACHE_CONTAINER_CACHE_MISS
+                .with_label_values(&[&self.name])
+                .inc();
+            let _timer = metrics::CACHE_CONTAINER_LOAD_CACHE
+                .with_label_values(&[&self.name])
+                .start_timer();
             moved_init(&moved_key)
                 .await
                 .transpose()
@@ -163,6 +172,9 @@ where
             metrics::CACHE_CONTAINER_CACHE_MISS
                 .with_label_values(&[&self.name])
                 .inc();
+            let _timer = metrics::CACHE_CONTAINER_LOAD_CACHE
+                .with_label_values(&[&self.name])
+                .start_timer();
 
             moved_init(&moved_key)
                 .await

src/common/meta/src/cache/registry.rs

@@ -14,6 +14,7 @@
 
 use std::sync::Arc;
 
+use anymap2::SendSyncAnyMap;
 use futures::future::join_all;
 
 use crate::cache_invalidator::{CacheInvalidator, Context};
@@ -21,6 +22,65 @@ use crate::error::Result;
 use crate::instruction::CacheIdent;
 
 pub type CacheRegistryRef = Arc<CacheRegistry>;
+pub type LayeredCacheRegistryRef = Arc<LayeredCacheRegistry>;
+
+/// [LayeredCacheRegistry] Builder.
+#[derive(Default)]
+pub struct LayeredCacheRegistryBuilder {
+    registry: LayeredCacheRegistry,
+}
+
+impl LayeredCacheRegistryBuilder {
+    /// Adds [CacheRegistry] into the next layer.
+    ///
+    /// During cache invalidation, [LayeredCacheRegistry] ensures sequential invalidation
+    /// of each layer (after the previous layer).
+    pub fn add_cache_registry(mut self, registry: CacheRegistry) -> Self {
+        self.registry.layers.push(registry);
+
+        self
+    }
+
+    /// Returns __cloned__ the value stored in the collection for the type `T`, if it exists.
+    pub fn get<T: Send + Sync + Clone + 'static>(&self) -> Option<T> {
+        self.registry.get()
+    }
+
+    /// Builds the [LayeredCacheRegistry]
+    pub fn build(self) -> LayeredCacheRegistry {
+        self.registry
+    }
+}
+
+/// [LayeredCacheRegistry] invalidate caches sequentially from the first layer.
+#[derive(Default)]
+pub struct LayeredCacheRegistry {
+    layers: Vec<CacheRegistry>,
+}
+
+#[async_trait::async_trait]
+impl CacheInvalidator for LayeredCacheRegistry {
+    async fn invalidate(&self, ctx: &Context, caches: &[CacheIdent]) -> Result<()> {
+        let mut results = Vec::with_capacity(self.layers.len());
+        for registry in &self.layers {
+            results.push(registry.invalidate(ctx, caches).await);
+        }
+        results.into_iter().collect::<Result<Vec<_>>>().map(|_| ())
+    }
+}
+
+impl LayeredCacheRegistry {
+    /// Returns __cloned__ the value stored in the collection for the type `T`, if it exists.
+    pub fn get<T: Send + Sync + Clone + 'static>(&self) -> Option<T> {
+        for registry in &self.layers {
+            if let Some(cache) = registry.get::<T>() {
+                return Some(cache);
+            }
+        }
+
+        None
+    }
+}
 
 /// [CacheRegistryBuilder] provides ability of
 /// - Register the `cache` which implements the [CacheInvalidator] trait into [CacheRegistry].
@@ -31,11 +91,13 @@ pub struct CacheRegistryBuilder {
 }
 
 impl CacheRegistryBuilder {
+    /// Adds the cache.
     pub fn add_cache<T: CacheInvalidator + 'static>(mut self, cache: Arc<T>) -> Self {
         self.registry.register(cache);
         self
     }
 
+    /// Builds [CacheRegistry].
     pub fn build(self) -> CacheRegistry {
         self.registry
     }
@@ -46,7 +108,7 @@ impl CacheRegistryBuilder {
 #[derive(Default)]
 pub struct CacheRegistry {
     indexes: Vec<Arc<dyn CacheInvalidator>>,
-    registry: anymap2::SendSyncAnyMap,
+    registry: SendSyncAnyMap,
 }
 
 #[async_trait::async_trait]
@@ -80,7 +142,7 @@ impl CacheRegistry {
 
 #[cfg(test)]
 mod tests {
-    use std::sync::atomic::{AtomicI32, Ordering};
+    use std::sync::atomic::{AtomicBool, AtomicI32, Ordering};
     use std::sync::Arc;
 
     use moka::future::{Cache, CacheBuilder};
@@ -89,7 +151,10 @@ mod tests {
     use crate::cache::*;
     use crate::instruction::CacheIdent;
 
-    fn test_cache(name: &str) -> CacheContainer<String, String, CacheIdent> {
+    fn test_cache(
+        name: &str,
+        invalidator: Invalidator<String, String, CacheIdent>,
+    ) -> CacheContainer<String, String, CacheIdent> {
         let cache: Cache<String, String> = CacheBuilder::new(128).build();
         let filter: TokenFilter<CacheIdent> = Box::new(|_| true);
         let counter = Arc::new(AtomicI32::new(0));
@@ -98,13 +163,14 @@ mod tests {
             moved_counter.fetch_add(1, Ordering::Relaxed);
             Box::pin(async { Ok(Some("hi".to_string())) })
         });
-        let invalidator: Invalidator<String, String, CacheIdent> =
-            Box::new(|_, _| Box::pin(async { Ok(()) }));
 
         CacheContainer::new(name.to_string(), cache, invalidator, init, filter)
     }
 
-    fn test_i32_cache(name: &str) -> CacheContainer<i32, String, CacheIdent> {
+    fn test_i32_cache(
+        name: &str,
+        invalidator: Invalidator<i32, String, CacheIdent>,
+    ) -> CacheContainer<i32, String, CacheIdent> {
         let cache: Cache<i32, String> = CacheBuilder::new(128).build();
         let filter: TokenFilter<CacheIdent> = Box::new(|_| true);
         let counter = Arc::new(AtomicI32::new(0));
@@ -113,8 +179,6 @@ mod tests {
             moved_counter.fetch_add(1, Ordering::Relaxed);
             Box::pin(async { Ok(Some("foo".to_string())) })
         });
-        let invalidator: Invalidator<i32, String, CacheIdent> =
-            Box::new(|_, _| Box::pin(async { Ok(()) }));
 
         CacheContainer::new(name.to_string(), cache, invalidator, init, filter)
     }
@@ -122,8 +186,12 @@ mod tests {
     #[tokio::test]
     async fn test_register() {
         let builder = CacheRegistryBuilder::default();
-        let i32_cache = Arc::new(test_i32_cache("i32_cache"));
-        let cache = Arc::new(test_cache("string_cache"));
+        let invalidator: Invalidator<_, String, CacheIdent> =
+            Box::new(|_, _| Box::pin(async { Ok(()) }));
+        let i32_cache = Arc::new(test_i32_cache("i32_cache", invalidator));
+        let invalidator: Invalidator<_, String, CacheIdent> =
+            Box::new(|_, _| Box::pin(async { Ok(()) }));
+        let cache = Arc::new(test_cache("string_cache", invalidator));
         let registry = builder.add_cache(i32_cache).add_cache(cache).build();
 
         let cache = registry
@@ -136,4 +204,45 @@ mod tests {
             .unwrap();
         assert_eq!(cache.name(), "string_cache");
     }
+
+    #[tokio::test]
+    async fn test_layered_registry() {
+        let builder = LayeredCacheRegistryBuilder::default();
+        // 1st layer
+        let counter = Arc::new(AtomicBool::new(false));
+        let moved_counter = counter.clone();
+        let invalidator: Invalidator<String, String, CacheIdent> = Box::new(move |_, _| {
+            let counter = moved_counter.clone();
+            Box::pin(async move {
+                assert!(!counter.load(Ordering::Relaxed));
+                counter.store(true, Ordering::Relaxed);
+                Ok(())
+            })
+        });
+        let cache = Arc::new(test_cache("string_cache", invalidator));
+        let builder =
+            builder.add_cache_registry(CacheRegistryBuilder::default().add_cache(cache).build());
+        // 2nd layer
+        let moved_counter = counter.clone();
+        let invalidator: Invalidator<i32, String, CacheIdent> = Box::new(move |_, _| {
+            let counter = moved_counter.clone();
+            Box::pin(async move {
+                assert!(counter.load(Ordering::Relaxed));
+                Ok(())
+            })
+        });
+        let i32_cache = Arc::new(test_i32_cache("i32_cache", invalidator));
+        let builder = builder
+            .add_cache_registry(CacheRegistryBuilder::default().add_cache(i32_cache).build());
+
+        let registry = builder.build();
+        let cache = registry
+            .get::<Arc<CacheContainer<i32, String, CacheIdent>>>()
+            .unwrap();
+        assert_eq!(cache.name(), "i32_cache");
+        let cache = registry
+            .get::<Arc<CacheContainer<String, String, CacheIdent>>>()
+            .unwrap();
+        assert_eq!(cache.name(), "string_cache");
+    }
 }

src/common/meta/src/cache/table.rs

@@ -12,14 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod composite_table_route;
 mod table_info;
 mod table_name;
 mod table_route;
-pub use composite_table_route::{
-    new_composite_table_route_cache, CompositeTableRoute, CompositeTableRouteCache,
-    CompositeTableRouteCacheRef,
-};
 pub use table_info::{new_table_info_cache, TableInfoCache, TableInfoCacheRef};
 pub use table_name::{new_table_name_cache, TableNameCache, TableNameCacheRef};
 pub use table_route::{new_table_route_cache, TableRoute, TableRouteCache, TableRouteCacheRef};

src/common/meta/src/cache/table/composite_table_route.rs

@@ -1,274 +0,0 @@
-// Copyright 2023 Greptime Team
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//     http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-use std::sync::Arc;
-
-use futures::future::BoxFuture;
-use moka::future::Cache;
-use snafu::OptionExt;
-use store_api::storage::TableId;
-
-use crate::cache::table::{TableRoute, TableRouteCacheRef};
-use crate::cache::{CacheContainer, Initializer};
-use crate::error;
-use crate::error::Result;
-use crate::instruction::CacheIdent;
-use crate::key::table_route::{LogicalTableRouteValue, PhysicalTableRouteValue};
-
-/// [CompositeTableRoute] stores all level routes of a table.
-/// - Stores [PhysicalTableRouteValue] for logical table.
-/// - Stores [LogicalTableRouteValue], [PhysicalTableRouteValue] for the logical table.
-#[derive(Clone)]
-pub enum CompositeTableRoute {
-    Physical(Arc<PhysicalTableRouteValue>),
-    Logical(Arc<LogicalTableRouteValue>, Arc<PhysicalTableRouteValue>),
-}
-
-impl CompositeTableRoute {
-    /// Returns true if it's physical table.
-    pub fn is_physical(&self) -> bool {
-        matches!(self, CompositeTableRoute::Physical(_))
-    }
-
-    /// Returns [PhysicalTableRouteValue] reference.
-    pub fn as_physical_table_route_ref(&self) -> &Arc<PhysicalTableRouteValue> {
-        match self {
-            CompositeTableRoute::Physical(route) => route,
-            CompositeTableRoute::Logical(_, route) => route,
-        }
-    }
-
-    /// Returns [LogicalTableRouteValue] reference if it's [CompositeTableRoute::Logical]; Otherwise returns [None].
-    pub fn as_logical_table_route_ref(&self) -> Option<&Arc<LogicalTableRouteValue>> {
-        match self {
-            CompositeTableRoute::Physical(_) => None,
-            CompositeTableRoute::Logical(route, _) => Some(route),
-        }
-    }
-}
-
-/// [CompositeTableRouteCache] caches the [TableId] to [CompositeTableRoute] mapping.
-pub type CompositeTableRouteCache = CacheContainer<TableId, Arc<CompositeTableRoute>, CacheIdent>;
-
-pub type CompositeTableRouteCacheRef = Arc<CompositeTableRouteCache>;
-
-/// Constructs a [CompositeTableRouteCache].
-pub fn new_composite_table_route_cache(
-    name: String,
-    cache: Cache<TableId, Arc<CompositeTableRoute>>,
-    table_route_cache: TableRouteCacheRef,
-) -> CompositeTableRouteCache {
-    let init = init_factory(table_route_cache);
-
-    CacheContainer::new(name, cache, Box::new(invalidator), init, Box::new(filter))
-}
-
-fn init_factory(
-    table_route_cache: TableRouteCacheRef,
-) -> Initializer<TableId, Arc<CompositeTableRoute>> {
-    Arc::new(move |table_id| {
-        let table_route_cache = table_route_cache.clone();
-        Box::pin(async move {
-            let table_route_value = table_route_cache
-                .get(*table_id)
-                .await?
-                .context(error::ValueNotExistSnafu)?;
-            match table_route_value.as_ref() {
-                TableRoute::Physical(physical_table_route) => Ok(Some(Arc::new(
-                    CompositeTableRoute::Physical(physical_table_route.clone()),
-                ))),
-                TableRoute::Logical(logical_table_route) => {
-                    let physical_table_id = logical_table_route.physical_table_id();
-                    let physical_table_route = table_route_cache
-                        .get(physical_table_id)
-                        .await?
-                        .context(error::ValueNotExistSnafu)?;
-
-                    let physical_table_route = physical_table_route
-                        .as_physical_table_route_ref()
-                        .with_context(|| error::UnexpectedSnafu {
-                            err_msg: format!(
-                                "Expected the physical table route, but got logical table route, table: {table_id}"
-                            ),
-                        })?;
-
-                    Ok(Some(Arc::new(CompositeTableRoute::Logical(
-                        logical_table_route.clone(),
-                        physical_table_route.clone(),
-                    ))))
-                }
-            }
-        })
-    })
-}
-
-fn invalidator<'a>(
-    cache: &'a Cache<TableId, Arc<CompositeTableRoute>>,
-    ident: &'a CacheIdent,
-) -> BoxFuture<'a, Result<()>> {
-    Box::pin(async move {
-        if let CacheIdent::TableId(table_id) = ident {
-            cache.invalidate(table_id).await
-        }
-        Ok(())
-    })
-}
-
-fn filter(ident: &CacheIdent) -> bool {
-    matches!(ident, CacheIdent::TableId(_))
-}
-
-#[cfg(test)]
-mod tests {
-    use std::collections::HashMap;
-    use std::sync::Arc;
-
-    use moka::future::CacheBuilder;
-    use store_api::storage::RegionId;
-
-    use super::*;
-    use crate::cache::new_table_route_cache;
-    use crate::ddl::test_util::create_table::test_create_table_task;
-    use crate::ddl::test_util::test_create_logical_table_task;
-    use crate::key::table_route::TableRouteValue;
-    use crate::key::TableMetadataManager;
-    use crate::kv_backend::memory::MemoryKvBackend;
-    use crate::peer::Peer;
-    use crate::rpc::router::{Region, RegionRoute};
-
-    #[tokio::test]
-    async fn test_cache_with_physical_table_route() {
-        let mem_kv = Arc::new(MemoryKvBackend::default());
-        let table_metadata_manager = TableMetadataManager::new(mem_kv.clone());
-        let cache = CacheBuilder::new(128).build();
-        let table_route_cache = Arc::new(new_table_route_cache(
-            "test".to_string(),
-            cache,
-            mem_kv.clone(),
-        ));
-        let cache = CacheBuilder::new(128).build();
-        let cache =
-            new_composite_table_route_cache("test".to_string(), cache, table_route_cache.clone());
-
-        let result = cache.get(1024).await.unwrap();
-        assert!(result.is_none());
-        let task = test_create_table_task("my_table", 1024);
-        let table_id = 10;
-        let region_id = RegionId::new(table_id, 1);
-        let peer = Peer::empty(1);
-        let region_routes = vec![RegionRoute {
-            region: Region::new_test(region_id),
-            leader_peer: Some(peer.clone()),
-            ..Default::default()
-        }];
-        table_metadata_manager
-            .create_table_metadata(
-                task.table_info.clone(),
-                TableRouteValue::physical(region_routes.clone()),
-                HashMap::new(),
-            )
-            .await
-            .unwrap();
-        let table_route = cache.get(1024).await.unwrap().unwrap();
-        assert_eq!(
-            (*table_route)
-                .clone()
-                .as_physical_table_route_ref()
-                .region_routes,
-            region_routes
-        );
-
-        assert!(table_route_cache.contains_key(&1024));
-        assert!(cache.contains_key(&1024));
-        cache
-            .invalidate(&[CacheIdent::TableId(1024)])
-            .await
-            .unwrap();
-        assert!(!cache.contains_key(&1024));
-    }
-
-    #[tokio::test]
-    async fn test_cache_with_logical_table_route() {
-        let mem_kv = Arc::new(MemoryKvBackend::default());
-        let table_metadata_manager = TableMetadataManager::new(mem_kv.clone());
-        let cache = CacheBuilder::new(128).build();
-        let table_route_cache = Arc::new(new_table_route_cache(
-            "test".to_string(),
-            cache,
-            mem_kv.clone(),
-        ));
-        let cache = CacheBuilder::new(128).build();
-        let cache =
-            new_composite_table_route_cache("test".to_string(), cache, table_route_cache.clone());
-
-        let result = cache.get(1024).await.unwrap();
-        assert!(result.is_none());
-        // Prepares table routes
-        let task = test_create_table_task("my_table", 1024);
-        let table_id = 10;
-        let region_id = RegionId::new(table_id, 1);
-        let peer = Peer::empty(1);
-        let region_routes = vec![RegionRoute {
-            region: Region::new_test(region_id),
-            leader_peer: Some(peer.clone()),
-            ..Default::default()
-        }];
-        table_metadata_manager
-            .create_table_metadata(
-                task.table_info.clone(),
-                TableRouteValue::physical(region_routes.clone()),
-                HashMap::new(),
-            )
-            .await
-            .unwrap();
-        let mut task = test_create_logical_table_task("logical");
-        task.table_info.ident.table_id = 1025;
-        table_metadata_manager
-            .create_logical_tables_metadata(vec![(
-                task.table_info,
-                TableRouteValue::logical(1024, vec![RegionId::new(1025, 0)]),
-            )])
-            .await
-            .unwrap();
-
-        // Gets logical table route
-        let table_route = cache.get(1025).await.unwrap().unwrap();
-        assert_eq!(
-            table_route
-                .as_logical_table_route_ref()
-                .unwrap()
-                .physical_table_id(),
-            1024
-        );
-        assert_eq!(
-            table_route.as_physical_table_route_ref().region_routes,
-            region_routes
-        );
-
-        assert!(!cache.contains_key(&1024));
-        // Gets physical table route
-        let table_route = cache.get(1024).await.unwrap().unwrap();
-        assert_eq!(
-            table_route.as_physical_table_route_ref().region_routes,
-            region_routes
-        );
-        assert!(table_route.is_physical());
-
-        cache
-            .invalidate(&[CacheIdent::TableId(1025)])
-            .await
-            .unwrap();
-        assert!(!cache.contains_key(&1025));
-    }
-}

src/common/meta/src/metrics.rs

@@ -89,4 +89,11 @@ lazy_static! {
         &["name"]
     )
     .unwrap();
+    /// Cache container load cache timer
+    pub static ref CACHE_CONTAINER_LOAD_CACHE: HistogramVec = register_histogram_vec!(
+        "greptime_meta_cache_container_load_cache",
+        "cache container load cache",
+        &["name"]
+    )
+    .unwrap();
 }