// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0. // Modifications Copyright 2022 PingCAP, Inc. use std::{ cmp::{max, min}, collections::HashMap, iter::Iterator, sync::{ Arc, atomic::{Ordering, Ordering::Release}, }, }; use api_version::{ ApiV2, KeyMode, KvFormat, api_v2::{KEYSPACE_PREFIX_LEN, is_whole_keyspace_range}, }; use bytes::{Buf, Bytes}; use collections::HashSet; use kvenginepb as pb; use schema::schema::StorageClassSpec; use slog_global::info; use crate::{ table::{ BoundedDataSet, DataBound, InnerKey, SnapVersion, columnar::ColumnarLevels, fts::FtsLevels, vector_index::VectorIndexes, }, table_id::{ get_table_id_from_data_bound, has_consistent_tracked_fts_indexes, keys_belong_to_same_table, merge_columnar_table_ids, }, *, }; #[derive(Debug)] pub struct CheckMergeResult { pub source_overbound: bool, pub target_overbound: bool, /// source/target region is required to be empty, but it's not. pub source_require_empty: bool, pub target_require_empty: bool, // source and target shards belong to same keyspace but with different encryption properties // (`ENCRYPTION_META_KEY` bytes and/or encryption key). pub inconsistent_encryption_key: bool, // source and target shards belong to same keyspace but with different storage class spec. pub inconsistent_storage_class_spec: bool, } impl Engine { /// Applies a split change set and materializes the resulting shard views. /// /// `skip_new_shard_ids` contains child shard IDs whose rfstore metadata /// write was intentionally skipped during merged-engine preprocessing. /// Those child shards must not be materialized here, or replay can /// recreate a stale lineage that rfstore has already advanced past. pub fn split( &self, mut cs: pb::ChangeSet, initial_seq: u64, skip_new_shard_ids: &[u64], ) -> Result<()> { let split = cs.take_split(); let sequence = cs.get_sequence(); let old_shard = self.get_shard_with_ver(cs.shard_id, cs.shard_ver)?; self.prepare_update_shard_version(&old_shard, sequence, true); let mut new_shards = vec![]; let new_shard_props = split.get_new_shards(); let new_ver = old_shard.ver + new_shard_props.len() as u64 - 1; let old_data = old_shard.get_data(); assert!( !old_data.has_txn_file_locks(), "{} shard with txn files are not allowed to split, lock_txn_files {:?}", old_shard.tag(), old_data.lock_txn_files ); let old_del_prefixes = old_shard.pending_ops.read().unwrap().del_prefixes.clone(); let sc_spec = old_shard.get_property(STORAGE_CLASS_KEY); for i in 0..=split.keys.len() { let new_shard_id = new_shard_props[i].shard_id; if new_shard_id != old_shard.id && skip_new_shard_ids.contains(&new_shard_id) { continue; } let (start_key, end_key) = get_splitting_start_end( old_shard.outer_start.chunk(), old_shard.outer_end.chunk(), split.get_keys(), i, ); let (range, inner_key_off) = if is_whole_keyspace_range(start_key, end_key) { (ShardRange::new(start_key, end_key), KEYSPACE_PREFIX_LEN) } else { debug!( "engine split"; "shard_id" => old_shard.id, "start_key" => format!("{:?}", start_key), "end_key" => format!("{:?}", end_key), "inner_key_off" => old_data.inner_key_off ); (ShardRange::new(start_key, end_key), old_data.inner_key_off) }; // Note: properties of new shards are processed in `build_split_pb`. let mut new_shard = Shard::new( self.get_engine_id(), &new_shard_props[i], new_ver, range, inner_key_off, self.opts.clone(), &self.master_key, self.encryption_key_manager.clone(), ); new_shard.parent_id = old_shard.id; if new_shard.id == old_shard.id { new_shard.set_active(old_shard.is_active()); store_u64(&new_shard.base_version, old_shard.get_base_version()); store_u64(&new_shard.meta_seq, sequence); store_u64(&new_shard.write_sequence, sequence); } else { store_u64( &new_shard.base_version, old_shard.get_base_version() + sequence, ); store_u64(&new_shard.meta_seq, initial_seq); store_u64(&new_shard.write_sequence, initial_seq); } if !old_del_prefixes.is_empty() { // We need to use the old shard's DEL_PREFIXES_KEY to overwrite the new shard's // DEL_PREFIXES_KEY. because the destroy_range compaction may have not // applied to the shard. let new_del_prefixes = old_del_prefixes.build_split(start_key, end_key); if !new_del_prefixes.is_empty() { new_shard.set_property(DEL_PREFIXES_KEY, &new_del_prefixes.marshal()); } } new_shard.set_property_opt(STORAGE_CLASS_KEY, sc_spec.as_deref()); new_shards.push(Arc::new(new_shard)); } let unconverted_l0s: HashSet = old_data .col_levels .unconverted_l0s .iter() .map(|l0| l0.id()) .collect(); for new_shard in &new_shards { if is_whole_keyspace_range( new_shard.range.outer_start.chunk(), new_shard.range.outer_end.chunk(), ) { // The newly split keyspace should be empty. continue; } let (min_table_id, max_table_id) = get_table_id_from_data_bound(new_shard.data_bound()); let columnar_table_ids: Vec<_> = old_data .columnar_table_ids .iter() .filter(|&&table_id| table_id >= min_table_id && table_id <= max_table_id) .copied() .collect(); let new_mem_tbls = new_shard.split_mem_tables(&old_data.mem_tbls); let mut new_l0s = vec![]; let mut new_unconverted_l0s = vec![]; for l0 in &old_data.l0_tbls { if new_shard.range.data_bound().overlap_bound(l0.data_bound()) { new_l0s.push(l0.clone()); if unconverted_l0s.contains(&l0.id()) && !columnar_table_ids.is_empty() { new_unconverted_l0s.push(l0.clone()); } } } let mut new_blob_tbl_map = HashMap::new(); for blob_tbl in old_data.blob_tbl_map.values() { if new_shard.overlap_bound(blob_tbl.data_bound()) { new_blob_tbl_map.insert(blob_tbl.id(), blob_tbl.clone()); } } let mut new_cfs = [ShardCf::new(0), ShardCf::new(1), ShardCf::new(2)]; for cf in 0..NUM_CFS { let old_scf = old_data.get_cf(cf); for lh in &old_scf.levels { let mut new_level_tbls = vec![]; for tbl in lh.tables.as_slice() { if new_shard.overlap_bound(tbl.data_bound()) { new_level_tbls.push(tbl.clone()); } } let new_level = LevelHandler::new(lh.level, new_level_tbls); new_cfs[cf].set_level(new_level); } } let schema_version = old_data.schema_version; let restore_version = old_data.restore_version; let schema_file = old_data.schema_file.clone(); let mut new_col_levels = ColumnarLevels::new(); new_col_levels.unconverted_l0s = new_unconverted_l0s; for col_level in &old_data.col_levels.levels { let new_col_level = &mut new_col_levels.levels[col_level.level]; for col_file in &col_level.files { if new_shard.overlap_bound(col_file.data_bound()) && !columnar_table_ids.is_empty() { new_col_level.files.push(col_file.clone()); } } } new_col_levels.l2_snap_version = old_data.col_levels.l2_snap_version; let mut new_vec_indexes = VectorIndexes::default(); for vec_index in old_data.vector_indexes.get_all() { if !columnar_table_ids.contains(&vec_index.table_id) { continue; } for file in &vec_index.files { if new_shard.overlap_bound(file.data_bound()) { new_vec_indexes.add_index_file(file.clone()); } } new_vec_indexes.update_snap_version( vec_index.table_id, vec_index.index_id, vec_index.col_id, vec_index.snap_version, ); } let shard_bound = new_shard.data_bound(); let columnar_l0_file_ids: Vec = new_col_levels.levels[0] .files .iter() .map(|file| file.id()) .collect(); let new_fts_levels = split_fts_levels( shard_bound, &columnar_table_ids, &columnar_l0_file_ids, &old_data.fts_levels, ); let mut builder = ShardDataBuilder::new(new_shard.get_data()); builder.set_mem_tbls(new_mem_tbls); builder.set_l0_tbls(new_l0s); builder.set_blob_tbls(new_blob_tbl_map); builder.set_cfs(new_cfs); builder.set_unloaded_tbls(old_data.unloaded_tbls.clone()); builder.with_new_limiter(); builder.set_schema(schema_version, restore_version, schema_file); builder.set_columnar_levels(new_col_levels); builder.set_vector_indexes(new_vec_indexes); builder.set_fts_levels(new_fts_levels); builder.set_columnar_table_ids(columnar_table_ids); builder.set_persisted_version(old_data.persisted_version); new_shard.set_data(builder.build()); } for shard in new_shards.drain(..) { let id = shard.id; if id != old_shard.id { self.insert_keyspace_shard(shard.keyspace_id, id); let shards = self.shards.pin(); // If the shard already exists, it must be created by ingest, and it maybe // newer than this one, we avoid insert it. let _ = shards.try_insert(id, shard.clone()); } else { let shards = self.shards.pin(); shards.insert(id, shard.clone()); } self.refresh_shard_states(&shard); let all_files = shard.get_all_files(); let all_col_files = shard.get_all_col_files(); let all_vec_files = shard.get_all_vec_idx_files(); let all_fts_files = shard.get_all_fts_files(); info!( "split new shard {}, start {:x}, end {:x}, all files {:?} (columnar {:?}, vector {:?}, fts {:?})", shard.tag(), shard.outer_start, shard.outer_end, all_files, all_col_files, all_vec_files, all_fts_files ); } Ok(()) } // `force_switch_mem_table` must be set `true` when & only when the shard need // initial flush, i.e. during split/prepare-merge/commit-merge. // See https://github.com/tidbcloud/cloud-storage-engine/issues/1557. pub(crate) fn prepare_update_shard_version( &self, shard: &Shard, sequence: u64, force_switch_mem_table: bool, ) { shard.write_sequence.store(sequence, Release); let version = shard.load_mem_table_snap_version(); // Switch the old shard mem-table, so the first mem-table is always empty. // ignore the read-only mem-table to be flushed. let the new shard handle it. self.switch_mem_table(shard, version, force_switch_mem_table, sequence); self.send_flush_msg(FlushMsg::Clear(shard.id)); self.send_compact_msg(CompactMsg::Clear(IdVer::new(shard.id, shard.ver))); } pub fn check_merge( &self, source_id: u64, source_ver: u64, target_id: u64, target_ver: u64, ) -> Result { let source_shard = self.get_shard_with_ver(source_id, source_ver)?; if !source_shard.get_initial_flushed() { return Err(Error::CheckMerge("source not initial flushed".to_string())); } if source_shard.has_txn_file_locks() { return Err(Error::CheckMerge("source has txn file locks".to_string())); } let source_sc_spec = StorageClassSpec::unmarshal(source_shard.get_property(STORAGE_CLASS_KEY).as_deref()); let target_shard = self.get_shard_with_ver(target_id, target_ver)?; if !target_shard.get_initial_flushed() { return Err(Error::CheckMerge("target not initial flushed".to_string())); } if !has_consistent_tracked_fts_indexes( source_shard.get_data().fts_levels.tracked_indexes_ref(), target_shard.get_data().fts_levels.tracked_indexes_ref(), source_shard.data_bound(), target_shard.data_bound(), ) { return Err(Error::CheckMerge( "source and target have inconsistent FTS indexes".to_string(), )); } let target_sc_spec = StorageClassSpec::unmarshal(target_shard.get_property(STORAGE_CLASS_KEY).as_deref()); let belongs_to_same_keyspace = ApiV2::is_belongs_to_same_keyspace( &source_shard.outer_start, &target_shard.outer_start, ); // Check if the source and target shards belong to the same keyspace but have // different encryption properties. This situation might occur during // keyspace restoration. In such cases, merging should be avoided. // If the source and target shards are from different keyspaces, the // encryption properties will be removed in `commit_merge`. let inconsistent_encryption_key = if belongs_to_same_keyspace { let source_encryption_key = source_shard.try_get_encryption_key()?; let target_encryption_key = target_shard.try_get_encryption_key()?; let encryption_key_mismatch = source_encryption_key .as_ref() .map(|k| k.cipher_text.clone()) != target_encryption_key .as_ref() .map(|k| k.cipher_text.clone()); let source_meta = source_shard.get_property(ENCRYPTION_META_KEY); let target_meta = target_shard.get_property(ENCRYPTION_META_KEY); let encryption_meta_mismatch = match (&source_meta, &target_meta) { (None, None) => false, (Some(a), Some(b)) if !a.is_empty() && !b.is_empty() => a != b, (Some(_), Some(_)) => true, (Some(_), None) | (None, Some(_)) => true, }; encryption_key_mismatch || encryption_meta_mismatch } else { false }; let inconsistent_sc_spec = if belongs_to_same_keyspace { if !source_sc_spec.is_specified() && !target_sc_spec.is_specified() { false } else { source_sc_spec != target_sc_spec || !keys_belong_to_same_table( source_shard.inner_start(), target_shard.inner_start(), ) } } else { source_sc_spec.is_specified() || target_sc_spec.is_specified() }; let (clear_source, clear_target) = need_clear_region_data_on_merge(&source_shard.outer_start, &target_shard.outer_start); let shard_is_empty = |shard: &Shard| -> bool { shard.is_empty() || shard.get_del_prefixes().cover_full_keyspace( ApiV2::get_keyspace_prefix(&shard.outer_start).unwrap_or_default(), ) }; let source_require_empty = clear_source && !shard_is_empty(source_shard.as_ref()); let target_require_empty = clear_target && !shard_is_empty(target_shard.as_ref()); Ok(CheckMergeResult { source_overbound: source_shard.has_over_bound_data(), target_overbound: target_shard.has_over_bound_data(), source_require_empty, target_require_empty, inconsistent_encryption_key, inconsistent_storage_class_spec: inconsistent_sc_spec, }) } pub fn prepare_merge(&self, shard_id: u64, shard_ver: u64, sequence: u64) -> Result<()> { let old_shard = self.get_shard_with_ver(shard_id, shard_ver).unwrap(); let old_data = old_shard.get_data(); if old_data.has_txn_file_locks() { warn!("{} prepare_merge denied, shard has txn file locks", old_shard.tag(); "txn_file_locks" => ?old_data.lock_txn_files); return Err(Error::Other( MERGE_REGION_WITH_TXN_FILE_LOCKS_ERR_MSG.into(), )); } if !old_data.col_levels.unconverted_l0s.is_empty() { warn!( "{} prepare_merge denied, unconverted l0s : {:?}", old_shard.tag(), old_data .col_levels .unconverted_l0s .iter() .map(|l0| l0.id()) .collect::>() ); return Err(Error::Other( MERGE_REGION_WITH_UNCONVERTED_L0S_ERR_MSG.into(), )); } self.prepare_update_shard_version(&old_shard, sequence, true); let mut new_shard = self.new_shard_version(&old_shard, sequence); // source shard may have non-empty mem-table, we need to flush them before // commit merge. The initial_flushed of the new shard is false, set the // parent for later initial flush. new_shard.parent_id = old_shard.id; info!("{} shard prepared merge", new_shard.tag()); self.insert_shard(Arc::new(new_shard)); Ok(()) } pub fn rollback_merge(&self, shard_id: u64, shard_ver: u64, sequence: u64) { let old_shard = self.get_shard_with_ver(shard_id, shard_ver).unwrap(); self.prepare_update_shard_version(&old_shard, sequence, false); let new_shard = self.new_shard_version(&old_shard, sequence); // There is no write during merging state, so we can directly set // initial_flushed to true. new_shard.initial_flushed.store(true, Ordering::Release); info!("{} shard rollback merge", new_shard.tag()); self.insert_shard_and_refresh(Arc::new(new_shard)); } pub fn commit_merge( &self, shard_id: u64, shard_ver: u64, source: &ChangeSet, sequence: u64, ) -> Result<()> { let old_shard = self.get_shard_with_ver(shard_id, shard_ver)?; let inner_key_off = old_shard.get_data().inner_key_off; self.prepare_update_shard_version(&old_shard, sequence, true); let source_snap = source.get_snapshot(); let belongs_to_same_keyspace = ApiV2::is_belongs_to_same_keyspace(&source_snap.outer_start, &old_shard.outer_start); if !belongs_to_same_keyspace { old_shard.del_property(ENCRYPTION_KEY); old_shard.del_property(ENCRYPTION_META_KEY); // Cross-keyspace commit-merge keeps the merge *target* shard id but removes its // encryption properties, so deregister the current encryption key for the merge // target shard to avoid leaving a stale `shard_id -> key` mapping in memory. // The merge *source* shard id will be deregistered when raftstore destroys it. self.encryption_key_manager .deregister_current_key_for_shard(shard_id); } let (clear_source, clear_target) = need_clear_region_data_on_merge(&source_snap.outer_start, &old_shard.outer_start); if clear_target { info!( "{} clear data of target shard on merge, target: {:?}", old_shard.tag(), old_shard.range, ); old_shard.set_property(DEL_PREFIXES_KEY, &[]); let source_sc = get_shard_property(STORAGE_CLASS_KEY, source_snap.get_properties()); old_shard.set_property(STORAGE_CLASS_KEY, source_sc.as_deref().unwrap_or_default()); old_shard.set_data_opt( ShardData::new_empty( old_shard.range.clone(), inner_key_off, old_shard.get_data().limiter.clone(), old_shard.get_writable_mem_table().implementation(), ), false, ); } let mut new_shard = self.new_shard_version(&old_shard, sequence); // TODO: Do we need to merge pending operations here? new_shard.range.outer_start = min( old_shard.outer_start.clone(), source_snap.outer_start.clone().into(), ); new_shard.range.outer_end = max( old_shard.outer_end.clone(), source_snap.outer_end.clone().into(), ); new_shard.ver = max(shard_ver, source.shard_ver) + 1; // make sure the new mem-table version is greater than source. let source_mem_tbl_version = source_snap.base_version + source.sequence; let target_mem_tbl_version = old_shard.get_base_version() + sequence; store_u64( &new_shard.base_version, max(source_mem_tbl_version, target_mem_tbl_version) - sequence, ); let data = if !clear_source { // merge source DEL_PREFIXES_KEY to new shard let source_del_prefixes = get_shard_property(DEL_PREFIXES_KEY, source_snap.get_properties()) .map(|v| Bytes::from(v)); let old_del_prefixes = old_shard.get_property(DEL_PREFIXES_KEY); if let Some(new_del_prefixes) = merge_del_prefixes_if_needed( source_del_prefixes, old_del_prefixes, new_shard.keyspace_id, ) { new_shard.set_property(DEL_PREFIXES_KEY, &new_del_prefixes); } let source_sc_spec = StorageClassSpec::unmarshal( get_shard_property(STORAGE_CLASS_KEY, source_snap.get_properties()).as_deref(), ); let target_sc_spec = old_shard.get_storage_class_spec(); if source_sc_spec != target_sc_spec { warn!("{} commit merge: storage class mismatch", old_shard.tag(); "source" => ?source, "target" => ?target_sc_spec); // Happen when source & target shards come from different // version (during upgrade). // TODO: uncomment the following assert after upgrade. // debug_assert!(false); } // merge shard data let old_data = old_shard.get_data(); let source_snap_version = SnapVersion::new( source_snap.get_base_version(), source_snap.get_data_sequence(), ); let new_snap_version = max(old_data.persisted_version, source_snap_version); let mem_tbls = old_data.mem_tbls.clone(); let mut blob_tbl_map = old_data.blob_tbl_map.as_ref().clone(); for v in source.blob_tables.values() { blob_tbl_map.insert(v.id(), v.clone()); } let mut l0_tbls = old_data.l0_tbls.clone(); for l0 in source.l0_tables.values() { l0_tbls.push(l0.clone()) } l0_tbls.sort_by(|a, b| b.snap_version().cmp(&a.snap_version())); let mut new_cf_builders = [ ShardCfBuilder::new(0), ShardCfBuilder::new(1), ShardCfBuilder::new(2), ]; for cf in 0..NUM_CFS { let old_scf = old_data.get_cf(cf); for level in 1..=CF_LEVELS[cf] { let old_level = old_scf.get_level(level); let cf_builder = &mut new_cf_builders[cf]; for tbl in old_level.tables.as_slice() { cf_builder.add_table(tbl.clone(), level); } } } for tbl_create in source_snap.get_table_creates() { let tbl = source.ln_tables.get(&tbl_create.id).unwrap().clone(); let cf_builder = &mut new_cf_builders[tbl_create.cf as usize]; cf_builder.add_table(tbl, tbl_create.level as usize); } let new_cfs = [ new_cf_builders[0].build(), new_cf_builders[1].build(), new_cf_builders[2].build(), ]; let mut unloaded_tbls = source.unloaded_tables.clone(); for (&id, tbl) in old_data.unloaded_tbls.iter() { unloaded_tbls.insert(id, tbl.clone()); } let lock_txn_files = old_data.lock_txn_files.clone(); assert!( source.lock_txn_files.is_empty(), "{} source shard with txn file locks are not allowed to be merged, source {:?}", old_shard.tag(), source ); let mut columnar_levels = old_data.col_levels.clone(); for columnar_create in source_snap.get_columnar_creates() { let col = source.col_files.get(&columnar_create.id).unwrap().clone(); columnar_levels.add_file(columnar_create.level as usize, col); } columnar_levels.l2_snap_version = max( old_data.col_levels.l2_snap_version, source_snap.columnar_l2_snap_version.into(), ); columnar_levels.sort(); let mut vector_indexes = old_data.vector_indexes.clone(); for source_vec_index in source_snap.get_vector_indexes() { for source_vec_idx_file_id in source_vec_index.files.iter().map(|f| f.id) { let vec_idx_file = source .vec_index_files .get(&source_vec_idx_file_id) .unwrap() .clone(); vector_indexes.add_index_file(vec_idx_file); } let old_snap_version = vector_indexes .get( source_vec_index.table_id, source_vec_index.index_id, source_vec_index.col_id, ) .map(|v| v.snap_version) .unwrap_or(SnapVersion::zero()); vector_indexes.update_snap_version( source_vec_index.table_id, source_vec_index.index_id, source_vec_index.col_id, std::cmp::max(old_snap_version, source_vec_index.snap_version.into()), ); } let mut fts_levels = Self::merge_fts(&old_data, source); let mut schema_version = old_data.schema_version; let mut restore_version = old_data.restore_version; let mut schema_file = old_data.schema_file.clone(); // If the target shard has no schema file and the source has schema file, we // should merge the schema file to target. if schema_file.is_none() { if let Some(source_schema_file) = source.schema_file.as_ref() { schema_version = source_schema_file.get_version(); restore_version = source_schema_file.get_restore_version(); schema_file = source.schema_file.clone(); } } let source_data_bound = DataBound::new( InnerKey::from_outer_key(source_snap.get_outer_start()), InnerKey::from_outer_key(source_snap.get_outer_end()), false, ); let columnar_table_ids = merge_columnar_table_ids( &source_snap.columnar_table_ids, &old_data.columnar_table_ids, source_data_bound, old_data.data_bound(), ); retain_columnar_and_vector( &columnar_table_ids, &mut columnar_levels, &mut vector_indexes, &mut fts_levels, ); let mut builder = ShardDataBuilder::new(old_data); builder.set_range(new_shard.range.clone()); if clear_target { // `inner_key_off` will be different when merge regions of different keyspaces. builder.set_inner_key_off(source_snap.inner_key_off as usize); } builder.set_mem_tbls(mem_tbls); builder.set_l0_tbls(l0_tbls); builder.set_blob_tbls(blob_tbl_map); builder.set_cfs(new_cfs); builder.set_unloaded_tbls(unloaded_tbls); builder.set_lock_txn_files(lock_txn_files); builder.set_columnar_levels(columnar_levels); builder.set_schema(schema_version, restore_version, schema_file); builder.set_vector_indexes(vector_indexes); builder.set_fts_levels(fts_levels); builder.set_columnar_table_ids(columnar_table_ids); builder.set_persisted_version(new_snap_version); builder.build() } else { info!( "{} clear data of source shard on merge, source start: {}, end: {}", old_shard.tag(), log_wrappers::hex_encode_upper(&source_snap.outer_start), &log_wrappers::hex_encode_upper(&source_snap.outer_end), ); let old_data = old_shard.get_data(); let mut builder = ShardDataBuilder::new(old_data); builder.set_range(new_shard.range.clone()); builder.build() }; new_shard.set_data(data); debug_assert_eq!(new_shard.range, new_shard.get_data().range); new_shard.parent_id = shard_id; let all_files = new_shard.get_all_files(); let all_col_files = new_shard.get_all_col_files(); let all_vec_files = new_shard.get_all_vec_idx_files(); let all_fts_files = new_shard.get_all_fts_files(); info!( "merged new shard {}, start {:x}, end {:x}, all files {:?} (columnar {:?}, vector {:?}, fts {:?})", new_shard.tag(), new_shard.outer_start, new_shard.outer_end, all_files, all_col_files, all_vec_files, all_fts_files, ); self.insert_shard_and_refresh(Arc::new(new_shard)); Ok(()) } fn merge_fts(old_data: &ShardData, source: &ChangeSet) -> FtsLevels { let mut fts_levels = (*old_data.fts_levels).clone(); let source_snap = source.get_snapshot(); for meta in source_snap.get_fts_index_metas() { fts_levels.track_index_meta(meta.get_table_id(), meta.get_index_id(), meta.clone()); } for tbl_idx in source_snap.get_fts_indexes() { fts_levels.track_legacy_index(tbl_idx.table_id, tbl_idx.index_id); } if !source_snap.get_fts_l0_files().is_empty() { let mut new_files = vec![]; for file_ref in source_snap.get_fts_l0_files() { if fts_levels .l0() .iter() .any(|file| file.id() == file_ref.get_id()) { continue; } if let Some(file) = source.fts_l0_files.get(&file_ref.get_id()) { new_files.push(file.clone()); } else { warn!( "FTS L0 file {} missing in changeset {:?}", file_ref.get_id(), source_snap ); debug_assert!(false, "missing FTS L0 file in changeset"); } } fts_levels.mut_l0(move |l0| l0.extend(new_files)); } if !source_snap.get_fts_l1_files().is_empty() { let mut new_files = vec![]; for file_ref in source_snap.get_fts_l1_files() { if fts_levels .l1() .iter() .any(|file| file.id() == file_ref.get_id()) { continue; } if let Some(file) = source.fts_l1_files.get(&file_ref.get_id()) { new_files.push(file.clone()); } else { warn!( "FTS L1 file {} missing in changeset {:?}", file_ref.get_id(), source_snap ); debug_assert!(false, "missing FTS L1 file in changeset"); } } fts_levels.mut_l1(move |l1| l1.extend(new_files)); } for file_ref in source_snap.get_fts_l2_files() { let Some(file) = source.fts_l2_files.get(&file_ref.get_id()) else { warn!( "FTS L2 file {} missing in changeset {:?}", file_ref.get_id(), source_snap ); debug_assert!(false, "missing FTS L2 file in changeset"); continue; }; if let Some(l2_files) = fts_levels.l2().get(file.props().get_lp_key()) { if l2_files.iter().any(|f| f.id() == file_ref.get_id()) { continue; } } fts_levels.insert_l2_file(file.clone()); } fts_levels.l0_snap_version = max( fts_levels.l0_snap_version, source_snap.get_fts_l0_snap_version().into(), ); if !fts_levels.is_empty() { // After shard merge, we advance `l0_snap_version` to the max of // source/target. This can make some source columnar L0 files look "covered" // (<= new watermark) even though they were never indexed on the source // side (their snap_version > source watermark). // // Those files must be marked as explicitly pending, otherwise the read // path may incorrectly skip brute-force and read stale index results. let source_l0_watermark = SnapVersion::from(source_snap.get_fts_l0_snap_version()); let mut source_pending_l0_ids = source_snap.get_fts_pending_l0_ids().to_vec(); source_pending_l0_ids.extend(source_snap.get_columnar_creates().iter().filter_map( |col| { if col.get_level() != 0 { return None; } let snap_version = SnapVersion::from(col.get_snap_version()); if snap_version > source_l0_watermark { Some(col.get_id()) } else { None } }, )); fts_levels.mut_pending_columnar_l0_ids(|ids| { ids.clear(); ids.extend(collect_pending_fts_l0_columnar_ids( old_data, &source_pending_l0_ids, )); }); } fts_levels } pub(crate) fn new_shard_version(&self, old_shard: &Shard, sequence: u64) -> Shard { let engine_id = self.get_engine_id(); let inner_key_off = old_shard.get_data().inner_key_off; let new_shard = Shard::new( engine_id, &old_shard.properties.to_pb(old_shard.id), old_shard.ver + 1, old_shard.range.clone(), inner_key_off, old_shard.opt.clone(), &self.master_key, self.encryption_key_manager.clone(), ); new_shard.set_data_opt(old_shard.get_data(), false); new_shard.set_active(old_shard.is_active()); store_u64(&new_shard.base_version, old_shard.get_base_version()); store_u64(&new_shard.meta_seq, sequence); store_u64(&new_shard.write_sequence, sequence); store_u64(&new_shard.estimated_size, old_shard.get_estimated_size()); store_u64( &new_shard.estimated_entries, old_shard.get_estimated_entries(), ); store_u64(&new_shard.max_ts, old_shard.get_max_ts()); store_u64( &new_shard.estimated_kv_size, old_shard.get_estimated_kv_size(), ); store_u64( &new_shard.estimated_ia_kv_size, old_shard.get_estimated_ia_kv_size(), ); *new_shard.store_usage_stats.write().unwrap() = old_shard.store_usage_stats.read().unwrap().clone(); new_shard } } fn retain_columnar_and_vector( columnar_table_ids: &[i64], col_levels: &mut ColumnarLevels, vector_indexes: &mut VectorIndexes, fts_levels: &mut FtsLevels, ) { col_levels.retain(|col| { let data_bound = col.data_bound(); let (min_table_id, max_table_id) = get_table_id_from_data_bound(data_bound); columnar_table_ids .iter() .any(|&id| id >= min_table_id && id <= max_table_id) }); if columnar_table_ids.is_empty() { col_levels.unconverted_l0s.clear(); } vector_indexes.retain(|idx| columnar_table_ids.contains(&idx.table_id)); let columnar_l0_ids: Vec = col_levels.levels[0] .files .iter() .map(|file| file.id()) .collect(); fts_levels.retain_for_tables(columnar_table_ids, &columnar_l0_ids); } fn split_fts_levels( shard_bound: DataBound<'_>, columnar_table_ids: &[i64], columnar_l0_file_ids: &[u64], old_levels: &FtsLevels, ) -> FtsLevels { if columnar_table_ids.is_empty() { return FtsLevels::default(); } let mut new_levels = FtsLevels::default(); new_levels.l0_snap_version = old_levels.l0_snap_version; for (table_id, index_metas) in old_levels.iter_tracked_indexes() { if columnar_table_ids.contains(table_id) { for (&index_id, meta) in index_metas { new_levels.track_index_meta(*table_id, index_id, meta.clone()); } } } new_levels.mut_l0(|l0| { l0.extend( old_levels .l0() .iter() .filter(|file| { fts_file_overlaps( shard_bound, file.props().get_smallest_key(), file.props().get_biggest_key(), ) }) .cloned(), ) }); new_levels.mut_l1(|l1| { l1.extend( old_levels .l1() .iter() .filter(|file| { fts_file_overlaps( shard_bound, file.props().get_smallest_key(), file.props().get_biggest_key(), ) }) .cloned(), ) }); new_levels.insert_l2_files(old_levels.l2().values().flat_map(|files_in_lp| { files_in_lp .iter() .filter(|file| { fts_file_overlaps( shard_bound, file.props().get_smallest_key(), file.props().get_biggest_key(), ) }) .cloned() })); if !old_levels.pending_columnar_l0_ids().is_empty() { let l0_id_set: HashSet = columnar_l0_file_ids.iter().copied().collect(); new_levels.mut_pending_columnar_l0_ids(|ids| { ids.extend( old_levels .pending_columnar_l0_ids() .iter() .copied() .filter(|id| l0_id_set.contains(id)), ); }); } new_levels } fn fts_file_overlaps(shard_bound: DataBound<'_>, smallest: &[u8], biggest: &[u8]) -> bool { let file_bound = DataBound::new( InnerKey::from_inner_buf(smallest), InnerKey::from_inner_buf(biggest), true, ); shard_bound.overlap_bound(file_bound) } pub fn get_split_shard_index(split_keys: &[Vec], key: &[u8]) -> usize { for i in 0..split_keys.len() { if key < split_keys[i].as_slice() { return i; } } split_keys.len() } /// Whether to clear the data of source and/or target regions on merge. /// /// Merging regions of different keyspaces only happens when the keyspace(s) /// has been deleted. The validity of which was ensured by PD. /// /// In this condition, we must clear the data of region in keyspace. /// /// Otherwise, as SSTs in keyspace do not have prefix (when enable key /// offset), the merged SSTs will violate data correctness. pub fn need_clear_region_data_on_merge( source_outer_start: &[u8], target_outer_start: &[u8], ) -> (bool /* clear_source */, bool /* clear_target */) { let is_same_keyspace_merge = ApiV2::is_belongs_to_same_keyspace(source_outer_start, target_outer_start); let clear_region = |key: &[u8]| -> bool { let key_mode = ApiV2::parse_key_mode(key); let in_keyspace = key_mode == KeyMode::Txn || key_mode == KeyMode::Raw; !is_same_keyspace_merge && in_keyspace }; ( clear_region(source_outer_start), clear_region(target_outer_start), ) } fn collect_pending_fts_l0_columnar_ids( data: &ShardData, source_pending_l0_ids: &[u64], ) -> Vec { let valid_ids: HashSet = data.col_levels.levels[0] .files .iter() .map(|f| f.id()) .collect(); // Existing pending file is still pending let mut pending: HashSet = data .fts_levels .pending_columnar_l0_ids() .iter() .filter(|id| valid_ids.contains(id)) .copied() .collect(); let old_watermark = data.fts_levels.l0_snap_version; // For files newer than old watermark, add them to pending ids. for col_file in &data.col_levels.levels[0].files { let l0_version = col_file.get_snap_version().unwrap_or_default(); if l0_version > old_watermark { pending.insert(col_file.id()); } } pending.extend(source_pending_l0_ids.iter().copied()); pending.into_iter().collect() } #[cfg(test)] mod tests { use std::{iter::Iterator, sync::Arc}; use protobuf::{Message, RepeatedField}; use super::*; use crate::table::{ DataBound, InnerKey, fts::test_util::{new_ded, new_packed}, }; #[test] fn test_need_clear_region_data_on_merge() { let cases: Vec<(&[u8], &[u8], (bool, bool))> = vec![ // Txn ( b"x0000", // source_outer_start b"x0001", // target_outer_start (false, false), // (clear_source, clear_target) ), (b"x0000", b"x0010", (true, true)), // Raw (b"r0000", b"r0001", (false, false)), (b"r0000", b"r0010", (false, false)), // Txn & Raw (b"x0000", b"r0000", (true, false)), // Txn & TiDB/Unknown (b"x0000", b"t", (true, false)), (b"t", b"x0000", (false, true)), (b"x0000", b"m", (true, false)), (b"m", b"x0000", (false, true)), (b"x0000", b"", (true, false)), (b"", b"x0000", (false, true)), // TiDB/Unknown (b"t", b"m", (false, false)), (b"m", b"", (false, false)), (b"", b"t", (false, false)), ]; for (idx, (source_outer_start, target_outer_start, expected)) in cases.into_iter().enumerate() { let res = need_clear_region_data_on_merge(source_outer_start, target_outer_start); assert_eq!(res, expected, "case {}", idx); } } #[test] fn test_split_fts_levels_keeps_files_when_untracked() { let mut old_levels = FtsLevels::default(); old_levels.mut_l1(|l1| { l1.push( new_packed(1, 1) .lp(1, 1, |d| { d(1, 1, false, "doc1"); }) .finish_as_file(), ); }); old_levels.insert_l2_file( new_ded(2) .lp(1, 1, |d| { d(1, 1, false, "doc1"); }) .finish_as_file(), ); // Use a very broad bound so the test is insensitive to key encoding. let shard_bound = DataBound::new( InnerKey::from_inner_buf(b""), InnerKey::from_inner_buf(b"\xff"), true, ); let new_levels = split_fts_levels(shard_bound, &[1], &[], &old_levels); assert!(new_levels.is_empty(), "expected no tracked indexes"); assert_eq!(new_levels.l1().len(), 1); assert_eq!( new_levels .l2() .values() .map(|files| files.len()) .sum::(), 1 ); } #[test] fn test_split_skips_materializing_stale_child_shards() { let (engine, _applier_tx) = crate::test_engine::new_test_engine(); let mut cs = pb::ChangeSet::new(); cs.set_shard_id(1); cs.set_shard_ver(1); cs.set_sequence(10); let mut split = pb::Split::new(); split.set_keys(RepeatedField::from_vec(vec![b"k50000".to_vec()])); let mut child = pb::Properties::new(); child.set_shard_id(2); let mut old = pb::Properties::new(); old.set_shard_id(1); split.set_new_shards(RepeatedField::from_vec(vec![child, old])); cs.set_split(split); let skip_new_shard_ids = vec![2]; engine.split(cs, 1, &skip_new_shard_ids).unwrap(); assert!(engine.get_shard(2).is_none()); assert_eq!(engine.get_shard(1).unwrap().ver, 2); } #[test] fn test_check_merge_rejects_inconsistent_encryption_meta() { ::test_util::init_log_for_test(); let (engine, _apply_tx) = crate::test_engine::new_test_engine_api_v2(); let key_builder = engine.key_builder().clone(); let inner_key_off = engine.get_shard(1).unwrap().get_data().inner_key_off; let master_key = engine.get_master_key(); let encryption_key_manager = engine.get_encryption_key_manager(); let make_shard = |id: u64, start: Vec, end: Vec| -> Arc { let mut props = pb::Properties::new(); props.shard_id = id; let shard = Shard::new( engine.get_engine_id(), &props, 1, ShardRange::new(&start, &end), inner_key_off, engine.core.opts.clone(), &master_key, encryption_key_manager.clone(), ); store_bool(&shard.active, true); store_bool(&shard.initial_flushed, true); Arc::new(shard) }; let shard2 = make_shard( 2, key_builder.i_to_outer_key(0), key_builder.i_to_outer_key(10), ); let shard3 = make_shard( 3, key_builder.i_to_outer_key(10), key_builder.i_to_outer_key(20), ); engine.insert_shard(shard2.clone()); engine.insert_shard(shard3.clone()); let meta_bytes = |file_id: u64| -> Vec { let mut meta = pb::EncryptionMeta::default(); meta.set_keyspace_id(shard2.keyspace_id); let mut epoch = pb::EncryptionEpoch::default(); epoch.set_file_id(file_id); epoch.set_data_key_id(1); meta.set_current(epoch); meta.write_to_bytes().unwrap() }; shard2.set_property(ENCRYPTION_META_KEY, &meta_bytes(10)); shard3.set_property(ENCRYPTION_META_KEY, &meta_bytes(11)); let res = engine.check_merge(2, 1, 3, 1).unwrap(); assert!(res.inconsistent_encryption_key); } // Two adjacent, flushed regions of the same table may carry incompatible // FTS metadata after an interrupted upgrade or restore. `check_merge` must // refuse so PD does not merge their physical index formats. #[test] fn test_check_merge_rejects_mixed_fts_field_layout() { use kvenginepb::fts::FtsFieldLayout; use crate::table::fts::new_legacy_tracked_index_meta; ::test_util::init_log_for_test(); let (engine, _apply_tx) = crate::test_engine::new_test_engine_api_v2(); let key_builder = engine.key_builder().clone(); let inner_key_off = engine.get_shard(1).unwrap().get_data().inner_key_off; let master_key = engine.get_master_key(); let encryption_key_manager = engine.get_encryption_key_manager(); let make_shard = |id: u64, start: Vec, end: Vec| -> Arc { let mut props = pb::Properties::new(); props.shard_id = id; let shard = Shard::new( engine.get_engine_id(), &props, 1, ShardRange::new(&start, &end), inner_key_off, engine.core.opts.clone(), &master_key, encryption_key_manager.clone(), ); store_bool(&shard.active, true); store_bool(&shard.initial_flushed, true); Arc::new(shard) }; let attach_tracked_fts = |shard: &Shard, field_layout: FtsFieldLayout| { let mut meta = new_legacy_tracked_index_meta(11, 7); meta.set_field_layout(field_layout); let mut levels = FtsLevels::default(); levels.track_index_meta(11, 7, meta); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_fts_levels(levels); shard.set_data(builder.build()); }; // Same table, adjacent regions. Bounds must decode as table keys so // `has_consistent_tracked_fts_indexes` actually compares this table // (synthetic `i_to_outer_key` ranges fall below `t` and skip FTS). let source = make_shard( 2, key_builder.gen_row_outer_key(11, 0), key_builder.gen_row_outer_key(11, 50), ); let target = make_shard( 3, key_builder.gen_row_outer_key(11, 50), key_builder.gen_row_outer_key(12, 0), ); attach_tracked_fts(&source, FtsFieldLayout::FtsFieldLayoutLegacyBody); attach_tracked_fts( &target, FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString, ); engine.insert_shard(source); engine.insert_shard(target); match engine.check_merge(2, 1, 3, 1) { Err(Error::CheckMerge(msg)) => { assert!( msg.contains("inconsistent FTS indexes"), "unexpected check_merge error: {}", msg ); } other => panic!( "mixed FTS field layouts should fail check_merge, got {:?}", other ), } // Control: once both sides persist the same layout, merge is allowed again. attach_tracked_fts( &engine.get_shard(2).unwrap(), FtsFieldLayout::FtsFieldLayoutLegacyBody, ); attach_tracked_fts( &engine.get_shard(3).unwrap(), FtsFieldLayout::FtsFieldLayoutLegacyBody, ); engine .check_merge(2, 1, 3, 1) .expect("aligned FTS layouts must be merge-compatible"); } }