// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0. // Modifications Copyright 2022 PingCAP, Inc. use std::{ collections::HashMap, fs, io::{Read, Seek, SeekFrom}, iter::Iterator, ops::Deref, os::unix::fs::{FileExt, MetadataExt}, path::{Path, PathBuf}, sync::{ Arc, atomic::{AtomicU64, Ordering::Relaxed}, }, }; use api_version::ApiV2; use bytes::{Buf, Bytes}; use cloud_encryption::EncryptionState; use dashmap::mapref::entry::Entry; use fail::fail_point; pub use file_system::IoType; use file_system::WithIoType; use kvenginepb::{TxnFileRef, TxnFileRefs}; use protobuf::Message; use schema::schema::StorageClassSpec; use tikv_util::{memory::MemoryLimiterGuard, mpsc::Receiver, time::Instant}; use txn_types::TimeStamp; use crate::{ EngineCore, apply::ChangeSet, context::IaCtx, dfs::FileType, disk_sync_progress::record_local_dir_sync_success, error::IoContext, ia::{ ia_auto_file::IaAutoFile, ia_file::{IaFile, table_meta_file_local_path}, }, limiter::{DfsLoadLimiter, DfsLoadLimiterPermit}, meta::FileMeta, metrics::{ENGINE_LEVEL_WRITE_VEC, PREPARE_COUNTER_VEC, observe_materialization_remote_read}, table::{ BoundedDataSet, Error as TableError, file::{FdCache, File, InMemFile, LocalFile}, fts::FtsCache, get_local_dir, schema_file::SchemaFile, sstable::{PROP_KEY_MAX_TS, SsTable, SsTableCore, SsTableProperty}, tiny_meta::{MetaPackReader, TypedTinyMeta}, vector_index::VectorIndexFile, }, *, }; pub struct PrepareOpts<'a> { pub use_direct_io: bool, /// I/O class used when materializing files that are missing locally. pub io_type: IoType, pub prepare_type: FilePrepareType, // The snap contains the current files that need to be reloaded. pub reload_snap: Option, pub table_filter: Option, // encryption_state will be ignored if cs is snapshot or restore_shard. pub encryption_state: EncryptionState, pub meta_pack_reader: Option<&'a MetaPackReader>, } impl Default for PrepareOpts<'_> { fn default() -> Self { Self { use_direct_io: false, io_type: IoType::Other, prepare_type: FilePrepareType::default(), reload_snap: None, table_filter: None, encryption_state: EncryptionState::default(), meta_pack_reader: None, } } } impl EngineCore { #[cfg(test)] #[allow(unused)] #[allow(deprecated)] pub(crate) fn new_for_test( fs: Arc, dfs_load_limiter: DfsLoadLimiter, opts: Arc, ) -> Self { use std::sync::{Mutex, atomic::AtomicBool}; use collections::HashSet; use dashmap::DashMap; use tikv_util::{mpsc, rolling_retriever::RollingRetriever}; use txn_types::ClusterGcStates; // Simplest possible IdAllocator for testing struct TestIdAllocator; #[async_trait::async_trait] impl IdAllocator for TestIdAllocator { fn alloc_id(&self, count: usize) -> Result> { Ok(vec![0; count]) // Just return zeros } async fn alloc_id_async(&self, count: usize) -> Result> { Ok(vec![0; count]) // Just return zeros } } // For testing, we'll create a minimal EngineCore with only the fields needed // for load_tables_by_ids Most fields will be initialized with dummy // values or defaults let (flush_tx, _) = mpsc::unbounded(); let (compact_tx, _) = mpsc::unbounded(); let (free_tx, _) = mpsc::unbounded(); // Create minimal components with default/dummy values for testing let cache = table::sstable::BlockCache::new( crate::table::sstable::BlockCacheType::Quick, opts.max_block_cache_size as u64, opts.table_builder_options.block_size, ); let columnar_meta_cache = table::columnar::ColumnarMetaCache::default(); let fd_cache = FdCache::new(1024); const FILE_LOCK_SLOTS: usize = 1024; let file_locks = (0..FILE_LOCK_SLOTS).map(|_| Mutex::new(())).collect(); // Create a dummy CompactionClient for testing let test_id_allocator = Arc::new(TestIdAllocator); let test_master_key = cloud_encryption::MasterKey::new(&[1u8; 32]); let test_encryption_key_manager = Arc::new(cloud_encryption::EncryptionKeyManager::new()); let test_security_mgr = Arc::new(security::SecurityManager::default()); // Clone values before moving them into the struct let local_dirs = opts.local_dirs.clone(); let local_dirs_clone = local_dirs.clone(); let fs_clone_1 = fs.clone(); let fs_clone_2 = fs.clone(); let cache_clone = cache.clone(); let txn_file_worker_pool_size = opts.txn_file_worker_pool_size; let ia_ctx = IaCtx::disabled(&opts.ia); let comp_client = CompactionClient::new( fs_clone_1, "".to_string(), 1, table::ChecksumType::Crc32, true, // allow_fallback_local test_id_allocator.clone(), test_master_key.clone(), test_encryption_key_manager.clone(), test_security_mgr, local_dirs, false, // for_restore columnar_meta_cache.clone(), ); Self { engine_id: AtomicU64::new(1), shards: papaya::HashMap::new(), keyspace_shards: DashMap::new(), opts, per_keyspace_configs: Arc::new(HashMap::new()), flush_tx, compact_tx, fs, cache, comp_client, id_allocator: test_id_allocator, rate_limiter: Arc::new(file_system::IoRateLimiter::new_for_test()), store_limiter: Arc::new(limiter::StoreLimiter::dummy()), free_tx, loaded: std::sync::atomic::AtomicBool::new(false), file_locks, shutting_down: std::sync::atomic::AtomicBool::new(false), fallback_to_legacy_gc_api: std::sync::atomic::AtomicBool::new(false), ks_safepoint_v2: None, cluster_gc_states: RollingRetriever::new_with(|| { Arc::new(ClusterGcStates::new( HashMap::default(), Instant::now_coarse(), )) }), master_key: test_master_key, encryption_key_manager: test_encryption_key_manager, txn_chunk_mgr: txn_chunk_manager::TxnChunkManager::new( local_dirs_clone.iter().map(|dir| dir.join("txn")).collect(), fs_clone_2, cache_clone, Some(fd_cache.clone()), txn_chunk_manager::with_pool_size(txn_file_worker_pool_size), txn_chunk_manager::TxnChunkManagerConfig::default(), ), ia_ctx, files_in_blacklist: Arc::new(HashSet::default()), schema_files: Arc::new(DashMap::new()), dfs_load_limiter, fd_cache, value_cache: None, columnar_meta_cache, managed_safe_ts: AtomicU64::new(0), meta_pack_scheduler: None, available_space_bytes: AtomicU64::new(0), local_disk_sync_progress: std::sync::OnceLock::new(), worker_handles: Default::default(), file_gc_compaction_trigger_limiter: FileGcCompactionTriggerLimiter::new(), } } pub fn prepare_change_set( &self, cs: kvenginepb::ChangeSet, opts: PrepareOpts<'_>, ) -> Result { let PrepareOpts { use_direct_io, io_type, prepare_type, reload_snap, table_filter, mut encryption_state, meta_pack_reader, } = opts; let mut ids: HashMap = HashMap::new(); let mut lock_txn_file_refs: Vec = vec![]; let mut cs = ChangeSet::new(cs); let mut snap = None; let mut schema_meta = None; if cs.has_flush() { let flush = cs.get_flush(); if flush.has_l0_create() { ids.insert( flush.get_l0_create().id, FileMeta::from_l0_table(flush.get_l0_create()), ); } for l0 in flush.get_l0_creates() { ids.insert(l0.get_id(), FileMeta::from_l0_table(l0)); } } if cs.has_compaction() { let comp = cs.get_compaction(); if !is_move_down(comp) { for tbl in &comp.table_creates { ids.insert(tbl.id, FileMeta::from_table(tbl)); } for bt in comp.get_blob_tables() { ids.insert(bt.get_id(), FileMeta::from_blob_table(bt)); } } else if comp.level == 0 && (prepare_type.not_local() || self.opts.ia.force_ia) { for tbl in &comp.table_creates { ids.insert(tbl.id, FileMeta::from_table(tbl)); } } } if cs.has_destroy_range() { for t in cs.get_destroy_range().get_table_creates() { ids.insert(t.id, FileMeta::from_table(t)); } for col in cs.get_destroy_range().get_columnar_creates() { ids.insert(col.id, FileMeta::from_columnar_table(col)); } } if cs.has_truncate_ts() { for t in cs.get_truncate_ts().get_table_creates() { ids.insert(t.id, FileMeta::from_table(t)); } for col in cs.get_truncate_ts().get_columnar_creates() { ids.insert(col.id, FileMeta::from_columnar_table(col)); } } if cs.has_trim_over_bound() { for t in cs.get_trim_over_bound().get_table_creates() { ids.insert(t.id, FileMeta::from_table(t)); } for col in cs.get_trim_over_bound().get_columnar_creates() { ids.insert(col.id, FileMeta::from_columnar_table(col)); } } if cs.has_snapshot() { snap = Some(cs.get_snapshot()); } if cs.has_initial_flush() { let snap = cs.get_initial_flush(); self.collect_snap_ids(snap, &mut ids); // Load schema file when prepare initial_flush if snap.has_schema_meta() { schema_meta = Some(snap.get_schema_meta()); } } if cs.has_restore_shard() { snap = Some(cs.get_restore_shard()); } if cs.has_ingest_files() { let ingest_files = cs.get_ingest_files(); for l0 in ingest_files.get_l0_creates() { ids.insert(l0.id, FileMeta::from_l0_table(l0)); } for tbl in ingest_files.get_table_creates() { ids.insert(tbl.id, FileMeta::from_table(tbl)); } for blob in ingest_files.get_blob_creates() { ids.insert(blob.id, FileMeta::from_blob_table(blob)); } } if cs.has_major_compaction() { let major_comp = cs.get_major_compaction(); for tbl in major_comp.get_sstable_change().get_table_creates() { ids.insert(tbl.get_id(), FileMeta::from_table(tbl)); } for blob in major_comp.get_new_blob_tables() { ids.insert(blob.get_id(), FileMeta::from_blob_table(blob)); } } if cs.has_update_schema_meta() { schema_meta = Some(cs.get_update_schema_meta()); } if let Some(reload_snap) = &reload_snap { self.collect_snap_ids(reload_snap, &mut ids); } if cs.has_columnar_compaction() { let columnar_comp = cs.get_columnar_compaction(); for tbl in columnar_comp.get_columnar_change().get_columnar_creates() { ids.insert(tbl.get_id(), FileMeta::from_columnar_table(tbl)); } } if cs.has_update_vector_index() { let update_vec_idx = cs.get_update_vector_index(); for vec_idx_file in update_vec_idx.get_added() { ids.insert( vec_idx_file.get_id(), FileMeta::from_vector_index_file(vec_idx_file), ); } } for file in cs.get_fts_update().get_l0_add_files() { ids.insert(file.get_id(), FileMeta::from_fts_l0_file(file)); } for file in cs.get_fts_update().get_l1_add_files() { ids.insert(file.get_id(), FileMeta::from_fts_l1_file(file)); } for file in cs.get_fts_update().get_l2_add_files() { ids.insert(file.get_id(), FileMeta::from_fts_l2_file(file)); } let mut prepare_type = prepare_type; if let Some(snap) = snap { self.collect_snap_ids(snap, &mut ids); lock_txn_file_refs.extend(collect_snap_lock_txn_file_refs(snap)); encryption_state = encryption_state_from_shard_properties( snap.get_properties(), self.encryption_key_manager.clone(), &self.master_key, ApiV2::get_u32_keyspace_id_by_key(snap.get_outer_start()), ); prepare_type = FilePrepareType::from_snapshot(snap); if snap.has_schema_meta() { schema_meta = Some(snap.get_schema_meta()); } } if let Some(schema_meta) = schema_meta { let schema_file_id = schema_meta.get_file_id(); if schema_file_id != 0 { // schema file is already loaded. if let Some(file) = self.schema_files.get(&schema_file_id) { cs.schema_file = Some(file.value().clone()); } else { ids.insert(schema_file_id, FileMeta::from_schema_meta()); } } }; let tag = ShardTag::new(self.get_engine_id(), IdVer::from_change_set(&cs)); if let Some(table_filter) = table_filter { info!( "{} is preparing change set (before table filter)", tag; "ids" => ?ids.keys(), ); cs.unloaded_tables = ids .extract_if(|_, tb| !table_filter(cs.shard_id, tb)) // !table_filter: table will not load .collect(); } if self.opts.ignore_columnar_table_load { ids.retain(|_, meta| { meta.file_type != FileType::Columnar && meta.file_type != FileType::VectorIndex && meta.file_type != FileType::FtsPackedFile && meta.file_type != FileType::FtsDedicatedFile }); } (|| fail_point!("before_track_pending_files", |_| {}))(); // Track pending files to avoid GC mistakenly deleting them before they are // applied to KvEngine. Will be untracked in the KvEngine::apply_change_set(). self.track_unapplied_pending_files( cs.shard_id, cs.sequence, PendingFileType::ChangeSet, &ids, ); info!( "{} is preparing change set, loading file by ids", tag; "ids" => ?ids.keys(), "encryption_state" => ?encryption_state, "prepare_type" => ?prepare_type, "io_type" => io_type.as_str(), ); self.load_tables_by_ids( tag, cs.shard_id, cs.shard_ver, &ids, &mut cs, use_direct_io, io_type, prepare_type, encryption_state.clone(), meta_pack_reader, )?; if !lock_txn_file_refs.is_empty() { cs.lock_txn_files = self.txn_chunk_mgr.load_txn_files_from_refs( cs.shard_id, cs.shard_ver, &lock_txn_file_refs, encryption_state.clone(), io_type, )?; info!( "{} is preparing change set, load lock txn files", tag; "lock_txn_files" => ?cs.lock_txn_files, "encryption_state" => ?encryption_state, ); } (|| { fail_point!("after_load_local_files", cs.flush.is_some(), |_| {}); })(); Ok(cs) } fn collect_snap_ids(&self, snap: &kvenginepb::Snapshot, ids: &mut HashMap) { for blob in snap.get_blob_creates() { ids.insert(blob.id, FileMeta::from_blob_table(blob)); } for l0 in snap.get_l0_creates() { ids.insert(l0.id, FileMeta::from_l0_table(l0)); } for ln in snap.get_table_creates() { ids.insert(ln.id, FileMeta::from_table(ln)); } for col in snap.get_columnar_creates() { ids.insert(col.id, FileMeta::from_columnar_table(col)); } for vec_idx in snap.get_vector_indexes() { for vec_idx_file in vec_idx.get_files() { ids.insert( vec_idx_file.id, FileMeta::from_vector_index_file(vec_idx_file), ); } } for file in snap.get_fts_l0_files() { ids.insert(file.id, FileMeta::from_fts_l0_file(file)); } for file in snap.get_fts_l1_files() { ids.insert(file.id, FileMeta::from_fts_l1_file(file)); } for file in snap.get_fts_l2_files() { ids.insert(file.id, FileMeta::from_fts_l2_file(file)); } } fn load_tables_by_ids( &self, tag: ShardTag, shard_id: u64, shard_ver: u64, ids: &HashMap, cs: &mut ChangeSet, use_direct_io: bool, io_type: IoType, shard_prepare_type: FilePrepareType, encryption_state: EncryptionState, meta_pack_reader: Option<&MetaPackReader>, ) -> Result<()> { let start_time = Instant::now_coarse(); let (result_tx, result_rx) = tikv_util::mpsc::bounded(ids.len()); let runtime = self.fs.get_runtime().handle(); let opts = dfs::Options::default().with_shard(shard_id, shard_ver); let current_ts: Option = matches!(&shard_prepare_type, FilePrepareType::AutoIa(_)) .then(|| TimeStamp::now().into_inner()); let available_permits = self.dfs_load_limiter.available_permits(); let mut join_set = tokio::task::JoinSet::new(); let mut msg_count = 0; for (&id, fm) in ids { let fs = self.fs.clone(); let prepare_type = if self.ia_ctx.is_enabled() { shard_prepare_type.convert_for_prepare_table(fm, self.opts.ia.force_ia) } else { FilePrepareType::Local { transit_from_ia: shard_prepare_type.transit_from_ia(), } }; let tiny_meta = meta_pack_reader .as_ref() .and_then(|reader| reader.get(id)) .map(|tm| tm.try_convert_to(fm.file_type)) .inspect(|_| { PREPARE_COUNTER_VEC.tiny_meta_hit.inc(); }) .unwrap_or_default(); // Try open local file first. let mut table_meta_file = None; let file: Option> = match &prepare_type { FilePrepareType::Local { .. } => { // TODO: support tiny meta. self.open_local_file(id, fm.file_type, None).ok().map(|f| { PREPARE_COUNTER_VEC.local_files.inc(); Arc::new(f) as _ }) } FilePrepareType::Ia => self .try_open_local_ia_file(tag, id, fm, use_direct_io, io_type, &tiny_meta)? .map(|f| { PREPARE_COUNTER_VEC.local_ia_files.inc(); Arc::new(f) as _ }), FilePrepareType::AutoIa(spec) => { let (ia_auto_f, table_meta_f) = self.try_open_local_auto_ia_file( tag, id, fm, use_direct_io, io_type, current_ts.unwrap(), spec, &TypedTinyMeta::None, // TODO: support tiny meta. )?; table_meta_file = table_meta_f; ia_auto_f.map(|f| { PREPARE_COUNTER_VEC.local_auto_ia_files.inc(); Arc::new(f) as _ }) } }; if let Some(file) = file { match cs.add_file( id, file, fm, self.cache.clone(), None, None, FtsCache::disabled(), encryption_state.clone(), self.columnar_meta_cache.clone(), tiny_meta, self.meta_pack_scheduler.as_ref(), ) { Ok(()) => { if fm.file_type == FileType::Schema { if let Some(schema_file) = cs.schema_file.as_ref() { debug_assert_eq!(schema_file.get_file_id(), id); self.schema_files.insert(id, schema_file.clone()); } } continue; } Err(e) => { if should_remove_corrupted_local_file(&e) { let path = self.local_file_path(id, fm.file_type); remove_corrupted_local_file(&path, id, &e.to_string()); } else { return Err(e); } } } } // Load from remote. let tx = result_tx.clone(); let fm = fm.clone(); join_set.spawn_on( Self::load_file_from_remote( self.dfs_load_limiter.clone(), fs, id, fm, io_type, prepare_type, opts, tag, table_meta_file, current_ts, self.transition_limiter(), tx, ), runtime, ); if msg_count < self.opts.dfs_load_concurrency_per_request { msg_count += 1; } else { self.recv_file_data( cs, use_direct_io, io_type, &result_rx, encryption_state.clone(), )?; } } for _ in 0..msg_count { self.recv_file_data( cs, use_direct_io, io_type, &result_rx, encryption_state.clone(), )?; } PREPARE_COUNTER_VEC.total_files.inc_by(ids.len() as u64); info!("{} load tables by ids", tag; "ids" => ids.len(), "takes" => ?start_time.saturating_elapsed(), "available_permits" => available_permits); Ok(()) } async fn load_file_from_remote( dfs_load_limiter: DfsLoadLimiter, fs: Arc, id: u64, fm: FileMeta, io_type: IoType, prepare_type: FilePrepareType, opts: dfs::Options, tag: ShardTag, table_meta_file: Option>, current_ts: Option, transition_limiter: crate::ia::transition::TransitionLimiter, tx: tikv_util::mpsc::Sender>, ) { let res = async { let transition_file_size = if prepare_type.transit_from_ia() { match fs.size(id, opts.with_type(fm.file_type)).await { Ok(total) => Some(total), Err(err) => { warn!( "failed to get dfs file size for IA to Standard transition, fallback to unbounded loading"; "file_id" => id, "err" => ?err ); None } } } else { None }; let transition_remote_read_charge = match transition_file_size { Some(total) => Some(transition_limiter.throttle_remote_read(total).await), None => None, }; let permit = dfs_load_limiter.acquire_permit().await; let memory_guard = if prepare_type.is_local() { let total = match transition_file_size { Some(total) => Some(total), None => match fs.size(id, opts.with_type(fm.file_type)).await { Ok(total) => Some(total), Err(err) => { warn!( "failed to get dfs file size, fallback to unbounded loading"; "file_id" => id, "err" => ?err ); None } }, }; match total { Some(total) => Some(dfs_load_limiter.acquire_memory_blocking(total).await), None => None, } } else { None }; let res = match prepare_type { FilePrepareType::Local { .. } => { Self::load_remote_table(id, fm.file_type, opts, fs.as_ref(), io_type) .await .map(|local_data| PreparedFileResult::Local { local_data }) } FilePrepareType::Ia => { Self::load_remote_table_meta(tag, id, &fm, opts, fs.as_ref(), io_type) .await .map(|table_meta_data| PreparedFileResult::Ia { table_meta_data }) } FilePrepareType::AutoIa(spec) => { Self::load_remote_auto_ia_file( tag, id, &fm, opts, fs.as_ref(), io_type, table_meta_file, current_ts.unwrap(), spec, ) .await } }; if let (true, Some(charge)) = (res.is_ok(), transition_remote_read_charge) { charge.commit(); } res.map(|prepared| LoadRemoteResult { id, fm, prepared, permit, memory_guard, }) } .await; let _ = tx.send(res); // memory_guard will be passed to recv_file_data and released // after file processing } fn recv_file_data( &self, cs: &mut ChangeSet, use_direct_io: bool, io_type: IoType, result_rx: &Receiver>, encryption_state: EncryptionState, ) -> Result<()> { let LoadRemoteResult { id, fm, prepared, permit, memory_guard: _memory_guard, } = result_rx.recv().unwrap()?; fail_point!("doing_recv_file_data"); fail_point!("return_recv_file_data", |_| { Ok(()) }); let file: Arc = match prepared { PreparedFileResult::Local { local_data } => self .save_and_open_table(id, &fm, local_data, permit, use_direct_io, io_type) .map(|f| { PREPARE_COUNTER_VEC.remote_files.inc(); Arc::new(f) as _ })?, PreparedFileResult::Ia { table_meta_data } => self .save_and_open_ia_file( id, &fm, table_meta_data, Some(permit), use_direct_io, io_type, ) .map(|f| { PREPARE_COUNTER_VEC.remote_ia_files.inc(); Arc::new(f) as _ })?, prepared @ PreparedFileResult::AutoIa { .. } => self .save_and_open_auto_ia_file(id, &fm, prepared, permit, use_direct_io, io_type) .map(|f| { PREPARE_COUNTER_VEC.remote_auto_ia_files.inc(); Arc::new(f) as _ })?, }; if fm.file_type == FileType::Schema { let schema_file = SchemaFile::open(file.clone())?; self.schema_files.insert(id, schema_file); } cs.add_file( file.id(), file, &fm, self.cache.clone(), None, None, FtsCache::disabled(), encryption_state, self.columnar_meta_cache.clone(), TypedTinyMeta::None, self.meta_pack_scheduler.as_ref(), )?; // _memory_guard will be dropped here, releasing the memory after file // processing is complete Ok(()) } /// Materializes unloaded tables using the I/O class selected by the caller. pub fn load_unloaded_tables( &self, shard_id: u64, shard_ver: u64, use_direct_io: bool, io_type: IoType, ) -> Result<()> { let shard = self.get_shard(shard_id).unwrap(); let tag = shard.tag(); let mut cs = ChangeSet::new(kvenginepb::ChangeSet::default()); let data = shard.get_data(); let mut builder = ShardDataBuilder::new(data.clone()); let load_tables = data .unloaded_tbls .iter() .filter(|&(_, tbl)| shard.overlap_bound(tbl.data_bound())) .map(|(id, tbl)| (*id, tbl.clone())) .collect(); info!("{} load_unloaded_tables: {:?}", tag, load_tables); self.load_tables_by_ids( tag, shard_id, shard_ver, &load_tables, &mut cs, use_direct_io, io_type, FilePrepareType::Local { transit_from_ia: false, }, shard.get_encryption_state(), None, )?; // level 0 let mut new_l0s = data.l0_tbls.clone(); // blob let mut new_blob_tbl_map = data.blob_tbl_map.as_ref().clone(); // columnar let mut new_columnar_levels = data.col_levels.clone(); // level n let mut scf_builders = vec![]; for cf in 0..NUM_CFS { let scf = ShardCfBuilder::new(cf); scf_builders.push(scf); } for (cf, shard_cf) in data.cfs.iter().enumerate() { for lh in &shard_cf.levels { for sst in &(*lh.tables) { let scf = &mut scf_builders.as_mut_slice()[cf]; scf.add_table(sst.clone(), lh.level); } } } // vector index let mut new_vec_indexes = data.vector_indexes.clone(); // fts let mut new_fts_levels = data.fts_levels.deref().clone(); let mut new_fts_l0_files = vec![]; let mut new_fts_l1_files = vec![]; for (id, tbl) in load_tables.into_iter() { match tbl.file_type { FileType::Sst => { if tbl.level == 0 { let l0 = cs.l0_tables.get(&id).unwrap().clone(); new_l0s.push(l0); } else { let sst = cs.ln_tables.get(&id).unwrap().clone(); let scf = &mut scf_builders.as_mut_slice()[tbl.get_cf() as usize]; scf.add_table(sst, tbl.get_level() as usize); } } FileType::Columnar => { let col_file = cs.col_files.get(&id).unwrap().clone(); new_columnar_levels.add_file(tbl.level as usize, col_file); } FileType::Blob => { let blob_file = cs.blob_tables.get(&id).unwrap().clone(); new_blob_tbl_map.insert(id, blob_file); } FileType::TxnChunk | FileType::Schema => unreachable!("not supported"), FileType::EncryptionDict => unreachable!("EncryptionDict not supported"), FileType::VectorIndex => { let vec_idx_file = cs.vec_index_files.get(&id).unwrap().clone(); new_vec_indexes.add_index_file(vec_idx_file); } FileType::FtsPackedFile => { if let Some(l0_file) = cs.fts_l0_files.get(&id) { new_fts_l0_files.push(l0_file.clone()); } else if let Some(l1_file) = cs.fts_l1_files.get(&id) { new_fts_l1_files.push(l1_file.clone()); } else { panic!("unexpected FTS file id {}", id); } } FileType::FtsDedicatedFile => { let l2_file = cs.fts_l2_files.get(&id).unwrap().clone(); new_fts_levels.insert_l2_file(l2_file); } } } if !new_fts_l0_files.is_empty() { new_fts_levels.mut_l0(move |l0| l0.extend(new_fts_l0_files)); } if !new_fts_l1_files.is_empty() { new_fts_levels.mut_l1(move |l1| l1.extend(new_fts_l1_files)); } new_l0s.sort_by(|a, b| b.snap_version().cmp(&a.snap_version())); let mut scfs = [ShardCf::new(0), ShardCf::new(1), ShardCf::new(2)]; for cf in 0..NUM_CFS { let scf = &mut scf_builders.as_mut_slice()[cf]; scfs[cf] = scf.build(); } new_columnar_levels.sort(); new_vec_indexes.sort(); builder.set_l0_tbls(new_l0s); builder.set_blob_tbls(new_blob_tbl_map); builder.set_cfs(scfs); builder.set_unloaded_tbls(HashMap::new()); builder.set_columnar_levels(new_columnar_levels); builder.set_vector_indexes(new_vec_indexes); builder.set_fts_levels(new_fts_levels); shard.set_data(builder.build()); Ok(()) } fn write_local_file( &self, id: u64, data: Bytes, use_direct_io: bool, io_type: IoType, file_type: FileType, ) -> Result<()> { let local_file_name = self.local_file_path(id, file_type); self.write_local_file_with_file_path(id, data, use_direct_io, io_type, &local_file_name) } fn write_local_file_with_file_path( &self, id: u64, data: Bytes, use_direct_io: bool, io_type: IoType, file_path: &Path, ) -> Result<()> { let start_time = Instant::now_coarse(); let tmp_file_name = Self::tmp_file_path(file_path); let _io_type_guard = WithIoType::new(io_type); if use_direct_io { let mut writer = file_system::DirectWriter::new(self.rate_limiter.clone(), IoType::Compaction); writer .write_to_file(data.chunk(), &tmp_file_name) .table_ctx(id, "write_local_file.direct_write_tmp")?; } else { let mut file = std::fs::File::create(&tmp_file_name) .table_ctx(id, "write_local_file.create_tmp")?; let synced_data = file_system::write_all_with_rate_limiter( &mut file, data.chunk(), &self.rate_limiter, IoType::Compaction, 256 * 1024, ) .table_ctx(id, "write_local_file.write_tmp")?; if synced_data < data.len() { file.sync_data() .table_ctx(id, "write_local_file.sync_tmp")?; } } fail_point!("before_rename_tmp_to_sst"); let rename_time = Instant::now_coarse(); std::fs::rename(tmp_file_name, file_path).table_ctx(id, "write_local_file.rename")?; let end_time = Instant::now_coarse(); if let Some(local_dir) = file_path.parent() { record_local_dir_sync_success(local_dir); } let takes_msg = format!( "{:?}(w:{:?},r:{:?})", // total(write,sync,rename) end_time.saturating_duration_since(start_time), rename_time.saturating_duration_since(start_time), end_time.saturating_duration_since(rename_time) ); info!( "{}: write local file {}", self.get_engine_id(), id; "size" => data.len(), "takes" => takes_msg, "path" => ?file_path.file_name().unwrap_or_default(), ); Ok(()) } pub fn write_local_file_if_not_exists( &self, id: u64, data: Bytes, file_type: FileType, ) -> Result<()> { let file_path = self.local_file_path(id, file_type); let _guard = self.lock_file(id); if !file_path.exists() { self.write_local_file(id, data, false, IoType::Other, file_type)?; } else { let meta = file_path.metadata().table_ctx(id, "metadata")?; if meta.size() != data.len() as u64 { return Err(table::Error::InvalidFileSize.into()); } } Ok(()) } pub fn read_local_file( &self, id: u64, file_type: FileType, start_off: u64, end_off: Option, check_exists_only: bool, ) -> Result { let _guard = self.lock_file(id); let path = self.local_file_path(id, file_type); if check_exists_only { let _ = fs::metadata(path).table_ctx(id, "read_local_file.metadata")?; return Ok(Bytes::new()); } let mut f = fs::File::open(path).table_ctx(id, "read_local_file.open")?; if let Some(end_off) = end_off { let mut buf = vec![0; (end_off - start_off) as usize]; f.read_exact_at(&mut buf, start_off).dfs_ctx(id, "read")?; Ok(buf.into()) } else { if start_off > 0 { f.seek(SeekFrom::Start(start_off)) .table_ctx(id, "read_local_file.seek")?; } let mut data = vec![]; f.read_to_end(&mut data).table_ctx(id, "read_local_file")?; Ok(data.into()) } } pub fn local_file_size(&self, id: u64, file_type: FileType) -> Result { let _guard = self.lock_file(id); let path = self.local_file_path(id, file_type); let meta = fs::metadata(path).table_ctx(id, "local_file_size.metadata")?; Ok(meta.size()) } pub(crate) fn open_local_file( &self, id: u64, file_type: FileType, file_size: Option, ) -> Result { let path = self.local_file_path(id, file_type); self.open_local_file_with_file_path(id, Some(self.fd_cache.clone()), path, file_size) } fn open_local_file_with_file_path( &self, id: u64, fd_cache: Option, file_path: PathBuf, file_size_opt: Option, ) -> Result { let _guard = self.lock_file(id); Ok(LocalFile::open_ext( id, file_path, file_size_opt, fd_cache, self.loaded.load(Relaxed), )?) } fn try_open_local_ia_file( &self, tag: ShardTag, id: u64, fm: &FileMeta, use_direct_io: bool, io_type: IoType, tiny_meta: &TypedTinyMeta, ) -> Result> { let IaCtx::Enabled(ia_mgr, data_dirs) = self.ia_ctx.clone() else { return Ok(None); }; let data_dir = get_local_dir(&data_dirs, id); let meta_file_path = table_meta_file_local_path(id, fm.file_type, data_dir.deref()); let mut table_meta_file = self.open_local_file_with_file_path( id, ia_mgr.get_meta_fd_cache(), meta_file_path.clone(), tiny_meta.meta_size(), ); if table_meta_file.is_err() { if let Ok(local_file) = self.open_local_file(id, fm.file_type, tiny_meta.file_size()) { let table_meta_off = fm.table_meta_off as u64; // Read table meta from local sst file. if let Ok(data) = local_file .read( table_meta_off, local_file.size() as usize - table_meta_off as usize, ) .map_err(Into::into) .and_then(|data| validate_table_meta_off(tag, id, fm, data)) { let data_len = data.len() as u64; self.write_local_file_with_file_path( id, data, use_direct_io, io_type, &meta_file_path, )?; table_meta_file = self.open_local_file_with_file_path( id, ia_mgr.get_meta_fd_cache(), meta_file_path, Some(data_len), ); }; } } if let Ok(table_meta_file) = table_meta_file { if let Ok(ia_file) = IaFile::open_with_tiny_meta( id, fm, Arc::new(table_meta_file), ia_mgr, tiny_meta, true, ) { return Ok(Some(ia_file)); } } Ok(None) } async fn load_remote_table( id: u64, ftype: FileType, opts: dfs::Options, fs: &dyn dfs::Dfs, io_type: IoType, ) -> Result { Self::read_remote_file(id, opts.with_type(ftype), fs, io_type).await } async fn read_remote_file( id: u64, opts: dfs::Options, fs: &dyn dfs::Dfs, io_type: IoType, ) -> Result { let start_time = Instant::now_coarse(); let result = fs.read_file(id, opts).await.map_err(Into::into); observe_materialization_remote_read( io_type, result.as_ref().ok().map(Bytes::len), start_time.saturating_elapsed(), ); result } fn save_and_open_table( &self, id: u64, fm: &FileMeta, data: Bytes, permit: DfsLoadLimiterPermit, use_direct_io: bool, io_type: IoType, ) -> Result { let data_len = data.len() as u64; self.write_local_file(id, data, use_direct_io, io_type, fm.file_type)?; drop(permit); ENGINE_LEVEL_WRITE_VEC .with_label_values(&[&fm.get_level().to_string()]) .inc_by(data_len); let file = self.open_local_file(id, fm.file_type, Some(data_len))?; Ok(file) } async fn load_remote_table_meta( tag: ShardTag, id: u64, fm: &FileMeta, opts: dfs::Options, fs: &dyn dfs::Dfs, io_type: IoType, ) -> Result { let opts = opts .with_type(fm.file_type) .with_start_off(fm.table_meta_off as u64); Self::read_remote_file(id, opts, fs, io_type) .await .and_then(|data| validate_table_meta_off(tag, id, fm, data)) } fn save_and_open_ia_file( &self, id: u64, fm: &FileMeta, table_meta_data: Bytes, permit: Option, use_direct_io: bool, io_type: IoType, ) -> Result { let IaCtx::Enabled(ia_mgr, data_dirs) = &self.ia_ctx else { unreachable!("ia_ctx should be enabled"); }; let data_len = table_meta_data.len() as u64; let data_dir = get_local_dir(data_dirs, id); let meta_file_path = table_meta_file_local_path(id, fm.file_type, data_dir.deref()); self.write_local_file_with_file_path( id, table_meta_data, use_direct_io, io_type, &meta_file_path, )?; drop(permit); ENGINE_LEVEL_WRITE_VEC .with_label_values(&[&fm.get_level().to_string()]) .inc_by(data_len); let meta_file = self.open_local_file_with_file_path( id, ia_mgr.get_meta_fd_cache(), meta_file_path, Some(data_len), )?; let table_meta_file = Arc::new(meta_file); Ok(IaFile::open(id, fm, table_meta_file, ia_mgr.clone(), true)?) } fn try_open_local_auto_ia_file( &self, tag: ShardTag, id: u64, fm: &FileMeta, use_direct_io: bool, io_type: IoType, current_ts: u64, spec: &StorageClassSpec, tiny_meta: &TypedTinyMeta, ) -> Result<( Option, Option>, // table_meta_file )> { debug_assert_eq!(fm.file_type, FileType::Sst); let ia_file = self.try_open_local_ia_file(tag, id, fm, use_direct_io, io_type, tiny_meta)?; let Some(ia_file) = ia_file else { return Ok((None, None)); }; let SsTableProperty::MaxTs(max_ts) = SsTableCore::extract_property_from_table_meta_file( ia_file.table_meta_file().as_ref(), PROP_KEY_MAX_TS, )?; let elapsed = TimeStamp::new(current_ts).duration_since(max_ts.into()); let target_sc = spec.target_storage_class(elapsed); Ok(if target_sc.is_ia() { let auto_ia_file = IaAutoFile::new(ia_file, None, max_ts, spec.clone()); (Some(auto_ia_file), None) } else if let Ok(local_file) = self.open_local_file(id, fm.file_type, tiny_meta.file_size()) { let auto_ia_file = IaAutoFile::new(ia_file, Some(Arc::new(local_file)), max_ts, spec.clone()); (Some(auto_ia_file), None) } else { (None, Some(ia_file.table_meta_file().clone())) }) } async fn load_remote_auto_ia_file( tag: ShardTag, id: u64, fm: &FileMeta, opts: dfs::Options, fs: &dyn dfs::Dfs, io_type: IoType, table_meta_file: Option>, current_ts: u64, spec: StorageClassSpec, ) -> Result { debug_assert_eq!(fm.file_type, FileType::Sst); let table_meta_file = if let Some(table_meta_file) = table_meta_file { table_meta_file } else { let table_meta = Self::load_remote_table_meta(tag, id, fm, opts, fs, io_type).await?; Arc::new(InMemFile::new(id, table_meta.clone())) }; let SsTableProperty::MaxTs(max_ts) = SsTableCore::extract_property_from_table_meta_file( table_meta_file.as_ref(), PROP_KEY_MAX_TS, )?; let elapsed = TimeStamp::new(current_ts).duration_since(max_ts.into()); let target_sc = spec.target_storage_class(elapsed); let table_data = if target_sc.is_ia() { None } else { Some(Self::load_remote_table(id, fm.file_type, opts, fs, io_type).await?) }; Ok(PreparedFileResult::AutoIa { table_meta_file, local_data: table_data, max_ts, spec, }) } fn save_and_open_auto_ia_file( &self, id: u64, fm: &FileMeta, prepared: PreparedFileResult, permit: DfsLoadLimiterPermit, use_direct_io: bool, io_type: IoType, ) -> Result { let PreparedFileResult::AutoIa { table_meta_file, local_data, max_ts, spec, } = prepared else { unreachable!(); }; let ia_file = if table_meta_file.path().is_none() { let table_meta_data = table_meta_file.read_all()?; self.save_and_open_ia_file(id, fm, table_meta_data, None, use_direct_io, io_type)? } else { let IaCtx::Enabled(ia_mgr, _) = &self.ia_ctx else { unreachable!("ia_ctx should be enabled"); }; IaFile::open(id, fm, table_meta_file, ia_mgr.clone(), true)? }; let local_file: Option> = if let Some(local_data) = local_data { let local_file = self.save_and_open_table(id, fm, local_data, permit, use_direct_io, io_type)?; Some(Arc::new(local_file) as _) } else { drop(permit); None }; Ok(IaAutoFile::new(ia_file, local_file, max_ts, spec)) } /// Loads a schema file using the I/O class selected by the caller. pub async fn load_schema_file(&self, id: u64, io_type: IoType) -> Result { if let Some(schema_file) = self.schema_files.get(&id) { return Ok(schema_file.clone()); } let local_schema_file_path = self.local_schema_file_path(id); if let Ok(data) = fs::read(local_schema_file_path.as_path()) { let in_mem_file = InMemFile::new(id, data.into()); match SchemaFile::open(Arc::new(in_mem_file)) { Ok(schema_file) => { return match self.schema_files.entry(id) { Entry::Occupied(e) => Ok(e.get().clone()), Entry::Vacant(e) => { e.insert(schema_file.clone()); Ok(schema_file) } }; } Err(e) => { if should_remove_corrupted_schema_file(&e) { remove_corrupted_local_file(&local_schema_file_path, id, &e.to_string()); } else { return Err(e.into()); } } } } let opts = dfs::Options::default().with_type(FileType::Schema); let data = Self::read_remote_file(id, opts, self.fs.as_ref(), io_type).await?; let in_mem_file = InMemFile::new(id, data.clone()); let schema_file = SchemaFile::open(Arc::new(in_mem_file))?; match self.schema_files.entry(id) { Entry::Occupied(e) => return Ok(e.get().clone()), Entry::Vacant(e) => { e.insert(schema_file.clone()); } } // Only the first one will write the file to disk. self.write_local_file_with_file_path(id, data, false, io_type, &local_schema_file_path)?; Ok(schema_file) } pub(crate) fn local_file_path(&self, file_id: u64, file_type: FileType) -> PathBuf { match file_type { FileType::Sst => self.local_sst_file_path(file_id), FileType::Blob => self.local_blob_file_path(file_id), FileType::Schema => self.local_schema_file_path(file_id), FileType::Columnar => self.local_columnar_file_path(file_id), FileType::TxnChunk => panic!("TxnChunk files are managed by TxnChunkManager"), FileType::EncryptionDict => self.local_encryption_dict_file_path(file_id), FileType::VectorIndex => self.local_vector_index_file_path(file_id), FileType::FtsPackedFile => self.local_fts_packed_file_path(file_id), FileType::FtsDedicatedFile => self.local_fts_dedicated_file_path(file_id), } } pub(crate) fn local_sst_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_sst_filename(file_id)) } pub(crate) fn local_blob_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_blob_filename(file_id)) } pub(crate) fn local_schema_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_schema_filename(file_id)) } pub(crate) fn local_columnar_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_columnar_filename(file_id)) } pub(crate) fn local_encryption_dict_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_encryption_dict_filename(file_id)) } pub(crate) fn local_vector_index_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_vector_index_filename(file_id)) } pub(crate) fn local_fts_packed_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_fts_packed_filename(file_id)) } pub(crate) fn local_fts_dedicated_file_path(&self, file_id: u64) -> PathBuf { get_local_dir(&self.opts.local_dirs, file_id).join(new_fts_dedicated_filename(file_id)) } fn tmp_file_path(file_path: &Path) -> PathBuf { lazy_static::lazy_static! { static ref TMP_FILE_ID: AtomicU64 = AtomicU64::default(); } let tmp_id = TMP_FILE_ID.fetch_add(1, std::sync::atomic::Ordering::SeqCst); file_path.with_extension(format!("{tmp_id}.tmp")) } } pub fn collect_snap_lock_txn_file_refs(snap: &kvenginepb::Snapshot) -> Vec { let props = snap.get_properties(); debug_assert_eq!(props.keys.len(), props.values.len()); for (key, val) in props.get_keys().iter().zip(props.get_values().iter()) { if key == TXN_FILE_REF { let mut txn_file_refs = TxnFileRefs::new(); txn_file_refs.merge_from_bytes(val).unwrap(); return txn_file_refs .take_txn_file_refs() .into_iter() .filter(|r| !r.lock_val_prefix.is_empty()) .collect::>(); } } vec![] } // For safety. Table meta offset should have been fixed. See // `ShardMeta::fix_table_meta_offset`. fn validate_table_meta_off(tag: ShardTag, id: u64, fm: &FileMeta, data: Bytes) -> Result { debug_assert!(fm.can_use_ia(), "{}: {}: {:?}", tag, id, fm); if fm.table_meta_off > 0 { return Ok(data); } warn!("{} validate_table_meta_off: table meta offset is 0", tag; "id" => id, "fm" => ?fm); match fm.file_type { FileType::Sst => SsTable::get_meta_data(&data).map_err(|err| { error!("{} validate_table_meta_off: get meta data failed", tag; "err" => ?err, "id" => id, "fm" => ?fm); Error::TableError(err) }), FileType::Columnar => { // Columnar must have meta offset. debug_assert!(false); Ok(data) } FileType::VectorIndex => VectorIndexFile::get_meta_data(&data).map_err(|err| { error!("{} validate_table_meta_off: vector index must have meta offset", tag; "id" => id, "fm" => ?fm); Error::TableError(err) }), _ => Ok(data), } } fn should_remove_corrupted_local_file(err: &Error) -> bool { matches!( err, Error::TableError( TableError::InvalidChecksum(_) | TableError::InvalidMagicNumber | TableError::InvalidFileSize, ) ) } fn should_remove_corrupted_schema_file(err: &TableError) -> bool { matches!( err, TableError::InvalidChecksum(_) | TableError::InvalidMagicNumber | TableError::InvalidFileSize, ) } /// Remove a corrupted local cache file so the next load falls through to S3. fn remove_corrupted_local_file(path: &std::path::Path, file_id: u64, err: &str) { let res = fs::remove_file(path); warn!( "remove corrupted local file"; "file_id" => file_id, "path" => ?path, "error" => err, "remove_result" => ?res, ); } #[derive(Debug, Clone)] pub enum FilePrepareType { Local { transit_from_ia: bool }, Ia, AutoIa(StorageClassSpec), } impl Default for FilePrepareType { fn default() -> Self { Self::Local { transit_from_ia: false, } } } impl FilePrepareType { fn from_storage_class_spec(sc_spec: StorageClassSpec) -> Self { if sc_spec.must_be_ia() { Self::Ia } else if sc_spec.can_be_ia() { Self::AutoIa(sc_spec) } else { Self::Local { transit_from_ia: false, } } } pub fn from_shard_meta(meta: &ShardMeta) -> Self { Self::from_storage_class_spec(meta.get_storage_class_spec()) } pub fn from_snapshot(snap: &kvenginepb::Snapshot) -> Self { Self::from_storage_class_spec(StorageClassSpec::unmarshal( get_shard_property(STORAGE_CLASS_KEY, snap.get_properties()).as_deref(), )) } pub fn not_local(&self) -> bool { !self.is_local() } pub fn is_local(&self) -> bool { matches!(self, Self::Local { .. }) } pub fn transit_from_ia(&self) -> bool { matches!( self, Self::Local { transit_from_ia: true } ) } pub fn convert_for_prepare_table(&self, fm: &FileMeta, force_ia: bool) -> Self { match fm.file_type { FileType::Sst if fm.can_use_ia() => { if force_ia { Self::Ia } else { self.clone() } } FileType::Columnar => Self::Ia, FileType::VectorIndex => Self::Ia, FileType::FtsPackedFile => Self::Ia, FileType::FtsDedicatedFile => Self::Ia, _ => Self::Local { transit_from_ia: self.transit_from_ia(), }, } } } enum PreparedFileResult { Local { local_data: Bytes, }, Ia { table_meta_data: Bytes, }, AutoIa { table_meta_file: Arc, local_data: Option, max_ts: u64, spec: StorageClassSpec, }, } struct LoadRemoteResult { id: u64, fm: FileMeta, prepared: PreparedFileResult, permit: DfsLoadLimiterPermit, memory_guard: Option, } #[cfg(test)] mod tests { use std::{str::FromStr, sync::Arc, time::Duration}; use bytes::Bytes; use tidb_query_datatype::{Collation::Utf8Mb4GeneralCi, FieldTypeTp}; use tikv_util::config::ReadableSize; use super::*; use crate::{ ia::transition::TransitionLimiter, limiter::DfsLoadLimiter, metrics::{ ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES, ENGINE_MATERIALIZATION_REMOTE_READ_DURATION, }, table::{ columnar::{new_common_handle_column_info, new_version_column_info}, schema_file::{Schema, SchemaBuf, build_schema_file}, }, }; #[test] fn test_prepare_opts_default_to_other_io() { assert_eq!(PrepareOpts::default().io_type, IoType::Other); } #[test] fn test_remote_read_metrics_use_explicit_io_type() { let fs = MockDfs::new(2048); // Use a class that no other prepare test emits, so global metric deltas // remain deterministic when the test suite runs concurrently. let io_type = IoType::LoadBalance; let label = io_type.as_str(); let bytes = ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(); let reads = ENGINE_MATERIALIZATION_REMOTE_READ_DURATION .with_label_values(&[label, "success"]) .get_sample_count(); tokio::runtime::Runtime::new().unwrap().block_on(async { let data = EngineCore::load_remote_table( 1, FileType::Sst, dfs::Options::default(), &fs, io_type, ) .await .unwrap(); assert_eq!(data.len(), 2048); }); assert_eq!( ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(), bytes + 2048 ); assert_eq!( ENGINE_MATERIALIZATION_REMOTE_READ_DURATION .with_label_values(&[label, "success"]) .get_sample_count(), reads + 1 ); } #[test] fn test_remote_read_error_records_no_logical_bytes() { let fs = MockDfs::failing(); let io_type = IoType::Gc; let label = io_type.as_str(); let bytes = ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(); let errors = ENGINE_MATERIALIZATION_REMOTE_READ_DURATION .with_label_values(&[label, "error"]) .get_sample_count(); let result = tokio::runtime::Runtime::new() .unwrap() .block_on(EngineCore::load_remote_table( 1, FileType::Sst, dfs::Options::default(), &fs, io_type, )); result.unwrap_err(); assert_eq!( ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(), bytes ); assert_eq!( ENGINE_MATERIALIZATION_REMOTE_READ_DURATION .with_label_values(&[label, "error"]) .get_sample_count(), errors + 1 ); } #[test] fn test_load_file_from_remote_forwards_io_type() { let file_size = 1024; let fs: Arc = Arc::new(MockDfs::new(file_size)); let cfg = KvEngineConfig::default(); let opts = Options::default(); let limiter = DfsLoadLimiter::new(&cfg, &opts); let transition_limiter = TransitionLimiter::new(ReadableSize(0), 1, Duration::from_secs(1)); let (tx, rx) = tikv_util::mpsc::bounded(1); let io_type = IoType::Export; let label = io_type.as_str(); let bytes = ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(); tokio::runtime::Runtime::new() .unwrap() .block_on(EngineCore::load_file_from_remote( limiter, fs, 1, FileMeta { file_type: FileType::Sst, ..Default::default() }, io_type, FilePrepareType::Local { transit_from_ia: false, }, dfs::Options::default(), ShardTag::new(1, IdVer::new(1, 1)), None, None, transition_limiter, tx, )); rx.recv().unwrap().unwrap(); assert_eq!( ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(), bytes + file_size as u64 ); } #[test] fn test_load_schema_file_observes_remote_read() { let schema = Schema::new(SchemaBuf::new( 10, new_common_handle_column_info(), new_version_column_info(), vec![new_column_info(3, true)], vec![], 0, vec![], vec![], StorageClassSpec::default(), None, )); let file_data = build_schema_file(1, 100, vec![schema], 0); let file_size = file_data.len() as u64; let fs: Arc = Arc::new(MockDfs::with_data(file_data)); let cfg = KvEngineConfig::default(); let tmp_dir = tempfile::tempdir().unwrap(); let mut opts = Options::default(); opts.local_dirs = vec![tmp_dir.path().to_path_buf()]; let limiter = DfsLoadLimiter::new(&cfg, &opts); let core = EngineCore::new_for_test(fs, limiter, Arc::new(opts)); let io_type = IoType::Compaction; let label = io_type.as_str(); let bytes = ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(); let schema_file = tokio::runtime::Runtime::new() .unwrap() .block_on(core.load_schema_file(10, io_type)) .unwrap(); assert_eq!(schema_file.get_file_id(), 10); assert_eq!( ENGINE_MATERIALIZATION_REMOTE_LOGICAL_READ_BYTES .with_label_values(&[label]) .get(), bytes + file_size ); } #[test] fn test_tmp_file_path() { for p in ["/var/lib/tikv/db/100.sst", "/var/lib/tikv/db/100"] { let file_path = PathBuf::from(p); let tmp_file_path = EngineCore::tmp_file_path(&file_path); let path_str = tmp_file_path.to_str().unwrap(); let segs = path_str.split('.').collect::>(); assert_eq!(segs.len(), 3); assert_eq!(segs[0], "/var/lib/tikv/db/100"); u64::from_str(segs[1]).unwrap(); assert_eq!(segs[2], "tmp"); } } /// Test load_tables_by_ids with memory limiter using 2 concurrent threads /// This test verifies that memory limiter works correctly when multiple /// threads concurrently call load_tables_by_ids /// - 2KB memory cap /// - 2KB files with 1s sleep per file /// - 2 threads concurrently executing load_tables_by_ids, 2 files per /// thread /// - Should take at least 2s due to memory limitation #[cfg(feature = "failpoints")] #[test] fn test_load_tables_by_ids_memory_limiter() { fail::cfg("doing_recv_file_data", "sleep(500)").unwrap(); fail::cfg("return_recv_file_data", "return").unwrap(); // Create a temporary directory for the test let temp_dir = tempfile::TempDir::new().unwrap(); // Create mock components with very small memory limit let config = KvEngineConfig { dfs_load_memory_ratio: 0.25, // Very small ratio to get ~2KB limit ..Default::default() }; let mut opts = Options::default(); opts.local_dirs = vec![temp_dir.path().to_path_buf()]; opts.dfs_load_memory_limit = (config.dfs_load_memory_ratio * 10f64 * 1024f64) as u64; // 2KB memory cap let mock_fs = Arc::new(MockDfs::new(2048)); // 2KB files with 1s sleep let dfs_load_limiter = DfsLoadLimiter::new(&config, &opts); // Create Arc that will be held until test ends let opts_arc = Arc::new(opts); // Create EngineCore for testing let engine_core = Arc::new(EngineCore::new_for_test( mock_fs.clone(), dfs_load_limiter, opts_arc, )); // Create file metadata for testing - 4 files total, 2 per thread let create_file_meta = |start_id: u64, count: usize| { let mut file_ids = HashMap::new(); for i in 0..count { let id = start_id + i as u64; file_ids.insert( id, FileMeta { file_type: FileType::Sst, table_meta_off: 0, ..Default::default() }, ); } file_ids }; let file_ids_1 = create_file_meta(1, 2); // Files 1, 2 let file_ids_2 = create_file_meta(3, 2); // Files 3, 4 let start_time = std::time::Instant::now(); // Helper function to spawn a thread that calls load_tables_by_ids let spawn_loading_thread = |engine_core: Arc, shard_id: u64, file_ids: HashMap| { std::thread::spawn(move || { let rt = tokio::runtime::Runtime::new().unwrap(); rt.block_on(async { let mut cs = ChangeSet::new(kvenginepb::ChangeSet::default()); let tag = ShardTag::new(engine_core.get_engine_id(), IdVer::new(shard_id, 1)); engine_core.load_tables_by_ids( tag, shard_id, 1, // shard_ver &file_ids, &mut cs, false, // use_direct_io IoType::Replication, FilePrepareType::Local { transit_from_ia: false, }, EncryptionState::Disabled, None, // meta_pack_reader ) }) }) }; // Spawn two concurrent threads let handle1 = spawn_loading_thread(engine_core.clone(), 1, file_ids_1); let handle2 = spawn_loading_thread(engine_core.clone(), 2, file_ids_2); // Wait for both threads to complete let result1 = handle1.join().unwrap(); let result2 = handle2.join().unwrap(); assert!(result1.is_ok(), "result1 is not ok. result1: {:?}", result1); assert!(result2.is_ok(), "result2 is not ok. result2: {:?}", result2); let elapsed = start_time.elapsed(); // Verify memory limiting worked by checking timing // With 2KB memory cap and 2KB files, only 1 file can be loaded at a time // With 4 files total (2 per thread) and 1s sleep per file, it should take at // least 2s because the memory limiter will serialize the file loading // across both threads assert!( elapsed >= Duration::from_secs(2), "Should take at least 2s due to memory limiting and concurrent execution, elapsed: {:?}", elapsed ); } // Mock DFS implementation for testing struct MockDfs { file_data: Bytes, fail_read: bool, } impl MockDfs { #[allow(unused)] fn new(file_size: usize) -> Self { Self { file_data: Bytes::from(vec![0u8; file_size]), fail_read: false, } } fn with_data(file_data: Vec) -> Self { Self { file_data: file_data.into(), fail_read: false, } } fn failing() -> Self { Self { file_data: Bytes::new(), fail_read: true, } } } #[async_trait::async_trait] impl dfs::Dfs for MockDfs { fn get_prefix(&self) -> String { "mock_dfs".to_string() } async fn size(&self, _file_id: u64, _opts: dfs::Options) -> dfs::Result { Ok(self.file_data.len() as u64) } async fn read_file(&self, _file_id: u64, _opts: dfs::Options) -> dfs::Result { if self.fail_read { return Err(dfs::Error::Other("injected read error".to_owned())); } Ok(self.file_data.clone()) } async fn exists(&self, _file_id: u64, _opts: dfs::Options) -> dfs::Result { Ok(true) } async fn create( &self, _file_id: u64, _data: Bytes, _opts: dfs::Options, ) -> dfs::Result<()> { Ok(()) } async fn remove(&self, _file_id: u64, _file_len: Option, _opts: dfs::Options) { // Mock implementation - just return } async fn permanently_remove(&self, _file_id: u64, _opts: dfs::Options) -> dfs::Result<()> { Ok(()) } fn get_runtime(&self) -> &tokio::runtime::Runtime { // Return a reference to the current runtime // This is a simplified implementation for testing std::thread_local! { static RUNTIME: std::cell::RefCell> = const {std::cell::RefCell::new(None)}; } RUNTIME.with(|rt| { let mut runtime = rt.borrow_mut(); if runtime.is_none() { *runtime = Some(tokio::runtime::Runtime::new().unwrap()); } unsafe { &*(runtime.as_ref().unwrap() as *const tokio::runtime::Runtime) } }) } } fn new_column_info(id: i64, is_int: bool) -> tipb::ColumnInfo { let mut ci = tipb::ColumnInfo::new(); ci.set_column_id(id); if is_int { ci.set_tp(FieldTypeTp::LongLong as i32); } else { ci.set_tp(FieldTypeTp::VarChar as i32); ci.set_collation(Utf8Mb4GeneralCi as i32); } ci } #[test] fn test_corrupted_schema_file_is_deleted() { // Build a valid schema file. let schema = Schema::new(SchemaBuf::new( 10, new_common_handle_column_info(), new_version_column_info(), vec![new_column_info(3, true)], vec![], 0, vec![], vec![], StorageClassSpec::default(), None, )); let mut data = build_schema_file(1, 100, vec![schema], 0); // Corrupt the data by flipping a byte in the body (before the footer). data[0] ^= 0xFF; // Write the corrupted file to a temp directory. let tmp_dir = tempfile::tempdir().unwrap(); let file_path = tmp_dir.path().join("schema_100.schm"); fs::write(&file_path, &data).unwrap(); assert!(file_path.exists()); // Simulate load_schema_file logic: read from disk, SchemaFile::open fails. let read_data = fs::read(&file_path).unwrap(); let in_mem_file = InMemFile::new(100, read_data.into()); let result = SchemaFile::open(Arc::new(in_mem_file)); // The fix: delete the corrupted file on error. let err = match result { Err(e) => e, Ok(_) => panic!("SchemaFile::open should fail on corrupted data"), }; assert!( matches!(err, TableError::InvalidChecksum(_)), "SchemaFile::open should fail with InvalidChecksum" ); remove_corrupted_local_file(&file_path, 100, &err.to_string()); // Verify the corrupted file is deleted from disk. assert!( !file_path.exists(), "corrupted file should be deleted after checksum error" ); } #[test] fn test_should_remove_local_file_only_for_corruption_errors() { let checksum_error = Error::TableError(TableError::InvalidChecksum("bad checksum".to_owned())); assert!(should_remove_corrupted_local_file(&checksum_error)); let magic_error = Error::TableError(TableError::InvalidMagicNumber); assert!(should_remove_corrupted_local_file(&magic_error)); let file_size_error = Error::TableError(TableError::InvalidFileSize); assert!(should_remove_corrupted_local_file(&file_size_error)); let io_error = Error::TableError(TableError::Io("not a corruption error".to_owned())); assert!(!should_remove_corrupted_local_file(&io_error)); let apply_error = Error::ApplyChangeSet("unexpected apply error".to_owned()); assert!(!should_remove_corrupted_local_file(&apply_error)); } #[test] fn test_should_remove_schema_file_only_for_corruption_errors() { let checksum_error = TableError::InvalidChecksum("schema checksum mismatch".to_owned()); assert!(should_remove_corrupted_schema_file(&checksum_error)); let magic_error = TableError::InvalidMagicNumber; assert!(should_remove_corrupted_schema_file(&magic_error)); let file_size_error = TableError::InvalidFileSize; assert!(should_remove_corrupted_schema_file(&file_size_error)); let io_error = TableError::Io("schema read io error".to_owned()); assert!(!should_remove_corrupted_schema_file(&io_error)); } }