From 3bd915865a4d5db7f40a9e56cefafd169e160875 Mon Sep 17 00:00:00 2001 From: Weny Xu Date: Thu, 4 Jun 2026 14:10:38 +0800 Subject: [PATCH] fix: improve remote WAL replay checkpoint handling (#8225) * fix: correct topic_latest_entry_id Signed-off-by: WenyXu * fix: use pruned wal id for replay checkpoint Signed-off-by: WenyXu * feat: persist remote wal checkpoints periodically Signed-off-by: WenyXu * test: cover remote wal topic latest entry id Signed-off-by: WenyXu * chore: apply suggestions Signed-off-by: WenyXu * fix: merge pruned wal id for migration checkpoint Signed-off-by: WenyXu --------- Signed-off-by: WenyXu --- src/common/meta/src/key/topic_name.rs | 34 +- src/common/meta/src/key/topic_region.rs | 80 ++++ src/datanode/src/utils.rs | 77 +++- .../src/procedure/region_migration.rs | 43 ++- .../upgrade_candidate_region.rs | 135 ++++++- src/meta-srv/src/region/flush_trigger.rs | 344 +++++++++++++++--- src/mito2/src/engine/flush_test.rs | 239 +++++++++++- src/mito2/src/region/opener.rs | 23 +- 8 files changed, 886 insertions(+), 89 deletions(-) diff --git a/src/common/meta/src/key/topic_name.rs b/src/common/meta/src/key/topic_name.rs index 99ae631a72..79454c7cd2 100644 --- a/src/common/meta/src/key/topic_name.rs +++ b/src/common/meta/src/key/topic_name.rs @@ -12,6 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::collections::HashMap; use std::fmt::{self, Display}; use serde::{Deserialize, Serialize}; @@ -27,7 +28,7 @@ use crate::key::{ use crate::kv_backend::KvBackendRef; use crate::kv_backend::txn::{Txn, TxnOp}; use crate::rpc::KeyValue; -use crate::rpc::store::{BatchPutRequest, RangeRequest}; +use crate::rpc::store::{BatchGetRequest, BatchPutRequest, RangeRequest}; #[derive(Debug, Clone, PartialEq)] pub struct TopicNameKey<'a> { @@ -205,6 +206,25 @@ impl TopicNameManager { .transpose() } + /// Batch get values for specific topics. + pub async fn batch_get( + &self, + topics: Vec>, + ) -> Result> { + let raw_keys = topics.iter().map(|key| key.to_bytes()).collect::>(); + let req = BatchGetRequest { keys: raw_keys }; + let resp = self.kv_backend.batch_get(req).await?; + + resp.kvs + .into_iter() + .map(|kv| { + let key = TopicNameKey::from_bytes(&kv.key)?; + let value = TopicNameValue::try_from_raw_value(&kv.value)?; + Ok((key.topic.to_string(), value)) + }) + .collect::>>() + } + /// Update the topic name key and value in the kv backend. pub async fn update( &self, @@ -295,5 +315,17 @@ mod tests { let err = manager.update(topic, 3, Some(value)).await.unwrap_err(); assert_matches!(err, error::Error::Unexpected { .. }); } + + let batch_topics = topics + .iter() + .take(2) + .map(|topic| TopicNameKey::new(topic)) + .chain(std::iter::once(TopicNameKey::new("missing-topic"))) + .collect::>(); + let values = manager.batch_get(batch_topics).await.unwrap(); + assert_eq!(values.len(), 2); + for topic in topics.iter().take(2) { + assert_eq!(values.get(topic).unwrap().pruned_entry_id, 1); + } } } diff --git a/src/common/meta/src/key/topic_region.rs b/src/common/meta/src/key/topic_region.rs index 9feac10804..e7db23cfc2 100644 --- a/src/common/meta/src/key/topic_region.rs +++ b/src/common/meta/src/key/topic_region.rs @@ -126,6 +126,34 @@ impl ReplayCheckpoint { metadata_entry_id, } } + + /// Merges the checkpoint with the topic pruned entry id. + pub fn merge_with_topic_pruned_entry_id( + checkpoint: Option, + pruned_entry_id: Option, + is_metric_engine: bool, + ) -> Option { + match (checkpoint, pruned_entry_id) { + (Some(checkpoint), Some(pruned_entry_id)) => Some(Self { + entry_id: checkpoint.entry_id.max(pruned_entry_id), + metadata_entry_id: if is_metric_engine { + Some( + checkpoint + .metadata_entry_id + .unwrap_or_default() + .max(pruned_entry_id), + ) + } else { + checkpoint.metadata_entry_id + }, + }), + (None, Some(pruned_entry_id)) => Some(Self { + entry_id: pruned_entry_id, + metadata_entry_id: is_metric_engine.then_some(pruned_entry_id), + }), + (checkpoint, None) => checkpoint, + } + } } impl TopicRegionValue { @@ -369,6 +397,58 @@ mod tests { use super::*; use crate::kv_backend::memory::MemoryKvBackend; + #[test] + fn test_merge_checkpoint_with_topic_pruned_entry_id_missing_pruned() { + let checkpoint = Some(ReplayCheckpoint::new(10, None)); + + assert_eq!( + ReplayCheckpoint::merge_with_topic_pruned_entry_id(checkpoint, None, true), + checkpoint + ); + } + + #[test] + fn test_merge_checkpoint_with_topic_pruned_entry_id_creates_checkpoint() { + assert_eq!( + ReplayCheckpoint::merge_with_topic_pruned_entry_id(None, Some(10), true), + Some(ReplayCheckpoint::new(10, Some(10))) + ); + } + + #[test] + fn test_merge_checkpoint_with_topic_pruned_entry_id_updates_both_ids() { + let checkpoint = ReplayCheckpoint::new(10, Some(5)); + + assert_eq!( + ReplayCheckpoint::merge_with_topic_pruned_entry_id(Some(checkpoint), Some(20), true), + Some(ReplayCheckpoint::new(20, Some(20))) + ); + } + + #[test] + fn test_merge_checkpoint_with_topic_pruned_entry_id_preserves_larger_ids() { + let checkpoint = ReplayCheckpoint::new(30, Some(40)); + + assert_eq!( + ReplayCheckpoint::merge_with_topic_pruned_entry_id(Some(checkpoint), Some(20), true), + Some(checkpoint) + ); + } + + #[test] + fn test_merge_checkpoint_with_topic_pruned_entry_id_for_mito() { + assert_eq!( + ReplayCheckpoint::merge_with_topic_pruned_entry_id(None, Some(10), false), + Some(ReplayCheckpoint::new(10, None)) + ); + + let checkpoint = ReplayCheckpoint::new(5, Some(8)); + assert_eq!( + ReplayCheckpoint::merge_with_topic_pruned_entry_id(Some(checkpoint), Some(10), false), + Some(ReplayCheckpoint::new(10, Some(8))) + ); + } + #[tokio::test] async fn test_topic_region_manager() { let kv_backend = Arc::new(MemoryKvBackend::default()); diff --git a/src/datanode/src/utils.rs b/src/datanode/src/utils.rs index 816ae021ba..88f4f1aa60 100644 --- a/src/datanode/src/utils.rs +++ b/src/datanode/src/utils.rs @@ -16,11 +16,16 @@ use std::collections::HashMap; use common_meta::DatanodeId; use common_meta::key::datanode_table::DatanodeTableManager; -use common_meta::key::topic_region::{TopicRegionKey, TopicRegionManager, TopicRegionValue}; +use common_meta::key::topic_name::{TopicNameKey, TopicNameManager, TopicNameValue}; +use common_meta::key::topic_region::{ + ReplayCheckpoint as MetadataReplayCheckpoint, TopicRegionKey, TopicRegionManager, + TopicRegionValue, +}; use common_meta::kv_backend::KvBackendRef; use common_meta::wal_provider::{extract_topic_from_wal_options, prepare_wal_options}; use futures::TryStreamExt; use snafu::ResultExt; +use store_api::metric_engine_consts::METRIC_ENGINE_NAME; use store_api::path_utils::table_dir; use store_api::region_request::{PathType, RegionOpenRequest, ReplayCheckpoint}; use store_api::storage::{RegionId, RegionNumber}; @@ -63,14 +68,45 @@ fn group_region_by_topic( } } +fn region_pruned_entry_ids( + topic_regions: &HashMap>, + topic_name_values: &HashMap, +) -> HashMap { + topic_regions + .iter() + .flat_map(|(topic, region_ids)| { + topic_name_values + .get(topic) + .into_iter() + .flat_map(move |value| { + region_ids + .iter() + .map(move |region_id| (*region_id, value.pruned_entry_id)) + }) + }) + .collect() +} + fn get_replay_checkpoint( region_id: RegionId, topic_region_values: &Option>, + pruned_entry_id: Option, + is_metric_engine: bool, ) -> Option { - let topic_region_values = topic_region_values.as_ref()?; - let topic_region_value = topic_region_values.get(®ion_id); - let replay_checkpoint = topic_region_value.and_then(|value| value.checkpoint); - replay_checkpoint.map(|checkpoint| ReplayCheckpoint { + let checkpoint = topic_region_values + .as_ref() + .and_then(|values| values.get(®ion_id)) + .and_then(|value| value.checkpoint) + .map(|checkpoint| { + MetadataReplayCheckpoint::new(checkpoint.entry_id, checkpoint.metadata_entry_id) + }); + + MetadataReplayCheckpoint::merge_with_topic_pruned_entry_id( + checkpoint, + pruned_entry_id, + is_metric_engine, + ) + .map(|checkpoint| ReplayCheckpoint { entry_id: checkpoint.entry_id, metadata_entry_id: checkpoint.metadata_entry_id, }) @@ -88,7 +124,8 @@ pub async fn build_region_open_requests( .await .context(GetMetadataSnafu)?; - let topic_region_manager = TopicRegionManager::new(kv_backend); + let topic_region_manager = TopicRegionManager::new(kv_backend.clone()); + let topic_name_manager = TopicNameManager::new(kv_backend); let mut topic_regions = HashMap::>::new(); let mut regions = vec![]; #[cfg(feature = "enterprise")] @@ -161,10 +198,36 @@ pub async fn build_region_open_requests( None }; + let topic_name_values = if !topic_regions.is_empty() { + let topics = topic_regions + .keys() + .map(|topic| TopicNameKey::new(topic)) + .collect::>(); + Some( + topic_name_manager + .batch_get(topics) + .await + .context(GetMetadataSnafu)?, + ) + } else { + None + }; + let region_pruned_entry_ids = topic_name_values + .as_ref() + .map(|values| region_pruned_entry_ids(&topic_regions, values)); + let mut leader_region_requests = Vec::with_capacity(regions.len()); for (region_id, engine, store_path, options) in regions { let table_dir = table_dir(&store_path, region_id.table_id()); - let checkpoint = get_replay_checkpoint(region_id, &topic_region_values); + let pruned_entry_id = region_pruned_entry_ids + .as_ref() + .and_then(|values| values.get(®ion_id).copied()); + let checkpoint = get_replay_checkpoint( + region_id, + &topic_region_values, + pruned_entry_id, + engine == METRIC_ENGINE_NAME, + ); info!("region_id: {}, checkpoint: {:?}", region_id, checkpoint); leader_region_requests.push(( region_id, diff --git a/src/meta-srv/src/procedure/region_migration.rs b/src/meta-srv/src/procedure/region_migration.rs index 6563da63bc..c0ab352fc6 100644 --- a/src/meta-srv/src/procedure/region_migration.rs +++ b/src/meta-srv/src/procedure/region_migration.rs @@ -39,6 +39,7 @@ use common_meta::ddl::RegionFailureDetectorControllerRef; use common_meta::instruction::CacheIdent; use common_meta::key::datanode_table::{DatanodeTableKey, DatanodeTableValue}; use common_meta::key::table_route::TableRouteValue; +use common_meta::key::topic_name::TopicNameKey; use common_meta::key::topic_region::{ReplayCheckpoint, TopicRegionKey}; use common_meta::key::{DeserializedValueWithBytes, TableMetadataManagerRef}; use common_meta::kv_backend::{KvBackendRef, ResettableKvBackendRef}; @@ -661,11 +662,15 @@ impl Context { .await; } - /// Fetches the replay checkpoints for the given topic region keys. - pub async fn get_replay_checkpoints( + /// Fetches replay checkpoints and merges them with topic pruned entry ids. + pub async fn get_replay_checkpoints_with_topic_pruned_entry_ids( &self, - topic_region_keys: Vec>, + region_topics: &[(RegionId, String, bool)], ) -> Result> { + let topic_region_keys = region_topics + .iter() + .map(|(region_id, topic, _)| TopicRegionKey::new(*region_id, topic)) + .collect::>(); let topic_region_values = self .table_metadata_manager .topic_region_manager() @@ -673,9 +678,37 @@ impl Context { .await .context(error::TableMetadataManagerSnafu)?; - let replay_checkpoints = topic_region_values + let topic_name_keys = region_topics + .iter() + .map(|(_, topic, _)| topic.as_str()) + .collect::>() .into_iter() - .flat_map(|(key, value)| value.checkpoint.map(|value| (key, value))) + .map(TopicNameKey::new) + .collect::>(); + let topic_name_values = self + .table_metadata_manager + .topic_name_manager() + .batch_get(topic_name_keys) + .await + .context(error::TableMetadataManagerSnafu)?; + + let replay_checkpoints = region_topics + .iter() + .filter_map(|(region_id, topic, is_metric_engine)| { + let checkpoint = topic_region_values + .get(region_id) + .and_then(|value| value.checkpoint); + let pruned_entry_id = topic_name_values + .get(topic) + .map(|value| value.pruned_entry_id); + + ReplayCheckpoint::merge_with_topic_pruned_entry_id( + checkpoint, + pruned_entry_id, + *is_metric_engine, + ) + .map(|checkpoint| (*region_id, checkpoint)) + }) .collect::>(); Ok(replay_checkpoints) diff --git a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs index 4c60215cf7..2ffd64bb47 100644 --- a/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs +++ b/src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs @@ -21,7 +21,6 @@ use common_meta::ddl::utils::parse_region_wal_options; use common_meta::instruction::{ Instruction, InstructionReply, UpgradeRegion, UpgradeRegionReply, UpgradeRegionsReply, }; -use common_meta::key::topic_region::TopicRegionKey; use common_meta::lock_key::RemoteWalLock; use common_meta::wal_provider::extract_topic_from_wal_options; use common_procedure::{Context as ProcedureContext, Status}; @@ -30,6 +29,7 @@ use common_telemetry::{error, info}; use common_wal::options::WalOptions; use serde::{Deserialize, Serialize}; use snafu::{OptionExt, ResultExt, ensure}; +use store_api::metric_engine_consts::METRIC_ENGINE_NAME; use tokio::time::{Instant, sleep}; use crate::error::{self, Result}; @@ -133,17 +133,14 @@ impl UpgradeCandidateRegion { &datanode_table_value.region_info.region_wal_options, ) { - region_topic.push((*region_id, topic)); + let is_metric_engine = + datanode_table_value.region_info.engine == METRIC_ENGINE_NAME; + region_topic.push((*region_id, topic, is_metric_engine)); } } let replay_checkpoints = ctx - .get_replay_checkpoints( - region_topic - .iter() - .map(|(region_id, topic)| TopicRegionKey::new(*region_id, topic)) - .collect(), - ) + .get_replay_checkpoints_with_topic_pruned_entry_ids(®ion_topic) .await?; // Build upgrade regions instruction. let mut upgrade_regions = Vec::with_capacity(region_ids.len()); @@ -358,8 +355,11 @@ mod tests { use common_meta::key::table_route::TableRouteValue; use common_meta::key::test_utils::new_test_table_info; + use common_meta::key::topic_name::TopicNameKey; + use common_meta::key::topic_region::{ReplayCheckpoint, TopicRegionKey, TopicRegionValue}; use common_meta::peer::Peer; use common_meta::rpc::router::{Region, RegionRoute}; + use common_wal::options::KafkaWalOptions; use store_api::storage::RegionId; use super::*; @@ -382,9 +382,28 @@ mod tests { ) } + fn kafka_wal_options(topic: &str) -> HashMap { + HashMap::from([( + 1, + serde_json::to_string(&WalOptions::Kafka(KafkaWalOptions { + topic: topic.to_string(), + })) + .unwrap(), + )]) + } + async fn prepare_table_metadata(ctx: &Context, wal_options: HashMap) { + prepare_table_metadata_with_engine(ctx, wal_options, "engine").await; + } + + async fn prepare_table_metadata_with_engine( + ctx: &Context, + wal_options: HashMap, + engine: &str, + ) { let region_id = ctx.persistent_ctx.region_ids[0]; - let table_info = new_test_table_info(region_id.table_id()); + let mut table_info = new_test_table_info(region_id.table_id()); + table_info.meta.engine = engine.to_string(); let region_routes = vec![RegionRoute { region: Region::new_test(region_id), leader_peer: Some(ctx.persistent_ctx.from_peer.clone()), @@ -401,6 +420,104 @@ mod tests { .unwrap(); } + #[tokio::test] + async fn test_build_upgrade_region_instruction_merges_topic_pruned_entry_id() { + let state = UpgradeCandidateRegion::default(); + let persistent_context = new_persistent_context(); + let env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let region_id = ctx.persistent_ctx.region_ids[0]; + let topic = "test_topic"; + prepare_table_metadata(&ctx, kafka_wal_options(topic)).await; + ctx.table_metadata_manager + .topic_region_manager() + .batch_put(&[( + TopicRegionKey::new(region_id, topic), + Some(TopicRegionValue::new(Some(ReplayCheckpoint::new(10, None)))), + )]) + .await + .unwrap(); + ctx.table_metadata_manager + .topic_name_manager() + .batch_put(vec![TopicNameKey::new(topic)]) + .await + .unwrap(); + let prev = ctx + .table_metadata_manager + .topic_name_manager() + .get(topic) + .await + .unwrap(); + ctx.table_metadata_manager + .topic_name_manager() + .update(topic, 20, prev) + .await + .unwrap(); + + let instruction = state + .build_upgrade_region_instruction(&mut ctx, Duration::from_secs(1)) + .await + .unwrap(); + let Instruction::UpgradeRegions(upgrade_regions) = instruction else { + unreachable!() + }; + + assert_eq!(upgrade_regions.len(), 1); + assert_eq!(upgrade_regions[0].replay_entry_id, Some(20)); + assert_eq!(upgrade_regions[0].metadata_replay_entry_id, None); + } + + #[tokio::test] + async fn test_build_upgrade_region_instruction_merges_metric_metadata_pruned_entry_id() { + let state = UpgradeCandidateRegion::default(); + let persistent_context = new_persistent_context(); + let env = TestingEnv::new(); + let mut ctx = env.context_factory().new_context(persistent_context); + let region_id = ctx.persistent_ctx.region_ids[0]; + let topic = "test_topic"; + prepare_table_metadata_with_engine(&ctx, kafka_wal_options(topic), METRIC_ENGINE_NAME) + .await; + ctx.table_metadata_manager + .topic_region_manager() + .batch_put(&[( + TopicRegionKey::new(region_id, topic), + Some(TopicRegionValue::new(Some(ReplayCheckpoint::new( + 10, + Some(5), + )))), + )]) + .await + .unwrap(); + ctx.table_metadata_manager + .topic_name_manager() + .batch_put(vec![TopicNameKey::new(topic)]) + .await + .unwrap(); + let prev = ctx + .table_metadata_manager + .topic_name_manager() + .get(topic) + .await + .unwrap(); + ctx.table_metadata_manager + .topic_name_manager() + .update(topic, 20, prev) + .await + .unwrap(); + + let instruction = state + .build_upgrade_region_instruction(&mut ctx, Duration::from_secs(1)) + .await + .unwrap(); + let Instruction::UpgradeRegions(upgrade_regions) = instruction else { + unreachable!() + }; + + assert_eq!(upgrade_regions.len(), 1); + assert_eq!(upgrade_regions[0].replay_entry_id, Some(20)); + assert_eq!(upgrade_regions[0].metadata_replay_entry_id, Some(20)); + } + #[tokio::test] async fn test_datanode_is_unreachable() { let state = UpgradeCandidateRegion::default(); diff --git a/src/meta-srv/src/region/flush_trigger.rs b/src/meta-srv/src/region/flush_trigger.rs index cdaa4339f4..d13fb9f244 100644 --- a/src/meta-srv/src/region/flush_trigger.rs +++ b/src/meta-srv/src/region/flush_trigger.rs @@ -45,6 +45,9 @@ const TICKER_INTERVAL: Duration = Duration::from_secs(60); /// The duration of the recent period. const RECENT_DURATION: Duration = Duration::from_secs(300); +/// The interval to periodically persist region checkpoints regardless of replay size. +const PERIODIC_CHECKPOINT_PERSIST_INTERVAL: Duration = Duration::from_secs(60 * 60); + /// [`Event`] represents various types of events that can be processed by the region flush ticker. /// /// Variants: @@ -84,6 +87,8 @@ pub struct RegionFlushTrigger { flush_trigger_size: ReadableSize, /// The checkpoint trigger size. checkpoint_trigger_size: ReadableSize, + /// The last timestamp in milliseconds when a region checkpoint was persisted. + last_checkpoint_persist_millis_by_region: HashMap, /// The receiver of events. receiver: Receiver, } @@ -123,6 +128,7 @@ impl RegionFlushTrigger { server_addr, flush_trigger_size, checkpoint_trigger_size, + last_checkpoint_persist_millis_by_region: HashMap::new(), receiver: rx, }; (region_flush_trigger, region_flush_ticker) @@ -147,14 +153,15 @@ impl RegionFlushTrigger { } } - async fn handle_tick(&self) { + async fn handle_tick(&mut self) { if let Err(e) = self.trigger_flush().await { error!(e; "Failed to trigger flush"); } } - async fn trigger_flush(&self) -> Result<()> { + async fn trigger_flush(&mut self) -> Result<()> { let now = Instant::now(); + let now_millis = current_time_millis(); let topics = self .table_metadata_manager .topic_name_manager() @@ -162,17 +169,19 @@ impl RegionFlushTrigger { .await .context(error::TableMetadataManagerSnafu)?; + let mut active_region_ids = HashSet::new(); for topic in &topics { - let Some((latest_entry_id, avg_record_size)) = self.retrieve_topic_stat(topic) else { - continue; - }; if let Err(e) = self - .flush_regions_in_topic(topic, latest_entry_id, avg_record_size) + .handle_topic(topic, now_millis, &mut active_region_ids) .await { - error!(e; "Failed to flush regions in topic: {}", topic); + error!(e; "Failed to handle regions in topic: {}", topic); } } + retain_checkpoint_persist_records( + &mut self.last_checkpoint_persist_millis_by_region, + &active_region_ids, + ); debug!( "Triggered flush for {} topics in {:?}", @@ -182,6 +191,79 @@ impl RegionFlushTrigger { Ok(()) } + async fn handle_topic( + &mut self, + topic: &str, + now_millis: i64, + active_region_ids: &mut HashSet, + ) -> Result<()> { + let topic_regions = self + .table_metadata_manager + .topic_region_manager() + .regions(topic) + .await + .context(error::TableMetadataManagerSnafu)?; + + if topic_regions.is_empty() { + debug!("No regions found for topic: {}", topic); + return Ok(()); + } + active_region_ids.extend(topic_regions.keys().copied()); + + let topic_stat = self.retrieve_topic_stat(topic); + let size_based_regions = topic_stat + .map(|(latest_entry_id, avg_record_size)| { + filter_regions_by_replay_size( + topic, + topic_regions.iter().map(|(region_id, value)| { + (*region_id, value.min_entry_id().unwrap_or_default()) + }), + avg_record_size as u64, + latest_entry_id, + self.checkpoint_trigger_size, + ) + }) + .unwrap_or_default(); + // Periodic checkpoint persistence is intentionally independent of topic stats freshness: + // Kafka retention can advance even when recent write stats are unavailable. + let periodic_regions = filter_regions_for_periodic_checkpoint( + topic_regions.keys().copied(), + &self.last_checkpoint_persist_millis_by_region, + now_millis, + PERIODIC_CHECKPOINT_PERSIST_INTERVAL, + ); + let regions_to_persist = merge_region_ids(size_based_regions, periodic_regions); + let region_manifests = self + .leader_region_registry + .batch_get(topic_regions.keys().cloned()); + + match self + .persist_region_checkpoints( + topic, + ®ions_to_persist, + &topic_regions, + ®ion_manifests, + ) + .await + { + // Only mark regions that were actually written to KV. If the checkpoint is stale, + // already persisted, or the write fails, the next tick should retry. + Ok(region_ids) => mark_checkpoint_persisted( + &mut self.last_checkpoint_persist_millis_by_region, + ®ion_ids, + now_millis, + ), + Err(err) => error!(err; "Failed to persist region checkpoints for topic: {}", topic), + } + + if let Some((latest_entry_id, avg_record_size)) = topic_stat { + self.flush_regions_in_topic(topic, latest_entry_id, avg_record_size, region_manifests) + .await?; + } + + Ok(()) + } + /// Retrieves the latest entry id and average record size of a topic. /// /// Returns `None` if the topic is not found or the latest entry id is not recent. @@ -226,7 +308,7 @@ impl RegionFlushTrigger { region_ids: &[RegionId], topic_regions: &HashMap, leader_regions: &HashMap, - ) -> Result<()> { + ) -> Result> { let regions = region_ids .iter() .flat_map(|region_id| match leader_regions.get(region_id) { @@ -237,27 +319,26 @@ impl RegionFlushTrigger { .cloned() .and_then(|value| value.checkpoint), ) - .map(|checkpoint| { - ( - TopicRegionKey::new(*region_id, topic), - Some(TopicRegionValue::new(Some(checkpoint))), - ) - }), + .map(|checkpoint| (*region_id, TopicRegionValue::new(Some(checkpoint)))), None => None, }) .collect::>(); - // The`chunks` will panic if chunks_size is zero, so we return early if there are no regions to persist. + // `chunks` will panic if chunk size is zero, so return early if there are no regions to persist. if regions.is_empty() { - return Ok(()); + return Ok(Vec::new()); } let max_txn_ops = self.table_metadata_manager.kv_backend().max_txn_ops(); let batch_size = max_txn_ops.min(regions.len()); for batch in regions.chunks(batch_size) { + let batch = batch + .iter() + .map(|(region_id, value)| (TopicRegionKey::new(*region_id, topic), Some(*value))) + .collect::>(); self.table_metadata_manager .topic_region_manager() - .batch_put(batch) + .batch_put(&batch) .await .context(error::TableMetadataManagerSnafu)?; } @@ -266,7 +347,10 @@ impl RegionFlushTrigger { .with_label_values(&[topic]) .inc_by(regions.len() as u64); - Ok(()) + Ok(regions + .into_iter() + .map(|(region_id, _)| region_id) + .collect()) } async fn flush_regions_in_topic( @@ -274,45 +358,8 @@ impl RegionFlushTrigger { topic: &str, latest_entry_id: u64, avg_record_size: usize, + region_manifests: HashMap, ) -> Result<()> { - let topic_regions = self - .table_metadata_manager - .topic_region_manager() - .regions(topic) - .await - .context(error::TableMetadataManagerSnafu)?; - - if topic_regions.is_empty() { - debug!("No regions found for topic: {}", topic); - return Ok(()); - } - - // Filters regions need to persist checkpoints. - let regions_to_persist = filter_regions_by_replay_size( - topic, - topic_regions - .iter() - .map(|(region_id, value)| (*region_id, value.min_entry_id().unwrap_or_default())), - avg_record_size as u64, - latest_entry_id, - self.checkpoint_trigger_size, - ); - let region_manifests = self - .leader_region_registry - .batch_get(topic_regions.keys().cloned()); - - if let Err(err) = self - .persist_region_checkpoints( - topic, - ®ions_to_persist, - &topic_regions, - ®ion_manifests, - ) - .await - { - error!(err; "Failed to persist region checkpoints for topic: {}", topic); - } - let regions = region_manifests .into_iter() .map(|(region_id, region)| (region_id, region.manifest.prunable_entry_id())) @@ -447,6 +494,56 @@ fn filter_regions_by_replay_size>( regions_to_flush } +/// Filters regions that need periodic checkpoint persistence. +fn filter_regions_for_periodic_checkpoint( + regions: I, + last_persisted: &HashMap, + now_millis: i64, + interval: Duration, +) -> Vec +where + I: Iterator, +{ + let interval_millis = interval.as_millis() as i64; + regions + .filter(|region_id| { + last_persisted + .get(region_id) + .is_none_or(|last_persist_millis| { + now_millis.saturating_sub(*last_persist_millis) >= interval_millis + }) + }) + .collect() +} + +/// Merges two region id lists and removes duplicates. +fn merge_region_ids(left: Vec, right: Vec) -> Vec { + left.into_iter() + .chain(right) + .collect::>() + .into_iter() + .collect() +} + +/// Marks checkpoint persistence timestamps for regions. +fn mark_checkpoint_persisted( + last_persisted: &mut HashMap, + region_ids: &[RegionId], + now_millis: i64, +) { + for region_id in region_ids { + last_persisted.insert(*region_id, now_millis); + } +} + +/// Retains checkpoint persistence records for active regions. +fn retain_checkpoint_persist_records( + last_persisted: &mut HashMap, + active_region_ids: &HashSet, +) { + last_persisted.retain(|region_id, _| active_region_ids.contains(region_id)); +} + /// Group regions by leader. /// /// The regions are grouped by the leader of the region. @@ -527,6 +624,71 @@ mod tests { assert!(!is_recent(now - 1001, now, Duration::from_secs(1))); } + #[test] + fn test_filter_regions_for_periodic_checkpoint() { + let now_millis = 10_000; + let interval = Duration::from_secs(5); + let regions = vec![region_id(1, 1), region_id(1, 2), region_id(1, 3)]; + let last_persisted = HashMap::from([ + (region_id(1, 1), now_millis - 4_000), + (region_id(1, 2), now_millis - 5_000), + ]); + + let result = filter_regions_for_periodic_checkpoint( + regions.into_iter(), + &last_persisted, + now_millis, + interval, + ); + + assert_eq!(result, vec![region_id(1, 2), region_id(1, 3)]); + } + + #[test] + fn test_merge_region_ids() { + let merged = merge_region_ids( + vec![region_id(1, 1), region_id(1, 2)], + vec![region_id(1, 2), region_id(1, 3)], + ); + let merged = merged.into_iter().collect::>(); + + assert_eq!( + merged, + HashSet::from([region_id(1, 1), region_id(1, 2), region_id(1, 3)]) + ); + } + + #[test] + fn test_mark_checkpoint_persisted() { + let now_millis = 10_000; + let mut last_persisted = HashMap::from([(region_id(1, 1), 1_000)]); + + mark_checkpoint_persisted( + &mut last_persisted, + &[region_id(1, 1), region_id(1, 2)], + now_millis, + ); + + assert_eq!(last_persisted.get(®ion_id(1, 1)), Some(&now_millis)); + assert_eq!(last_persisted.get(®ion_id(1, 2)), Some(&now_millis)); + } + + #[test] + fn test_retain_checkpoint_persist_records() { + let mut last_persisted = HashMap::from([ + (region_id(1, 1), 1_000), + (region_id(1, 2), 2_000), + (region_id(1, 3), 3_000), + ]); + let active_regions = HashSet::from([region_id(1, 1), region_id(1, 3)]); + + retain_checkpoint_persist_records(&mut last_persisted, &active_regions); + + assert_eq!(last_persisted.len(), 2); + assert!(last_persisted.contains_key(®ion_id(1, 1))); + assert!(last_persisted.contains_key(®ion_id(1, 3))); + } + fn region_id(table: u32, region: u32) -> RegionId { RegionId::new(table, region) } @@ -735,6 +897,76 @@ mod tests { assert!(result.is_none()); } + #[tokio::test] + async fn test_persist_region_checkpoints_returns_written_region_ids() { + let kv_backend = Arc::new(MemoryKvBackend::new()); + let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend.clone())); + let leader_region_registry = Arc::new(LeaderRegionRegistry::new()); + let topic_stats_registry = Arc::new(TopicStatsRegistry::default()); + let mailbox_sequence = SequenceBuilder::new( + "test_persist_region_checkpoints_returns_written_region_ids", + kv_backend, + ) + .build(); + let mailbox_ctx = MailboxContext::new(mailbox_sequence); + + let (trigger, _ticker) = RegionFlushTrigger::new( + table_metadata_manager.clone(), + leader_region_registry, + topic_stats_registry, + mailbox_ctx.mailbox().clone(), + "127.0.0.1:3002".to_string(), + ReadableSize(1), + ReadableSize(1), + ); + + let topic = "test_topic"; + let region_to_write = region_id(1, 1); + let region_already_persisted = region_id(1, 2); + let region_without_leader = region_id(1, 3); + let topic_regions = HashMap::from([ + (region_to_write, TopicRegionValue::new(None)), + ( + region_already_persisted, + TopicRegionValue::new(Some(ReplayCheckpoint::new(100, None))), + ), + (region_without_leader, TopicRegionValue::new(None)), + ]); + let leader_regions = HashMap::from([ + (region_to_write, mito_leader_region(100)), + (region_already_persisted, mito_leader_region(100)), + ]); + + let written_region_ids = trigger + .persist_region_checkpoints( + topic, + &[ + region_to_write, + region_already_persisted, + region_without_leader, + ], + &topic_regions, + &leader_regions, + ) + .await + .unwrap(); + + assert_eq!(written_region_ids, vec![region_to_write]); + let persisted = table_metadata_manager + .topic_region_manager() + .get(TopicRegionKey::new(region_to_write, topic)) + .await + .unwrap() + .unwrap(); + assert_eq!(persisted.checkpoint, Some(ReplayCheckpoint::new(100, None))); + let skipped = table_metadata_manager + .topic_region_manager() + .get(TopicRegionKey::new(region_already_persisted, topic)) + .await + .unwrap(); + assert!(skipped.is_none()); + } + #[tokio::test] async fn test_send_flush_instructions_payload_includes_remote_wal_prune_reason() { let kv_backend = Arc::new(MemoryKvBackend::new()); diff --git a/src/mito2/src/engine/flush_test.rs b/src/mito2/src/engine/flush_test.rs index b86e75c72a..9279d7def9 100644 --- a/src/mito2/src/engine/flush_test.rs +++ b/src/mito2/src/engine/flush_test.rs @@ -14,6 +14,7 @@ //! Flush tests for mito engine. +use std::collections::HashMap; use std::sync::Arc; use std::sync::atomic::{AtomicI64, Ordering}; use std::time::Duration; @@ -21,11 +22,14 @@ use std::time::Duration; use api::v1::Rows; use common_recordbatch::RecordBatches; use common_time::util::current_time_millis; -use common_wal::options::WAL_OPTIONS_KEY; +use common_wal::options::{KafkaWalOptions, WAL_OPTIONS_KEY, WalOptions}; use rstest::rstest; use rstest_reuse::{self, apply}; use store_api::region_engine::RegionEngine; -use store_api::region_request::{RegionFlushRequest, RegionRequest}; +use store_api::region_request::{ + PathType, RegionCloseRequest, RegionFlushRequest, RegionOpenRequest, RegionRequest, + ReplayCheckpoint, +}; use store_api::storage::{RegionId, ScanRequest}; use crate::config::MitoConfig; @@ -319,8 +323,6 @@ async fn test_flush_empty_with_format(flat_format: bool) { async fn test_flush_reopen_region(factory: Option) { use std::collections::HashMap; - use common_wal::options::{KafkaWalOptions, WalOptions}; - common_telemetry::init_default_ut_logging(); let Some(factory) = factory else { return; @@ -394,6 +396,235 @@ async fn test_flush_reopen_region(factory: Option) { assert_eq!(5, version_data.committed_sequence); } +#[apply(single_kafka_log_store_factory)] +async fn test_skip_remote_wal_replay_sets_topic_latest_entry_id(factory: Option) { + common_telemetry::init_default_ut_logging(); + let Some(factory) = factory else { + return; + }; + + let mut env = TestEnv::new().await.with_log_store_factory(factory.clone()); + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let topic = prepare_test_for_kafka_log_store(&factory).await; + let request = CreateRequestBuilder::new() + .kafka_topic(topic.clone()) + .build(); + let table_dir = request.table_dir.clone(); + let column_schemas = rows_schema(&request); + let options = kafka_wal_options(&topic); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + + let rows = Rows { + schema: column_schemas, + rows: build_rows_for_key("a", 0, 2, 0), + }; + put_rows(&engine, region_id, rows).await; + flush_region(&engine, region_id, None).await; + // The empty flush updates `topic_latest_entry_id` from KafkaLogStore's topic stats. + flush_region(&engine, region_id, None).await; + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(1, region.topic_latest_entry_id.load(Ordering::Relaxed)); + + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: PathType::Bare, + options, + skip_wal_replay: true, + checkpoint: Some(ReplayCheckpoint { + entry_id: 1, + metadata_entry_id: None, + }), + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(1, region.topic_latest_entry_id.load(Ordering::Relaxed)); +} + +#[apply(single_kafka_log_store_factory)] +async fn test_remote_wal_open_with_replayed_memtable_sets_topic_latest_entry_id( + factory: Option, +) { + common_telemetry::init_default_ut_logging(); + let Some(factory) = factory else { + return; + }; + + let mut env = TestEnv::new().await.with_log_store_factory(factory.clone()); + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let topic = prepare_test_for_kafka_log_store(&factory).await; + let request = CreateRequestBuilder::new() + .kafka_topic(topic.clone()) + .build(); + let table_dir = request.table_dir.clone(); + let column_schemas = rows_schema(&request); + let options = kafka_wal_options(&topic); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows = Rows { + schema: column_schemas.clone(), + rows: build_rows_for_key("a", 0, 2, 0), + }; + put_rows(&engine, region_id, rows).await; + flush_region(&engine, region_id, None).await; + + let rows = Rows { + schema: column_schemas, + rows: build_rows_for_key("b", 0, 2, 0), + }; + put_rows(&engine, region_id, rows).await; + + let engine = env.reopen_engine(engine, MitoConfig::default()).await; + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: PathType::Bare, + options, + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert_eq!(1, region.version().flushed_entry_id); + assert!(!region.version().memtables.is_empty()); + assert_eq!(1, region.topic_latest_entry_id.load(Ordering::Relaxed)); +} + +#[apply(single_kafka_log_store_factory)] +async fn test_remote_wal_open_without_replayed_memtable_sets_topic_latest_entry_id( + factory: Option, +) { + common_telemetry::init_default_ut_logging(); + let Some(factory) = factory else { + return; + }; + + let mut env = TestEnv::new().await.with_log_store_factory(factory.clone()); + let engine = env.create_engine(MitoConfig::default()).await; + let region_id = RegionId::new(1, 1); + env.get_schema_metadata_manager() + .register_region_table_info( + region_id.table_id(), + "test_table", + "test_catalog", + "test_schema", + None, + env.get_kv_backend(), + ) + .await; + + let topic = prepare_test_for_kafka_log_store(&factory).await; + let request = CreateRequestBuilder::new() + .kafka_topic(topic.clone()) + .build(); + let table_dir = request.table_dir.clone(); + let column_schemas = rows_schema(&request); + let options = kafka_wal_options(&topic); + + engine + .handle_request(region_id, RegionRequest::Create(request)) + .await + .unwrap(); + let rows = Rows { + schema: column_schemas, + rows: build_rows_for_key("a", 0, 2, 0), + }; + put_rows(&engine, region_id, rows).await; + flush_region(&engine, region_id, None).await; + // The empty flush updates `topic_latest_entry_id` from KafkaLogStore's topic stats. + flush_region(&engine, region_id, None).await; + engine + .handle_request(region_id, RegionRequest::Close(RegionCloseRequest {})) + .await + .unwrap(); + + engine + .handle_request( + region_id, + RegionRequest::Open(RegionOpenRequest { + engine: String::new(), + table_dir, + path_type: PathType::Bare, + options, + skip_wal_replay: false, + checkpoint: None, + requirements: Default::default(), + }), + ) + .await + .unwrap(); + + let region = engine.get_region(region_id).unwrap(); + assert!(region.version().memtables.is_empty()); + assert_eq!(1, region.topic_latest_entry_id.load(Ordering::Relaxed)); +} + +fn kafka_wal_options(topic: &Option) -> HashMap { + topic + .as_ref() + .map(|topic| { + HashMap::from([( + WAL_OPTIONS_KEY.to_string(), + serde_json::to_string(&WalOptions::Kafka(KafkaWalOptions { + topic: topic.clone(), + })) + .unwrap(), + )]) + }) + .unwrap_or_default() +} + #[derive(Debug)] pub(crate) struct MockTimeProvider { now: AtomicI64, diff --git a/src/mito2/src/region/opener.rs b/src/mito2/src/region/opener.rs index 412172aead..e09a2e2aa0 100644 --- a/src/mito2/src/region/opener.rs +++ b/src/mito2/src/region/opener.rs @@ -339,6 +339,7 @@ impl RegionOpener { let version = VersionBuilder::new(metadata, mutable) .options(options) + .flushed_entry_id(flushed_entry_id) .build(); let version_control = Arc::new(VersionControl::new(version)); let access_layer = Arc::new(AccessLayer::new( @@ -371,7 +372,7 @@ impl RegionOpener { last_flush_millis: AtomicI64::new(now), last_schedule_compaction_millis: AtomicI64::new(now), time_provider: self.time_provider.clone(), - topic_latest_entry_id: AtomicU64::new(0), + topic_latest_entry_id: AtomicU64::new(flushed_entry_id), written_bytes: Arc::new(AtomicU64::new(0)), stats: self.stats, })) @@ -534,11 +535,11 @@ impl RegionOpener { let flushed_entry_id = version.flushed_entry_id; let version_control = Arc::new(VersionControl::new(version)); + let replay_from_entry_id = self + .replay_checkpoint + .unwrap_or_default() + .max(flushed_entry_id); let topic_latest_entry_id = if !self.skip_wal_replay { - let replay_from_entry_id = self - .replay_checkpoint - .unwrap_or_default() - .max(flushed_entry_id); info!( "Start replaying memtable at replay_from_entry_id: {} for region {}, manifest version: {}, flushed entry id: {}, elapsed: {:?}", replay_from_entry_id, @@ -561,7 +562,11 @@ impl RegionOpener { // Only set after the WAL replay is completed. if provider.is_remote_wal() && version_control.current().version.memtables.is_empty() { - wal.store().latest_entry_id(&provider).unwrap_or(0) + wal.store() + .latest_entry_id(&provider) + .unwrap_or(replay_from_entry_id) + } else if provider.is_remote_wal() { + replay_from_entry_id } else { 0 } @@ -574,7 +579,11 @@ impl RegionOpener { now.elapsed() ); - 0 + if provider.is_remote_wal() { + replay_from_entry_id + } else { + 0 + } }; if let Some(committed_in_manifest) = manifest.committed_sequence {