fix: schema cache invalidation (#5067)

* fix: use SchemaCache to locate database metadata

* main:
 Refactor SchemaMetadataManager to use TableInfoCacheRef

 - Replace TableInfoManagerRef with TableInfoCacheRef in SchemaMetadataManager
 - Update DatanodeBuilder to pass TableInfoCacheRef to SchemaMetadataManager
 - Rename error MissingCacheRegistrySnafu to MissingCacheSnafu in datanode module
 - Adjust tests to use new mock_schema_metadata_manager with TableInfoCacheRef

* fix/schema-cache-invalidation: Add cache module and integrate cache registry into datanode

 • Implement build_datanode_cache_registry function to create cache registry for datanode
 • Integrate cache registry into datanode by modifying DatanodeBuilder and HeartbeatTask
 • Refactor InvalidateTableCacheHandler to InvalidateCacheHandler and move to common-meta crate
 • Update Cargo.toml to include cache as a dev-dependency for datanode
 • Adjust related modules (flownode, frontend, tests-integration, standalone) to use new cache handler and registry
 • Remove obsolete handler module from frontend crate

* fix: fuzz imports

* chore: add some doc for cahce builder functions

* refactor: change table info cache to table schema cache

* fix: remove unused variants

* fix fuzz

* chore: apply suggestion

Co-authored-by: Weny Xu <wenymedia@gmail.com>

* chore: apply suggestion

Co-authored-by: Weny Xu <wenymedia@gmail.com>

* fix: compile

---------

Co-authored-by: dennis zhuang <killme2008@gmail.com>
Co-authored-by: Weny Xu <wenymedia@gmail.com>

src/cache/src/lib.rs

@@ -19,9 +19,9 @@ use std::time::Duration;
 
 use catalog::kvbackend::new_table_cache;
 use common_meta::cache::{
-    new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache,
-    new_table_route_cache, new_view_info_cache, CacheRegistry, CacheRegistryBuilder,
-    LayeredCacheRegistryBuilder,
+    new_schema_cache, new_table_flownode_set_cache, new_table_info_cache, new_table_name_cache,
+    new_table_route_cache, new_table_schema_cache, new_view_info_cache, CacheRegistry,
+    CacheRegistryBuilder, LayeredCacheRegistryBuilder,
 };
 use common_meta::kv_backend::KvBackendRef;
 use moka::future::CacheBuilder;
@@ -37,9 +37,47 @@ pub const TABLE_INFO_CACHE_NAME: &str = "table_info_cache";
 pub const VIEW_INFO_CACHE_NAME: &str = "view_info_cache";
 pub const TABLE_NAME_CACHE_NAME: &str = "table_name_cache";
 pub const TABLE_CACHE_NAME: &str = "table_cache";
+pub const SCHEMA_CACHE_NAME: &str = "schema_cache";
+pub const TABLE_SCHEMA_NAME_CACHE_NAME: &str = "table_schema_name_cache";
 pub const TABLE_FLOWNODE_SET_CACHE_NAME: &str = "table_flownode_set_cache";
 pub const TABLE_ROUTE_CACHE_NAME: &str = "table_route_cache";
 
+/// Builds cache registry for datanode, including:
+/// - Schema cache.
+/// - Table id to schema name cache.
+pub fn build_datanode_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry {
+    // Builds table id schema name cache that never expires.
+    let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build();
+    let table_id_schema_cache = Arc::new(new_table_schema_cache(
+        TABLE_SCHEMA_NAME_CACHE_NAME.to_string(),
+        cache,
+        kv_backend.clone(),
+    ));
+
+    // Builds schema cache
+    let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY)
+        .time_to_live(DEFAULT_CACHE_TTL)
+        .time_to_idle(DEFAULT_CACHE_TTI)
+        .build();
+    let schema_cache = Arc::new(new_schema_cache(
+        SCHEMA_CACHE_NAME.to_string(),
+        cache,
+        kv_backend.clone(),
+    ));
+
+    CacheRegistryBuilder::default()
+        .add_cache(table_id_schema_cache)
+        .add_cache(schema_cache)
+        .build()
+}
+
+/// Builds cache registry for frontend and datanode, including:
+/// - Table info cache
+/// - Table name cache
+/// - Table route cache
+/// - Table flow node cache
+/// - View cache
+/// - Schema cache
 pub fn build_fundamental_cache_registry(kv_backend: KvBackendRef) -> CacheRegistry {
     // Builds table info cache
     let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY)
@@ -95,12 +133,30 @@ pub fn build_fundamental_cache_registry(kv_backend: KvBackendRef) -> CacheRegist
         kv_backend.clone(),
     ));
 
+    // Builds schema cache
+    let cache = CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY)
+        .time_to_live(DEFAULT_CACHE_TTL)
+        .time_to_idle(DEFAULT_CACHE_TTI)
+        .build();
+    let schema_cache = Arc::new(new_schema_cache(
+        SCHEMA_CACHE_NAME.to_string(),
+        cache,
+        kv_backend.clone(),
+    ));
+
+    let table_id_schema_cache = Arc::new(new_table_schema_cache(
+        TABLE_SCHEMA_NAME_CACHE_NAME.to_string(),
+        CacheBuilder::new(DEFAULT_CACHE_MAX_CAPACITY).build(),
+        kv_backend,
+    ));
     CacheRegistryBuilder::default()
         .add_cache(table_info_cache)
         .add_cache(table_name_cache)
         .add_cache(table_route_cache)
         .add_cache(view_info_cache)
         .add_cache(table_flownode_set_cache)
+        .add_cache(schema_cache)
+        .add_cache(table_id_schema_cache)
         .build()
 }
 

src/cmd/src/datanode.rs

@@ -16,10 +16,12 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
+use cache::build_datanode_cache_registry;
 use catalog::kvbackend::MetaKvBackend;
 use clap::Parser;
 use common_base::Plugins;
 use common_config::Configurable;
+use common_meta::cache::LayeredCacheRegistryBuilder;
 use common_telemetry::logging::TracingOptions;
 use common_telemetry::{info, warn};
 use common_version::{short_version, version};
@@ -300,9 +302,17 @@ impl StartCommand {
             client: meta_client.clone(),
         });
 
+        // Builds cache registry for datanode.
+        let layered_cache_registry = Arc::new(
+            LayeredCacheRegistryBuilder::default()
+                .add_cache_registry(build_datanode_cache_registry(meta_backend.clone()))
+                .build(),
+        );
+
         let mut datanode = DatanodeBuilder::new(opts.clone(), plugins)
             .with_meta_client(meta_client)
             .with_kv_backend(meta_backend)
+            .with_cache_registry(layered_cache_registry)
             .build()
             .await
             .context(StartDatanodeSnafu)?;

src/cmd/src/flownode.rs

@@ -22,6 +22,7 @@ use common_base::Plugins;
 use common_config::Configurable;
 use common_grpc::channel_manager::ChannelConfig;
 use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder};
+use common_meta::heartbeat::handler::invalidate_table_cache::InvalidateCacheHandler;
 use common_meta::heartbeat::handler::parse_mailbox_message::ParseMailboxMessageHandler;
 use common_meta::heartbeat::handler::HandlerGroupExecutor;
 use common_meta::key::flow::FlowMetadataManager;
@@ -30,7 +31,6 @@ use common_telemetry::info;
 use common_telemetry::logging::TracingOptions;
 use common_version::{short_version, version};
 use flow::{FlownodeBuilder, FlownodeInstance, FrontendInvoker};
-use frontend::heartbeat::handler::invalidate_table_cache::InvalidateTableCacheHandler;
 use meta_client::{MetaClientOptions, MetaClientType};
 use servers::Mode;
 use snafu::{OptionExt, ResultExt};
@@ -288,9 +288,7 @@ impl StartCommand {
 
         let executor = HandlerGroupExecutor::new(vec![
             Arc::new(ParseMailboxMessageHandler),
-            Arc::new(InvalidateTableCacheHandler::new(
-                layered_cache_registry.clone(),
-            )),
+            Arc::new(InvalidateCacheHandler::new(layered_cache_registry.clone())),
         ]);
 
         let heartbeat_task = flow::heartbeat::HeartbeatTask::new(

src/cmd/src/frontend.rs

@@ -24,13 +24,13 @@ use common_base::Plugins;
 use common_config::Configurable;
 use common_grpc::channel_manager::ChannelConfig;
 use common_meta::cache::{CacheRegistryBuilder, LayeredCacheRegistryBuilder};
+use common_meta::heartbeat::handler::invalidate_table_cache::InvalidateCacheHandler;
 use common_meta::heartbeat::handler::parse_mailbox_message::ParseMailboxMessageHandler;
 use common_meta::heartbeat::handler::HandlerGroupExecutor;
 use common_telemetry::info;
 use common_telemetry::logging::TracingOptions;
 use common_time::timezone::set_default_timezone;
 use common_version::{short_version, version};
-use frontend::heartbeat::handler::invalidate_table_cache::InvalidateTableCacheHandler;
 use frontend::heartbeat::HeartbeatTask;
 use frontend::instance::builder::FrontendBuilder;
 use frontend::instance::{FrontendInstance, Instance as FeInstance};
@@ -328,9 +328,7 @@ impl StartCommand {
 
         let executor = HandlerGroupExecutor::new(vec![
             Arc::new(ParseMailboxMessageHandler),
-            Arc::new(InvalidateTableCacheHandler::new(
-                layered_cache_registry.clone(),
-            )),
+            Arc::new(InvalidateCacheHandler::new(layered_cache_registry.clone())),
         ]);
 
         let heartbeat_task = HeartbeatTask::new(

src/cmd/src/standalone.rs

@@ -497,6 +497,7 @@ impl StartCommand {
 
         let datanode = DatanodeBuilder::new(dn_opts, plugins.clone())
             .with_kv_backend(kv_backend.clone())
+            .with_cache_registry(layered_cache_registry.clone())
             .build()
             .await
             .context(StartDatanodeSnafu)?;

src/common/meta/src/cache.rs

@@ -24,7 +24,8 @@ pub use registry::{
     LayeredCacheRegistryBuilder, LayeredCacheRegistryRef,
 };
 pub use table::{
-    new_table_info_cache, new_table_name_cache, new_table_route_cache, new_view_info_cache,
-    TableInfoCache, TableInfoCacheRef, TableNameCache, TableNameCacheRef, TableRoute,
-    TableRouteCache, TableRouteCacheRef, ViewInfoCache, ViewInfoCacheRef,
+    new_schema_cache, new_table_info_cache, new_table_name_cache, new_table_route_cache,
+    new_table_schema_cache, new_view_info_cache, SchemaCache, SchemaCacheRef, TableInfoCache,
+    TableInfoCacheRef, TableNameCache, TableNameCacheRef, TableRoute, TableRouteCache,
+    TableRouteCacheRef, TableSchemaCache, TableSchemaCacheRef, ViewInfoCache, ViewInfoCacheRef,
 };

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

@@ -12,12 +12,16 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+mod schema;
 mod table_info;
 mod table_name;
 mod table_route;
+mod table_schema;
 mod view_info;
 
+pub use schema::{new_schema_cache, SchemaCache, SchemaCacheRef};
 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};
+pub use table_schema::{new_table_schema_cache, TableSchemaCache, TableSchemaCacheRef};
 pub use view_info::{new_view_info_cache, ViewInfoCache, ViewInfoCacheRef};

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

@@ -0,0 +1,73 @@
+// 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_util::future::BoxFuture;
+use moka::future::Cache;
+use snafu::OptionExt;
+
+use crate::cache::{CacheContainer, Initializer};
+use crate::error::ValueNotExistSnafu;
+use crate::instruction::CacheIdent;
+use crate::key::schema_name::{SchemaManager, SchemaName, SchemaNameKey, SchemaNameValue};
+use crate::kv_backend::KvBackendRef;
+
+pub type SchemaCache = CacheContainer<SchemaName, Arc<SchemaNameValue>, CacheIdent>;
+pub type SchemaCacheRef = Arc<SchemaCache>;
+
+/// Constructs a [SchemaCache].
+pub fn new_schema_cache(
+    name: String,
+    cache: Cache<SchemaName, Arc<SchemaNameValue>>,
+    kv_backend: KvBackendRef,
+) -> SchemaCache {
+    let schema_manager = SchemaManager::new(kv_backend.clone());
+    let init = init_factory(schema_manager);
+
+    CacheContainer::new(name, cache, Box::new(invalidator), init, Box::new(filter))
+}
+
+fn init_factory(schema_manager: SchemaManager) -> Initializer<SchemaName, Arc<SchemaNameValue>> {
+    Arc::new(move |schema_name| {
+        let manager = schema_manager.clone();
+        Box::pin(async move {
+            let schema_value = manager
+                .get(SchemaNameKey {
+                    catalog: &schema_name.catalog_name,
+                    schema: &schema_name.schema_name,
+                })
+                .await?
+                .context(ValueNotExistSnafu)?
+                .into_inner();
+            Ok(Some(Arc::new(schema_value)))
+        })
+    })
+}
+
+fn invalidator<'a>(
+    cache: &'a Cache<SchemaName, Arc<SchemaNameValue>>,
+    ident: &'a CacheIdent,
+) -> BoxFuture<'a, crate::error::Result<()>> {
+    Box::pin(async move {
+        if let CacheIdent::SchemaName(schema_name) = ident {
+            cache.invalidate(schema_name).await
+        }
+        Ok(())
+    })
+}
+
+fn filter(ident: &CacheIdent) -> bool {
+    matches!(ident, CacheIdent::SchemaName(_))
+}

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

@@ -0,0 +1,76 @@
+// 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.
+
+//! Cache for table id to schema name mapping.
+
+use std::sync::Arc;
+
+use futures_util::future::BoxFuture;
+use moka::future::Cache;
+use snafu::OptionExt;
+use store_api::storage::TableId;
+
+use crate::cache::{CacheContainer, Initializer};
+use crate::error;
+use crate::instruction::CacheIdent;
+use crate::key::schema_name::SchemaName;
+use crate::key::table_info::TableInfoManager;
+use crate::kv_backend::KvBackendRef;
+
+pub type TableSchemaCache = CacheContainer<TableId, Arc<SchemaName>, CacheIdent>;
+pub type TableSchemaCacheRef = Arc<TableSchemaCache>;
+
+/// Constructs a [TableSchemaCache].
+pub fn new_table_schema_cache(
+    name: String,
+    cache: Cache<TableId, Arc<SchemaName>>,
+    kv_backend: KvBackendRef,
+) -> TableSchemaCache {
+    let table_info_manager = TableInfoManager::new(kv_backend);
+    let init = init_factory(table_info_manager);
+
+    CacheContainer::new(name, cache, Box::new(invalidator), init, Box::new(filter))
+}
+
+fn init_factory(table_info_manager: TableInfoManager) -> Initializer<TableId, Arc<SchemaName>> {
+    Arc::new(move |table_id| {
+        let table_info_manager = table_info_manager.clone();
+        Box::pin(async move {
+            let raw_table_info = table_info_manager
+                .get(*table_id)
+                .await?
+                .context(error::ValueNotExistSnafu)?
+                .into_inner()
+                .table_info;
+
+            Ok(Some(Arc::new(SchemaName {
+                catalog_name: raw_table_info.catalog_name,
+                schema_name: raw_table_info.schema_name,
+            })))
+        })
+    })
+}
+
+/// Never invalidates table id schema cache.
+fn invalidator<'a>(
+    _cache: &'a Cache<TableId, Arc<SchemaName>>,
+    _ident: &'a CacheIdent,
+) -> BoxFuture<'a, error::Result<()>> {
+    Box::pin(std::future::ready(Ok(())))
+}
+
+/// Never invalidates table id schema cache.
+fn filter(_ident: &CacheIdent) -> bool {
+    false
+}

src/common/meta/src/heartbeat/handler.rs

@@ -21,6 +21,7 @@ use common_telemetry::error;
 use crate::error::Result;
 use crate::heartbeat::mailbox::{IncomingMessage, MailboxRef};
 
+pub mod invalidate_table_cache;
 pub mod parse_mailbox_message;
 #[cfg(test)]
 mod tests;