diff --git a/src/catalog/src/kvbackend/manager.rs b/src/catalog/src/kvbackend/manager.rs index dc86c703d88..a4d8214a55b 100644 --- a/src/catalog/src/kvbackend/manager.rs +++ b/src/catalog/src/kvbackend/manager.rs @@ -29,7 +29,7 @@ use common_meta::key::TableMetadataManagerRef; use common_meta::key::catalog_name::CatalogNameKey; use common_meta::key::flow::FlowMetadataManager; use common_meta::key::schema_name::SchemaNameKey; -use common_meta::key::table_info::{TableInfoManager, TableInfoValue}; +use common_meta::key::table_info::TableInfoValue; use common_meta::key::table_name::TableNameKey; use common_meta::kv_backend::KvBackendRef; use common_procedure::ProcedureManagerRef; @@ -112,7 +112,7 @@ impl KvBackendCatalogManager { // Override logical table's partition key indices with physical table's. async fn override_logical_table_partition_key_indices( table_route_cache: &TableRouteCacheRef, - table_info_manager: &TableInfoManager, + table_info_cache: &TableInfoCacheRef, table: TableRef, ) -> Result { // If the table is not a metric table, return the table directly. @@ -125,7 +125,7 @@ impl KvBackendCatalogManager { .await .context(TableMetadataManagerSnafu)? && let TableRoute::Logical(logical_route) = &*table_route_value - && let Some(physical_table_info_value) = table_info_manager + && let Some(physical_table_info) = table_info_cache .get(logical_route.physical_table_id()) .await .context(TableMetadataManagerSnafu)? @@ -135,15 +135,13 @@ impl KvBackendCatalogManager { let mut phy_part_cols_not_in_logical_table = vec![]; // Remap partition key indices from physical table to logical table - new_table_info.meta.partition_key_indices = physical_table_info_value - .table_info + new_table_info.meta.partition_key_indices = physical_table_info .meta .partition_key_indices .iter() .filter_map(|&physical_index| { // Get the column name from the physical table using the physical index - physical_table_info_value - .table_info + physical_table_info .meta .schema .column_schemas() @@ -321,9 +319,13 @@ impl CatalogManager for KvBackendCatalogManager { self.cache_registry.get().context(CacheNotFoundSnafu { name: "table_route_cache", })?; + let table_info_cache: TableInfoCacheRef = + self.cache_registry.get().context(CacheNotFoundSnafu { + name: "table_info_cache", + })?; return Self::override_logical_table_partition_key_indices( &table_route_cache, - self.table_metadata_manager.table_info_manager(), + &table_info_cache, table, ) .await @@ -453,6 +455,10 @@ impl CatalogManager for KvBackendCatalogManager { self.cache_registry.get().context(CacheNotFoundSnafu { name: "table_route_cache", }); + let table_info_cache: Result = + self.cache_registry.get().context(CacheNotFoundSnafu { + name: "table_info_cache", + }); common_runtime::spawn_global(async move { let table_route_cache = match table_route_cache { @@ -462,6 +468,13 @@ impl CatalogManager for KvBackendCatalogManager { return; } }; + let table_info_cache = match table_info_cache { + Ok(table_info_cache) => table_info_cache, + Err(e) => { + let _ = tx.send(Err(e)).await; + return; + } + }; let table_id_stream = metadata_manager .table_name_manager() @@ -490,6 +503,7 @@ impl CatalogManager for KvBackendCatalogManager { let tx = tx.clone(); let semaphore = semaphore.clone(); let table_route_cache = table_route_cache.clone(); + let table_info_cache = table_info_cache.clone(); common_runtime::spawn_global(async move { // we don't explicitly close the semaphore so just ignore the potential error. let _ = semaphore.acquire().await; @@ -509,7 +523,7 @@ impl CatalogManager for KvBackendCatalogManager { for table in table_info_values.into_values().map(build_table) { let table = Self::override_logical_table_partition_key_indices( &table_route_cache, - metadata_manager.table_info_manager(), + &table_info_cache, table, ) .await; @@ -645,3 +659,86 @@ impl SystemCatalog { } } } + +#[cfg(test)] +mod tests { + use std::sync::atomic::{AtomicUsize, Ordering}; + + use common_meta::cache::{CacheContainer, Initializer}; + use common_meta::key::table_route::LogicalTableRouteValue; + use moka::future::Cache as FutureCache; + use table::metadata::TableInfo; + + use super::*; + + fn metric_table_info( + table_id: TableId, + table_name: &str, + partition_key_indices: Vec, + ) -> TableInfo { + let mut table_info = + common_meta::ddl::test_util::create_table::test_create_table_task(table_name, table_id) + .table_info; + table_info.meta.engine = METRIC_ENGINE_NAME.to_string(); + table_info.meta.partition_key_indices = partition_key_indices; + table_info + } + + #[tokio::test] + async fn test_logical_table_uses_cached_physical_table_info() { + const LOGICAL_TABLE_ID: TableId = 1; + const PHYSICAL_TABLE_ID: TableId = 2; + + let route = Arc::new(TableRoute::Logical(Arc::new(LogicalTableRouteValue::new( + PHYSICAL_TABLE_ID, + )))); + let route_initializer: Initializer> = Arc::new(move |_| { + let route = route.clone(); + Box::pin(async move { Ok(Some(route)) }) + }); + let table_route_cache = Arc::new(CacheContainer::new( + "test_table_route_cache".to_string(), + FutureCache::new(16), + Box::new(|_, _| Box::pin(async { Ok(()) })), + route_initializer, + |_| true, + )); + + let load_count = Arc::new(AtomicUsize::new(0)); + let physical_table_info = + Arc::new(metric_table_info(PHYSICAL_TABLE_ID, "physical", vec![1])); + let table_info_initializer: Initializer> = { + let load_count = load_count.clone(); + let table_info = physical_table_info.clone(); + Arc::new(move |_| { + load_count.fetch_add(1, Ordering::Relaxed); + let table_info = table_info.clone(); + Box::pin(async move { Ok(Some(table_info)) }) + }) + }; + let table_info_cache = Arc::new(CacheContainer::new( + "test_table_info_cache".to_string(), + FutureCache::new(16), + Box::new(|_, _| Box::pin(async { Ok(()) })), + table_info_initializer, + |_| true, + )); + + let logical_table = DistTable::table(Arc::new(metric_table_info( + LOGICAL_TABLE_ID, + "logical", + vec![], + ))); + for _ in 0..2 { + let table = KvBackendCatalogManager::override_logical_table_partition_key_indices( + &table_route_cache, + &table_info_cache, + logical_table.clone(), + ) + .await + .unwrap(); + assert_eq!(table.table_info().meta.partition_key_indices, vec![1]); + } + assert_eq!(load_count.load(Ordering::Relaxed), 1); + } +}