rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-07-07 22:40:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: improve remote WAL replay checkpoint handling (#8225)

Browse Source 
* fix: correct topic_latest_entry_id

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: use pruned wal id for replay checkpoint

Signed-off-by: WenyXu <wenymedia@gmail.com>

* feat: persist remote wal checkpoints periodically

Signed-off-by: WenyXu <wenymedia@gmail.com>

* test: cover remote wal topic latest entry id

Signed-off-by: WenyXu <wenymedia@gmail.com>

* chore: apply suggestions

Signed-off-by: WenyXu <wenymedia@gmail.com>

* fix: merge pruned wal id for migration checkpoint

Signed-off-by: WenyXu <wenymedia@gmail.com>

---------

Signed-off-by: WenyXu <wenymedia@gmail.com>
This commit is contained in:
Weny Xu
2026-06-04 14:10:38 +08:00
committed by GitHub
parent da9d314fa4
commit 3bd915865a
8 changed files with 886 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
src/common/meta/src/key/topic_name.rs
View File

@@ -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<TopicNameKey<'_>>,
) -> Result<HashMap<String, TopicNameValue>> {
let raw_keys = topics.iter().map(|key| key.to_bytes()).collect::<Vec<_>>();
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::<Result<HashMap<_, _>>>()
}
/// 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::<Vec<_>>();
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);
}
}
}

80
src/common/meta/src/key/topic_region.rs
View File

@@ -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<Self>,
pruned_entry_id: Option<u64>,
is_metric_engine: bool,
) -> Option<Self> {
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());

77
src/datanode/src/utils.rs
View File

@@ -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<String, Vec<RegionId>>,
topic_name_values: &HashMap<String, TopicNameValue>,
) -> HashMap<RegionId, u64> {
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<HashMap<RegionId, TopicRegionValue>>,
pruned_entry_id: Option<u64>,
is_metric_engine: bool,
) -> Option<ReplayCheckpoint> {
let topic_region_values = topic_region_values.as_ref()?;
let topic_region_value = topic_region_values.get(&region_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(&region_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::<String, Vec<RegionId>>::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::<Vec<_>>();
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(&region_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,

43
src/meta-srv/src/procedure/region_migration.rs
View File

@@ -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<TopicRegionKey<'_>>,
region_topics: &[(RegionId, String, bool)],
) -> Result<HashMap<RegionId, ReplayCheckpoint>> {
let topic_region_keys = region_topics
.iter()
.map(|(region_id, topic, _)| TopicRegionKey::new(*region_id, topic))
.collect::<Vec<_>>();
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::<HashSet<_>>()
.into_iter()
.flat_map(|(key, value)| value.checkpoint.map(|value| (key, value)))
.map(TopicNameKey::new)
.collect::<Vec<_>>();
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::<HashMap<_, _>>();
Ok(replay_checkpoints)

135
src/meta-srv/src/procedure/region_migration/upgrade_candidate_region.rs
View File

@@ -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(&region_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<u32, String> {
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<u32, String>) {
prepare_table_metadata_with_engine(ctx, wal_options, "engine").await;
}
async fn prepare_table_metadata_with_engine(
ctx: &Context,
wal_options: HashMap<u32, String>,
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();

344
src/meta-srv/src/region/flush_trigger.rs
View File

@@ -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<RegionId, i64>,
/// The receiver of events.
receiver: Receiver<Event>,
}
@@ -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<RegionId>,
) -> 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,
&regions_to_persist,
&topic_regions,
&region_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,
&region_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<RegionId, TopicRegionValue>,
leader_regions: &HashMap<RegionId, LeaderRegion>,
) -> Result<()> {
) -> Result<Vec<RegionId>> {
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::<Vec<_>>();
// 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::<Vec<_>>();
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<RegionId, LeaderRegion>,
) -> 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,
&regions_to_persist,
&topic_regions,
&region_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<I: Iterator<Item = (RegionId, u64)>>(
regions_to_flush
}
/// Filters regions that need periodic checkpoint persistence.
fn filter_regions_for_periodic_checkpoint<I>(
regions: I,
last_persisted: &HashMap<RegionId, i64>,
now_millis: i64,
interval: Duration,
) -> Vec<RegionId>
where
I: Iterator<Item = RegionId>,
{
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<RegionId>, right: Vec<RegionId>) -> Vec<RegionId> {
left.into_iter()
.chain(right)
.collect::<HashSet<_>>()
.into_iter()
.collect()
}
/// Marks checkpoint persistence timestamps for regions.
fn mark_checkpoint_persisted(
last_persisted: &mut HashMap<RegionId, i64>,
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<RegionId, i64>,
active_region_ids: &HashSet<RegionId>,
) {
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::<HashSet<_>>();
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(&region_id(1, 1)), Some(&now_millis));
assert_eq!(last_persisted.get(&region_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(&region_id(1, 1)));
assert!(last_persisted.contains_key(&region_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());

239
src/mito2/src/engine/flush_test.rs
View File

@@ -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<LogStoreFactory>) {
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<LogStoreFactory>) {
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<LogStoreFactory>) {
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<LogStoreFactory>,
) {
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<LogStoreFactory>,
) {
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<String>) -> HashMap<String, String> {
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,

23
src/mito2/src/region/opener.rs
View File

@@ -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 {