// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0. // Modifications Copyright 2022 PingCAP, Inc. mod test_columnar; mod test_fts; mod test_gc_lock_extra_cf; mod test_ia_auto_file; mod test_ia_file; mod test_storage_class; mod test_tombstone_triggered_compaction; mod test_txn_file; use std::{ collections::{HashMap, HashSet}, iter::Iterator, ops::Deref, rc::Rc, sync::Arc, thread, time::Duration, vec, }; use api_version::ApiV2; use bytes::{Buf, Bytes}; use cloud_encryption::{EncryptionKey, EncryptionState, MasterKey}; use dashmap::DashMap; use rand::prelude::*; use rstest::rstest; use tikv_util::{config::ReadableDuration, mpsc, time::Instant}; use txn_types::{ClusterGcStates, DEFAULT_KEYSPACE_ID, GcState, NULL_KEYSPACE_ID}; pub use crate::test_engine::*; use crate::{ shard::ShardDataBuilder, table::{ BIT_DELETE, BoundedDataSet, ChecksumType, DataBound, InnerKey, OP_PUT, file::{File, InMemFile}, memtable::{CfTable, MemTableImplementation}, sstable::{BlockCache, L0Builder, L0Table, SsTable}, }, tests::test_txn_file::{build_txn_chunk, make_lock_prefix, make_txn_file_refs}, *, }; const EXTRA_CF_KVENGINE_FUZZ_DURATION_ENV: &str = "CSE_EXTRA_CF_KVENGINE_FUZZ_DURATION"; const EXTRA_CF_KVENGINE_DEFAULT_ROUNDS: usize = 8; const EXTRA_CF_KVENGINE_CASES_PER_ROUND: usize = 64; const EXTRA_CF_KVENGINE_REPORT_INTERVAL: Duration = Duration::from_secs(10); // NOTE: In async test, access underlying file with async methods is not // checked, as `load_data` can only generate tables with `LocalFile`. // TODO: generate sstables with `IaFile`. #[maybe_async::test] async fn test_engine() { ::test_util::init_log_for_test(); for implementation in [ MemTableImplementation::Arena, MemTableImplementation::Crossbeam, ] { test_engine_with_memtable_implementation(implementation).await; } } #[maybe_async::both] async fn test_engine_with_memtable_implementation(implementation: MemTableImplementation) { let (engine, applier_tx) = new_test_engine_opt_with_custom_options(false, DEF_BLOCK_SIZE, "", |options| { options.enable_crossbeam_memtable = implementation == MemTableImplementation::Crossbeam; }); assert_eq!( engine .get_shard(1) .unwrap() .get_writable_mem_table() .implementation(), implementation ); // FIXME(youjiali1995): split has bugs. // // let mut keys = vec![]; // for i in &[1000, 3000, 6000, 9000] { // keys.push(i_to_key(*i, // engine.opts.blob_table_build_options.min_blob_size)); } // let mut splitter = Splitter::new(keys.clone(), applier_tx.clone()); // let handle = thread::spawn(move || { // splitter.run(); // }); let (begin, end) = (0, 1000); load_data( begin, end, 1, applier_tx, engine.opts.blob_table_build_options.min_blob_size, ); // handle.join().unwrap(); wait_for_compact_to_level_1_plus(&engine, begin, Duration::from_secs(10)); check_get( begin, end, 2, &[0, 1, 2], &engine, true, None, engine.opts.blob_table_build_options.min_blob_size, ) .await; check_iterator(begin, end, &engine).await; } #[maybe_async::test] async fn test_large_value_survives_memtable_flush() { const OLD_ARENA_LIMIT: usize = 16 * 1024 * 1024; ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_opt_with_custom_options(false, DEF_BLOCK_SIZE, "", |options| { options.enable_crossbeam_memtable = true; }); let initial_mem_table_count = engine.get_shard_stat(1).mem_table_count; let key = b"large-value"; let value = vec![b'x'; OLD_ARENA_LIMIT + 1]; let versioned_key = b"versioned-value"; let deleted_key = b"deleted-value"; let mut batch = WriteBatch::new(1); batch.put(WRITE_CF, key, &value, 0, &[], 1); batch.put(WRITE_CF, versioned_key, b"version-1", 0, &[], 1); batch.put(WRITE_CF, deleted_key, b"before-delete", 0, &[], 1); write_data(batch, &applier_tx); let mut batch = WriteBatch::new(1); batch.put(WRITE_CF, versioned_key, b"version-2", 0, &[], 2); batch.put(WRITE_CF, deleted_key, &[], BIT_DELETE, &[], 2); write_data(batch, &applier_tx); let shard = engine.get_shard(1).unwrap(); assert_eq!( shard.get_writable_mem_table().implementation(), crate::table::memtable::MemTableImplementation::Crossbeam ); { let snap = shard.new_snap_access(); let item = snap.get(WRITE_CF, key, u64::MAX).await; assert_eq!(item.get_value(), value); } switch_mem_table(&applier_tx); assert!(try_wait( || engine.get_shard_stat(1).mem_table_count == initial_mem_table_count, 10, )); assert_eq!( shard.get_writable_mem_table().implementation(), crate::table::memtable::MemTableImplementation::Crossbeam ); let snap = shard.new_snap_access(); let item = snap.get(WRITE_CF, key, u64::MAX).await; assert_eq!(item.get_value(), value); let item = snap.get(WRITE_CF, versioned_key, 1).await; assert_eq!(item.get_value(), b"version-1"); let item = snap.get(WRITE_CF, versioned_key, u64::MAX).await; assert_eq!(item.get_value(), b"version-2"); let item = snap.get(WRITE_CF, deleted_key, 1).await; assert_eq!(item.get_value(), b"before-delete"); let item = snap.get(WRITE_CF, deleted_key, u64::MAX).await; assert!(item.is_deleted()); } #[test] fn test_default_engine_uses_crossbeam_memtable() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); assert_eq!( shard.get_writable_mem_table().implementation(), crate::table::memtable::MemTableImplementation::Crossbeam ); } #[test] fn test_destroy_range() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine(); let mem_table_count = engine.get_shard_stat(1).mem_table_count; load_data( 10, 50, 1, applier_tx.clone(), engine.opts.blob_table_build_options.min_blob_size, ); // Unsafe destroy keys [10, 30). for prefix in [10, 20] { let mut wb = WriteBatch::new(1); let key = i_to_key(prefix, engine.opts.blob_table_build_options.min_blob_size); wb.set_property(DEL_PREFIXES_KEY, &key.as_bytes()[..key.len() - 1]); write_data(wb, &applier_tx); } assert!(!engine.get_shard(1).unwrap().get_del_prefixes().is_empty()); // Memtable is switched because it contains data covered by the delete-prefixes. let stats = engine.get_shard_stat(1); assert_eq!( stats.mem_table_count + stats.l0_table_count + stats.blob_table_count, mem_table_count + 1 ); let wait_for_destroying_range = || { for _ in 0..30 { let shard = engine.get_shard(1).unwrap(); let data = shard.get_data(); let del_prefixes = shard.get_del_prefixes(); if del_prefixes.is_empty() { break; } if Shard::ready_to_destroy_range(&del_prefixes, &data) { engine.trigger_compact(shard.id_ver()); } thread::sleep(Duration::from_millis(100)); } let shard = engine.get_shard(1).unwrap(); let length = shard .get_property(DEL_PREFIXES_KEY) .map(|v| v.len()) .unwrap_or_default(); // Delete-prefixes is cleaned up. assert_eq!(length, 0); }; wait_for_destroying_range(); // After destroying range, key [10, 30) should be removed. check_get( 10, 30, 2, &[0, 1, 2], &engine, false, None, engine.opts.blob_table_build_options.min_blob_size, ); check_get( 30, 50, 2, &[0, 1, 2], &engine, true, None, engine.opts.blob_table_build_options.min_blob_size, ); check_iterator(30, 50, &engine); // Trigger L0 compaction. for i in 1..=10 { load_data( 50 + (i - 1) * 10, 50 + i * 10, 1, applier_tx.clone(), engine.opts.blob_table_build_options.min_blob_size, ); let mut wb = WriteBatch::new(1); wb.set_switch_mem_table(); write_data(wb, &applier_tx); } // Waiting for L0 compaction. for _ in 0..30 { if engine.get_shard_stat(1).l0_table_count == 0 { break; } thread::sleep(Duration::from_millis(100)); } assert!(engine.get_shard_stat(1).l0_table_count < 10); // Unsafe destroy keys [100, 150). let mut wb = WriteBatch::new(1); let key = i_to_key(100, engine.opts.blob_table_build_options.min_blob_size); wb.set_property(DEL_PREFIXES_KEY, &key.as_bytes()[..key.len() - 2]); write_data(wb, &applier_tx); wait_for_destroying_range(); check_get( 100, 150, 2, &[0, 1, 2], &engine, false, None, engine.opts.blob_table_build_options.min_blob_size, ); check_get( 50, 100, 2, &[0, 1, 2], &engine, true, None, engine.opts.blob_table_build_options.min_blob_size, ); // Clean all data. let mut wb = WriteBatch::new(1); wb.set_property(DEL_PREFIXES_KEY, b"key"); write_data(wb, &applier_tx); wait_for_destroying_range(); check_get( 10, 150, 2, &[0, 1, 2], &engine, false, None, engine.opts.blob_table_build_options.min_blob_size, ); // No data exists and delete-prefixes can be cleaned too. let mut wb = WriteBatch::new(1); wb.set_property(DEL_PREFIXES_KEY, b"key"); write_data(wb, &applier_tx); wait_for_destroying_range(); } #[test] fn test_get_gc_safe_point_handles_keyspace_level_and_unified_gc() { ::test_util::init_log_for_test(); let ks_safe_point_v2_cache = Arc::new(DashMap::new()); ks_safe_point_v2_cache.insert(7, 107); let (engine, _) = new_test_engine_opt_with_legacy_gc_safe_point_cache( false, DEF_BLOCK_SIZE, "", ks_safe_point_v2_cache, |_| {}, ); #[allow(deprecated)] engine.update_managed_safe_ts(90); let now = Instant::now(); let mut cluster_gc_states = ClusterGcStates::new(HashMap::default(), now); cluster_gc_states.keyspace_gc_states.insert( NULL_KEYSPACE_ID, GcState::new(NULL_KEYSPACE_ID, false, 100.into(), 100.into(), vec![], now), ); cluster_gc_states.keyspace_gc_states.insert( 1, GcState::new(1, true, 101.into(), 101.into(), vec![], now), ); cluster_gc_states.keyspace_gc_states.insert( 2, GcState::new(2, false, 102.into(), 102.into(), vec![], now), ); engine.update_cluster_gc_states_cache(cluster_gc_states); assert_eq!(engine.get_gc_safe_point(NULL_KEYSPACE_ID).into_inner(), 100); assert_eq!( engine.get_gc_safe_point(DEFAULT_KEYSPACE_ID).into_inner(), 100 ); assert_eq!(engine.get_gc_safe_point(1).into_inner(), 101); assert_eq!(engine.get_gc_safe_point(2).into_inner(), 100); assert_eq!(engine.get_gc_safe_point(7).into_inner(), 0); assert_eq!(engine.get_gc_safe_point(8).into_inner(), 0); } #[test] fn test_new_shard_version_keeps_store_usage_stats() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine(); load_data( 0, 1000, 1, applier_tx, engine.opts.blob_table_build_options.min_blob_size, ); wait_for_compact_to_level_1_plus(&engine, 0, Duration::from_secs(10)); let old_stats = engine.collect_store_usage_stats(); assert!( old_stats.row_size > 0, "expected persisted data before cloning shard version: {:?}", old_stats ); let shard = engine.get_shard(1).unwrap(); let new_shard = engine.new_shard_version(&shard, shard.get_write_sequence() + 1); engine.insert_shard(Arc::new(new_shard)); let new_stats = engine.collect_store_usage_stats(); assert_eq!(new_stats.mem_tables_size, old_stats.mem_tables_size); assert_eq!( new_stats.mem_tables_memory_usage, old_stats.mem_tables_memory_usage ); assert_eq!(new_stats.row_size, old_stats.row_size); assert_eq!(new_stats.leader_row_size, old_stats.leader_row_size); assert_eq!( new_stats.leader_columnar_size, old_stats.leader_columnar_size ); assert_eq!(new_stats.leader_fts_size, old_stats.leader_fts_size); assert_eq!(new_stats.leader_vector_size, old_stats.leader_vector_size); assert_eq!(new_stats.ia, old_stats.ia); assert_eq!(new_stats.cf_level_sizes, old_stats.cf_level_sizes); } #[test] fn test_shard_registry_reclaims_single_replaced_shard() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let old_shard = engine.get_shard(1).unwrap(); let old_shard_ref = Arc::downgrade(&old_shard); let new_shard = engine.new_shard_version(&old_shard, old_shard.get_write_sequence() + 1); // The load worker sends its completion token before dropping the loaded shard. for _ in 0..100_000 { if Arc::strong_count(&old_shard) == 2 { break; } thread::yield_now(); } assert_eq!(Arc::strong_count(&old_shard), 2); let replaced_shard = engine.insert_shard(Arc::new(new_shard)); drop(old_shard); drop(replaced_shard); assert!( old_shard_ref.upgrade().is_none(), "a replaced shard should not wait for another registry retirement" ); } #[test] fn test_collect_store_usage_stats_matches_full_stats() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine(); load_data( 0, 1000, 1, applier_tx.clone(), engine.opts.blob_table_build_options.min_blob_size, ); wait_for_compact_to_level_1_plus(&engine, 0, Duration::from_secs(10)); load_data( 1000, 1050, 2, applier_tx, engine.opts.blob_table_build_options.min_blob_size, ); let ok = try_wait( || { let store_usage_stats = engine.collect_store_usage_stats(); let all_shard_stats = engine.get_all_shard_stats(); let mut leader_row = 0u64; let mut leader_columnar = 0u64; let mut leader_fts = 0u64; let mut leader_vector = 0u64; for shard in &all_shard_stats { if !shard.active { continue; } leader_row += shard.l0_table_size; for cf in &shard.cfs { for level in &cf.levels { leader_row += level.data_size; } } leader_columnar += shard .columnar_levels .iter() .map(|l| l.data_size) .sum::(); leader_fts += shard.fts_levels.iter().map(|l| l.data_size).sum::(); leader_vector += shard.vector_indexes.data_size; } let engine_stats = engine.get_engine_stats(all_shard_stats); store_usage_stats.mem_tables_size == engine_stats.mem_tables_size && store_usage_stats.row_size + engine_stats .columnar_levels .iter() .map(|l| l.data_size) .sum::() + engine_stats .fts_levels .iter() .map(|l| l.data_size) .sum::() + engine_stats.vector_indexes.data_size == engine_stats.total_data_size && store_usage_stats.leader_row_size == leader_row && store_usage_stats.leader_columnar_size == leader_columnar && store_usage_stats.leader_fts_size == leader_fts && store_usage_stats.leader_vector_size == leader_vector && store_usage_stats.ia == engine_stats.ia && (0..NUM_CFS).all(|cf| { store_usage_stats.cf_level_sizes[cf][..CF_LEVELS[cf]] == engine_stats.cf_level_sizes[cf].as_slice()[..] }) }, 10, ); assert!(ok, "store usage stats should match full engine stats"); } #[test] fn test_get_gc_safe_point_uses_legacy_cache_only_after_fallback() { ::test_util::init_log_for_test(); let ks_safe_point_v2_cache = Arc::new(DashMap::new()); ks_safe_point_v2_cache.insert(7, 107); let (engine, _) = new_test_engine_opt_with_legacy_gc_safe_point_cache( false, DEF_BLOCK_SIZE, "", ks_safe_point_v2_cache, |_| {}, ); #[allow(deprecated)] engine.update_managed_safe_ts(90); assert_eq!(engine.get_gc_safe_point(NULL_KEYSPACE_ID).into_inner(), 0); assert_eq!(engine.get_gc_safe_point(7).into_inner(), 0); assert_eq!(engine.get_gc_safe_point(8).into_inner(), 0); engine.set_fallback_to_legacy_gc_api(true); assert_eq!(engine.get_gc_safe_point(NULL_KEYSPACE_ID).into_inner(), 90); assert_eq!(engine.get_gc_safe_point(7).into_inner(), 107); assert_eq!(engine.get_gc_safe_point(8).into_inner(), 90); } #[test] fn test_get_gc_safe_point_handles_missing_null_keyspace_state() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); #[allow(deprecated)] engine.update_managed_safe_ts(90); let now = Instant::now(); let mut cluster_gc_states = ClusterGcStates::new(HashMap::default(), now); cluster_gc_states.keyspace_gc_states.insert( DEFAULT_KEYSPACE_ID, GcState::new( DEFAULT_KEYSPACE_ID, false, 101.into(), 101.into(), vec![], now, ), ); cluster_gc_states.keyspace_gc_states.insert( 2, GcState::new(2, false, 102.into(), 102.into(), vec![], now), ); cluster_gc_states.keyspace_gc_states.insert( 3, GcState::new(3, true, 103.into(), 103.into(), vec![], now), ); engine.update_cluster_gc_states_cache(cluster_gc_states); assert_eq!( engine.get_gc_safe_point(DEFAULT_KEYSPACE_ID).into_inner(), 101 ); assert_eq!(engine.get_gc_safe_point(2).into_inner(), 102); assert_eq!(engine.get_gc_safe_point(3).into_inner(), 103); assert_eq!(engine.get_gc_safe_point(8).into_inner(), 0); } // This test case construct a LSM tree and trigger a particular compaction to // verify deleted entry will not reappear after the compaction. // In the LSM true, L3 contains data, L2 contains tombstone, L1 contains data. // The range for overlap check should use both L1 and L2 instead of only L1. // The tombstone entries will be discarded when the compaction is // non-overlapping. If only use L1's range for overlap check, some of the L2's // tombstone will be lost, cause deleted entries reappear. #[test] fn test_lost_tombstone_issue() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let mut cf_builder = ShardCfBuilder::new(0); let mut saved_vals: Vec>> = Vec::new(); cf_builder.add_table( new_table(&engine, 11, 0, 100, 101, false, &mut saved_vals), 3, ); cf_builder.add_table( new_table(&engine, 12, 50, 150, 102, true, &mut saved_vals), 2, ); cf_builder.add_table( new_table(&engine, 13, 120, 200, 103, false, &mut saved_vals), 1, ); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); shard.set_data(builder.build()); let pri = CompactionPriority::L1Plus { cf: 0, score: 2.0, level: 1, file_gc: false, }; let mut guard = shard.compaction_priority.write().unwrap(); *guard = Some(pri); drop(guard); let now = Instant::now(); let mut cluster_gc_states = ClusterGcStates::new(HashMap::default(), now); cluster_gc_states.keyspace_gc_states.insert( NULL_KEYSPACE_ID, GcState::new(NULL_KEYSPACE_ID, false, 104.into(), 104.into(), vec![], now), ); engine.update_cluster_gc_states_cache(cluster_gc_states); engine.trigger_compact(IdVer::new(1, 1)); thread::sleep(Duration::from_secs(1)); check_get(50, 100, 104, &[0], &engine, false, None, 0); } // This test case construct a three level LSM tree and with different version // and add tombstone with latest version. Iterator should return all versions of // the key, excluding the tombstoned key. #[maybe_async::test] async fn test_read_iterator_all_versions() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let cf = 0; let mut cf_builder = ShardCfBuilder::new(cf); let mut saved_vals: Vec>> = Vec::new(); // 0..50 has only version 101 // 50..70 has version 101 and 102 // 70..90 has version 103 marked deleted, and older version 101 and 102 // 90..100 has version 101 and 102 // 100..150 has version 102 cf_builder.add_table( new_table(&engine, 11, 0, 100, 101, false, &mut saved_vals).await, 3, ); cf_builder.add_table( new_table(&engine, 12, 50, 150, 102, false, &mut saved_vals).await, 2, ); cf_builder.add_table( new_table(&engine, 13, 70, 90, 103, true, &mut saved_vals).await, 1, ); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(cf, false, true, None, false).await; iter.seek(shard.outer_start.chunk()).await; let mut expected_keys = Vec::with_capacity(160); for i in 0..50 { expected_keys.push(i); } for i in 50..70 { expected_keys.push(i); expected_keys.push(i); } // 70..90 should not returned for i in 90..100 { expected_keys.push(i); expected_keys.push(i); } for i in 100..150 { expected_keys.push(i); } let mut expected_keys_iter = expected_keys.iter(); while iter.valid() { let key = engine .key_builder .i_to_outer_key(expected_keys_iter.next().unwrap().to_owned()); assert_eq!(iter.key(), key.as_slice()); assert_eq!(iter.val(), key.repeat(2).as_slice()); iter.next().await; } } // Basic sanity for SnapAccess::new_delta_iterator over WRITE_CF. // Build two versions of the same key range and verify: // - since_ts = 0 returns both versions (2x keys) // - since_ts > older_version returns only the newer version (1x keys) // - at the equality boundary, stable sources are pruned because WRITE_CF delta // scanning has no key bound. #[test] fn test_delta_iterator_basic() { let (engine, _) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); // Build write CF with two non-deleted versions for keys [0, 50). let cf = WRITE_CF; let mut cf_builder = ShardCfBuilder::new(cf); let mut saved_vals: Vec>> = Vec::new(); // Older version = 100, Newer version = 200. cf_builder.add_table( new_table(&engine, 1001, 0, 50, 100, false, &mut saved_vals), 3, ); cf_builder.add_table( new_table(&engine, 1002, 0, 50, 200, false, &mut saved_vals), 2, ); // Wire into shard data (other CFs empty). let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); shard.set_data(builder.build()); // Create snapshot and set a very large managed_ts so parse_item doesn't filter // by read_ts. let snap = SnapAccess::new(&shard); let lower = Bytes::from(snap.get_start_key().to_vec()); let upper = Bytes::from(snap.get_end_key().to_vec()); // since_ts = 0 -> expect two versions per key let mut it = snap.new_delta_iterator(false, 0); it.set_range(lower.clone(), upper.clone()); let mut cnt = 0; while it.valid() { cnt += 1; it.next(); } assert_eq!(cnt, 50 /* keys */ * 2 /* versions */); // since_ts = 150 -> expect only version 200 let mut it2 = snap.new_delta_iterator(false, 150); it2.set_range(lower.clone(), upper.clone()); let mut cnt2 = 0; while it2.valid() { cnt2 += 1; it2.next(); } assert_eq!(cnt2, 50); // since_ts = 100 -> expect only version 200 let mut it3 = snap.new_delta_iterator(false, 100); it3.set_range(lower.clone(), upper.clone()); let mut cnt3 = 0; while it3.valid() { cnt3 += 1; it3.next(); } assert_eq!(cnt3, 50); // since_ts = 200 -> expect no newer version let mut it4 = snap.new_delta_iterator(false, 200); it4.set_range(lower, upper); let mut cnt4 = 0; while it4.valid() { cnt4 += 1; it4.next(); } assert_eq!(cnt4, 0); } #[test] fn test_new_delta_write_iterator_filters_mixed_tables() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let since_ts = 150; let mut saved_vals: Vec>> = Vec::new(); let mut write_cf_builder = ShardCfBuilder::new(WRITE_CF); write_cf_builder.add_table(new_table(&engine, 11, 0, 5, 120, false, &mut saved_vals), 1); write_cf_builder.add_table( new_table(&engine, 12, 5, 10, 200, false, &mut saved_vals), 1, ); write_cf_builder.add_table( new_table(&engine, 13, 10, 15, 250, false, &mut saved_vals), 1, ); write_cf_builder.add_table( new_table_ext( &engine, 14, 15, 20, |i| if i < 18 { 140 } else { 160 }, false, &mut saved_vals, ), 1, ); write_cf_builder.add_table( new_table(&engine, 15, 20, 25, 210, false, &mut saved_vals), 1, ); write_cf_builder.add_table( new_table(&engine, 16, 20, 30, 150, false, &mut saved_vals), 2, ); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([write_cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_delta_write_iterator(since_ts); iter.rewind(); let mut expected = 5; while iter.valid() { let expected_key = engine.key_builder.i_to_inner_key(expected); assert_eq!(iter.key(), expected_key.as_ref()); let expected_version = if expected < 10 { 200 } else if expected < 15 { 250 } else if expected < 18 { 140 // caller should filter out this } else if expected < 20 { 160 } else { 210 }; assert_eq!(iter.value().version, expected_version); expected += 1; iter.next(); } assert_eq!(expected, 25); } #[test] fn test_iterator_range_clamped_to_snapshot_bounds() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key1, b"v1", 0, &[], 1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key2, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, false, false, Some(u64::MAX), true); iter.set_range(Bytes::from(key1.clone()), Bytes::from(key4.clone())); let mut keys = Vec::new(); while iter.valid() { keys.push(iter.key().to_vec()); iter.next(); } assert_eq!(keys, vec![key2, key3]); } #[test] fn test_iterator_range_outside_snapshot_upper_should_not_leak() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let key5 = engine.key_builder.i_to_outer_key(5); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key1, b"v1", 0, &[], 1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); wb.put(WRITE_CF, &key4, b"v4", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key2, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, false, false, Some(u64::MAX), true); iter.set_range(Bytes::from(key1.clone()), Bytes::from(key5.clone())); let mut keys = Vec::new(); while iter.valid() { keys.push(iter.key().to_vec()); iter.next(); } assert_eq!(keys, vec![key2, key3]); } #[test] fn test_iterator_seek_before_snapshot_start_should_not_leak() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key1, b"v1", 0, &[], 1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key2, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, false, false, Some(u64::MAX), true); iter.seek(&key1); assert!(iter.valid()); assert_eq!(iter.key(), key2.as_slice()); } #[test] fn test_reverse_iterator_seek_after_snapshot_end_should_not_leak() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let key5 = engine.key_builder.i_to_outer_key(5); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); wb.put(WRITE_CF, &key4, b"v4", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key2, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, true, false, Some(u64::MAX), true); iter.seek(&key5); assert!(iter.valid()); assert_eq!(iter.key(), key3.as_slice()); } #[test] fn test_reverse_iterator_seek_before_snapshot_start_should_be_invalid() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key2, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, true, false, Some(u64::MAX), true); iter.seek(&key1); assert!(!iter.valid()); } #[test] fn test_reverse_iterator_range_outside_snapshot_lower_should_not_leak() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let key5 = engine.key_builder.i_to_outer_key(5); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key1, b"v1", 0, &[], 1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); wb.put(WRITE_CF, &key4, b"v4", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key2, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, true, false, Some(u64::MAX), true); iter.set_range(Bytes::from(key1.clone()), Bytes::from(key5.clone())); let mut keys = Vec::new(); while iter.valid() { keys.push(iter.key().to_vec()); iter.next(); } assert_eq!(keys, vec![key3, key2]); } #[test] fn test_iterator_rewind_respects_range_forward() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key1, b"v1", 0, &[], 1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key1, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, false, false, Some(u64::MAX), true); iter.set_range(Bytes::from(key2.clone()), Bytes::from(key4.clone())); iter.rewind(); let mut keys = Vec::new(); while iter.valid() { keys.push(iter.key().to_vec()); iter.next(); } assert_eq!(keys, vec![key2, key3]); } #[test] fn test_iterator_rewind_respects_range_reverse() { ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_api_v2(); let shard = engine.get_shard(1).unwrap(); let key1 = engine.key_builder.i_to_outer_key(1); let key2 = engine.key_builder.i_to_outer_key(2); let key3 = engine.key_builder.i_to_outer_key(3); let key4 = engine.key_builder.i_to_outer_key(4); let mut wb = WriteBatch::new(1); wb.put(WRITE_CF, &key1, b"v1", 0, &[], 1); wb.put(WRITE_CF, &key2, b"v2", 0, &[], 1); wb.put(WRITE_CF, &key3, b"v3", 0, &[], 1); write_data(wb, &applier_tx); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(ShardRange::new(&key1, &key4)); shard.set_data(builder.build()); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(WRITE_CF, true, false, Some(u64::MAX), true); iter.set_range(Bytes::from(key2.clone()), Bytes::from(key4.clone())); iter.rewind(); let mut keys = Vec::new(); while iter.valid() { keys.push(iter.key().to_vec()); iter.next(); } assert_eq!(keys, vec![key3, key2]); } #[test] fn test_lock_cf_repeatable_read() { let enable_inner_key_off = true; ::test_util::init_log_for_test(); let (engine, applier_tx) = new_test_engine_opt(enable_inner_key_off, DEF_BLOCK_SIZE, TABLE_KEY_PREFIX); let kb = engine.key_builder(); let shard = engine.get_shard(1).unwrap(); let verify_locks = |snap: &SnapAccess, start: usize, end: usize, version: u64, del: bool| { let tag = format!("{start}->{end}, {version}, {del}"); let mut it = snap.new_iterator(LOCK_CF, false, false, None, false); it.seek(&kb.i_to_outer_key(start)); if del { assert!( !it.valid() || it.key() >= kb.i_to_outer_key(end).as_slice(), "{}", tag ); return; } for i in start..end { assert!(it.valid(), "{}", tag); assert_eq!(it.key(), kb.i_to_outer_key(i).as_slice(), "{}", tag); assert_eq!(it.version(), version, "{}", tag); assert_eq!(table::is_deleted(it.meta()), del, "{}", tag); it.next(); } }; // Prepare level data. { let mut cf_builder = ShardCfBuilder::new(LOCK_CF); let mut saved_vals: Vec>> = Vec::new(); cf_builder.add_table( new_table(&engine, 1000, 100, 300, 1000, false, &mut saved_vals), 2, ); cf_builder.add_table( new_table(&engine, 2000, 110, 300, 2000, true, &mut saved_vals), 1, ); let l0_file3 = new_l0table_file( &engine, 3000, [0, 120, 0], [0, 300, 0], 3000, [false, false, false], ); let l0_tbl3 = L0Table::new(l0_file3, BlockCache::None, false, EncryptionState::Disabled) .unwrap() .unwrap(); let l0_file4 = new_l0table_file( &engine, 4000, [0, 130, 0], [0, 300, 0], 4000, [false, true, false], ); let l0_tbl4 = L0Table::new(l0_file4, BlockCache::None, false, EncryptionState::Disabled) .unwrap() .unwrap(); let mut cfs = shard.get_data().cfs.clone(); cfs[LOCK_CF] = cf_builder.build(); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_l0_tbls(vec![l0_tbl4, l0_tbl3]); builder.set_cfs(cfs); shard.set_data(builder.build()); shard.set_base_version(10000); } let base_ver = shard.get_base_version(); // Write mem table. let mem_tbl_ver5 = load_data_ext( &engine, [0, 140, 0], [0, 300, 0], 5000, // useless [false, false, false], &applier_tx, ) + base_ver; // Write txn file lock. let txn_file_ver6 = { let chunk_id = 6000; build_txn_chunk(&engine, 150, 300, chunk_id, |_| OP_PUT, None); let primary = kb.i_to_outer_key(150); let mut wb = WriteBatch::new(1); let txn_file_refs = make_txn_file_refs( 6000, // useless vec![chunk_id], make_lock_prefix(primary.clone(), 6000), vec![], ); wb.set_property(TXN_FILE_REF, &txn_file_refs); engine .txn_chunk_mgr .prepare(chunk_id, EncryptionState::Disabled, IoType::Other) .unwrap(); write_data(wb, &applier_tx) } + base_ver; let snap = shard.new_snap_access(); print_locks(&snap, false, None); verify_locks(&snap, 100, 110, 1000, false); verify_locks(&snap, 110, 120, 2000, true); verify_locks(&snap, 120, 130, 3000, false); verify_locks(&snap, 130, 140, 4000, true); verify_locks(&snap, 140, 150, mem_tbl_ver5, false); verify_locks( &snap, 150, 300, snap.get_mem_table_snap_version().into_inner(), false, ); // Write more mem table data. { switch_mem_table(&applier_tx); load_data_ext( &engine, [0, 160, 0], [0, 300, 0], 7000, // useless [false, false, false], &applier_tx, ); } // Verify repeatable read. verify_locks(&snap, 140, 150, mem_tbl_ver5, false); verify_locks( &snap, 150, 300, snap.get_mem_table_snap_version().into_inner(), false, ); // Check latest snapshot. let snap = shard.new_snap_access(); verify_locks(&snap, 140, 150, mem_tbl_ver5, false); verify_locks(&snap, 150, 160, txn_file_ver6, false); verify_locks( &snap, 160, 300, snap.get_mem_table_snap_version().into_inner(), false, ); } #[rstest] #[case::enable_key_off(true)] #[case::disable_key_off(false)] fn test_level_overlapping_tables(#[case] enable_inner_key_off: bool) { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(enable_inner_key_off, DEF_BLOCK_SIZE, ""); let shard = engine.get_shard(1).unwrap(); let mut cf_builder = ShardCfBuilder::new(0); let mut saved_vals: Vec>> = Vec::new(); let table_ranges: Vec<(usize, usize)> = vec![(10, 20), (50, 100), (120, 200)]; for (i, (start, end)) in table_ranges.into_iter().enumerate() { cf_builder.add_table( new_table( &engine, i as u64, start, end, i as u64 + 100, false, &mut saved_vals, ), 1, // level ); } let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); let data = builder.build(); let cf0 = data.get_cf(0); let level1 = cf0.get_level(1); let get_overlapping_tables = |start: usize, end: usize| -> (usize, usize) { let inner_start = engine.key_builder.i_to_inner_key(start); let inner_end = engine.key_builder.i_to_inner_key(end); let bound = DataBound::new(inner_start.as_ref(), inner_end.as_ref(), false); bound.get_overlap_data_sets(&level1.tables) }; assert_eq!(get_overlapping_tables(0, 10), (0, 0)); assert_eq!(get_overlapping_tables(0, 20), (0, 1)); assert_eq!(get_overlapping_tables(0, 50), (0, 1)); assert_eq!(get_overlapping_tables(0, 100), (0, 2)); assert_eq!(get_overlapping_tables(0, 120), (0, 2)); assert_eq!(get_overlapping_tables(0, 200), (0, 3)); assert_eq!(get_overlapping_tables(0, 300), (0, 3)); assert_eq!(get_overlapping_tables(10, 20), (0, 1)); assert_eq!(get_overlapping_tables(10, 50), (0, 1)); assert_eq!(get_overlapping_tables(10, 100), (0, 2)); assert_eq!(get_overlapping_tables(10, 120), (0, 2)); assert_eq!(get_overlapping_tables(10, 200), (0, 3)); assert_eq!(get_overlapping_tables(10, 300), (0, 3)); assert_eq!(get_overlapping_tables(20, 50), (1, 1)); assert_eq!(get_overlapping_tables(20, 100), (1, 2)); assert_eq!(get_overlapping_tables(20, 120), (1, 2)); assert_eq!(get_overlapping_tables(20, 200), (1, 3)); assert_eq!(get_overlapping_tables(20, 300), (1, 3)); assert_eq!(get_overlapping_tables(50, 100), (1, 2)); assert_eq!(get_overlapping_tables(50, 120), (1, 2)); assert_eq!(get_overlapping_tables(50, 200), (1, 3)); assert_eq!(get_overlapping_tables(50, 300), (1, 3)); assert_eq!(get_overlapping_tables(100, 120), (2, 2)); assert_eq!(get_overlapping_tables(100, 200), (2, 3)); assert_eq!(get_overlapping_tables(100, 300), (2, 3)); assert_eq!(get_overlapping_tables(120, 200), (2, 3)); assert_eq!(get_overlapping_tables(120, 300), (2, 3)); assert_eq!(get_overlapping_tables(200, 300), (3, 3)); } #[rstest] #[case::enable_key_off(true)] #[case::disable_key_off(false)] fn test_get_suggest_split_key(#[case] enable_inner_key_off: bool) { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(enable_inner_key_off, 4096, ""); let shard = engine.get_shard(1).unwrap(); let range = ShardRange::new( &engine.key_builder.i_to_outer_key(30), &engine.key_builder.i_to_outer_key(600), ); let inner_key_off = shard.get_data().inner_key_off; let shard = Shard::new( shard.engine_id, &shard.properties.to_pb(shard.id), shard.ver, range.clone(), inner_key_off, engine.opts.clone(), &engine.master_key, engine.encryption_key_manager.clone(), ); let cases: Vec<( Vec<(usize, usize)>, Vec<(usize, usize)>, // table ranges // expected suggest split key, inner key off (enable,disable) Option, )> = vec![ (vec![], vec![(20, 50)], None), (vec![], vec![(20, 50), (100, 150), (150, 180)], Some(150)), ( vec![], vec![ (20, 50), (100, 150), // in range (150, 180), // in range (180, 300), // in range (300, 400), ], Some(180), ), ( vec![], vec![ (20, 50), (100, 150), // in range, key in block ], Some(100), ), ( vec![(50, 400), (60, 460), (70, 670)], // L0 only vec![], Some(320), ), ( vec![(70, 370)], // L0 + L1+ vec![(50, 80), (80, 100), (100, 120)], Some(100), ), ( vec![(80, 380), (190, 570)], // more L0 + L1+ vec![(50, 80), (80, 100), (100, 120)], Some(330), ), ]; let mut id_alloc = 1000; for (idx, (l0_ranges, table_ranges, expect_split_key)) in cases.into_iter().enumerate() { let mut l0_tables = vec![]; for (start, end) in l0_ranges { let id = id_alloc; id_alloc += 1; let (begins, ends) = if idx % 2 == 0 { ([start, 0, 0], [end, 0, 0]) } else { ([0, start, 0], [0, end, 0]) }; let l0_file = new_l0table_file(&engine, id, begins, ends, id, [false, false, false]); let l0_tbl = L0Table::new(l0_file, BlockCache::None, false, EncryptionState::Disabled) .unwrap() .unwrap(); l0_tables.push(l0_tbl); } let mut cf_builder = ShardCfBuilder::new(idx % 2); let mut saved_vals: Vec>> = Vec::new(); for (i, (start, end)) in table_ranges.into_iter().enumerate() { cf_builder.add_table( new_table( &engine, i as u64 + 1, start, end, 100 + i as u64, false, &mut saved_vals, ), 1, ); } let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(range.clone()); builder.set_l0_tbls(l0_tables); builder.set_cfs([cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); shard.set_data_opt(builder.build(), false); let key = shard.get_suggest_split_key(1); assert_eq!( key.map(|k| k.to_vec()), expect_split_key.map(|i| engine.key_builder.i_to_outer_key(i)), "case {}", idx ); } } #[rstest] #[case::enable_key_off(true)] #[case::disable_key_off(false)] fn test_get_evenly_split_keys(#[case] enable_inner_key_off: bool) { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(enable_inner_key_off, 4096, ""); let shard = engine.get_shard(1).unwrap(); let range = ShardRange::new( &engine.key_builder.i_to_outer_key(30), &engine.key_builder.i_to_outer_key(600), ); let inner_key_off = shard.get_data().inner_key_off; let shard = Shard::new( shard.engine_id, &shard.properties.to_pb(shard.id), shard.ver, range.clone(), inner_key_off, engine.opts.clone(), &engine.master_key, engine.encryption_key_manager.clone(), ); let cases: Vec<( Vec<(usize, usize)>, // l0 ranges Vec<(usize, usize)>, // table ranges usize, // split count // expected evenly split keys Option>, )> = vec![ (vec![], vec![(20, 50), (50, 100)], 1, None), (vec![], vec![(20, 50), (50, 100)], 2, None), ( vec![], vec![(0, 20), (20, 50), (100, 150), (150, 180)], 2, Some(vec![150]), ), ( vec![], vec![ (20, 50), (100, 150), // in range (150, 180), // in range (180, 300), // in range (300, 400), ], 1, None, ), ( vec![], vec![ (20, 50), (100, 150), // in range (150, 180), // in range (180, 300), // in range ], 2, Some(vec![180]), ), ( vec![], vec![ (100, 120), // in range (120, 150), // in range (150, 180), // in range (180, 300), // in range ], 3, Some(vec![150, 180]), ), ( vec![], vec![ (20, 50), (100, 150), // in range (150, 180), // in range (180, 300), // in range (300, 400), ], 4, Some(vec![150, 180, 300]), ), ( vec![], vec![ (20, 50), (100, 150), // in range (150, 180), // in range (180, 300), // in range (300, 400), ], 5, Some(vec![150, 180, 300]), ), ( vec![], vec![ (20, 50), (100, 150), // in range (150, 180), // in range (180, 300), // in range (300, 400), ], 6, Some(vec![150, 180, 300]), ), ( vec![(100, 400), (200, 500), (300, 600), (400, 700)], // L0 only vec![], 4, Some(vec![450, 550]), ), ( vec![(100, 400), (150, 550), (300, 700), (200, 800)], // L0 only vec![], 4, Some(vec![400, 450, 550]), ), ( vec![(100, 800), (200, 900)], // L0 + L1+ vec![(30, 80), (80, 100), (100, 130)], 4, Some(vec![100, 350, 450]), ), ]; let mut id_alloc = 1000; for (idx, (l0_ranges, table_ranges, split_count, expect_split_keys)) in cases.into_iter().enumerate() { let mut l0_tables = vec![]; for (start, end) in l0_ranges { let id = id_alloc; id_alloc += 1; let (begins, ends) = if idx % 2 == 0 { ([start, 0, 0], [end, 0, 0]) } else { ([0, start, 0], [0, end, 0]) }; let l0_file = new_l0table_file(&engine, id, begins, ends, id, [false, false, false]); let l0_tbl = L0Table::new(l0_file, BlockCache::None, false, EncryptionState::Disabled) .unwrap() .unwrap(); l0_tables.push(l0_tbl); } let mut cf_builder = ShardCfBuilder::new(idx % 2); let mut saved_vals: Vec>> = Vec::new(); for (i, (start, end)) in table_ranges.into_iter().enumerate() { cf_builder.add_table( new_table( &engine, i as u64 + 1, start, end, 100 + i as u64, false, &mut saved_vals, ), 1, // level ); } let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(range.clone()); builder.set_l0_tbls(l0_tables); builder.set_cfs([cf_builder.build(), ShardCf::new(1), ShardCf::new(2)]); shard.set_data_opt(builder.build(), false); let split_keys = shard.get_evenly_split_keys(split_count, 1); assert_eq!( split_keys.map(|keys| keys.into_iter().map(|k| k.to_vec()).collect::>()), expect_split_keys.map(|keys| keys .into_iter() .map(|i| engine.key_builder.i_to_outer_key(i)) .collect::>()), "case {}", idx ); } } fn new_approx_split_test_shard(engine: &TestEngine) -> (Shard, ShardRange) { let base_shard = engine.get_shard(1).unwrap(); let range = ShardRange::new( &engine.key_builder.i_to_outer_key(30), &engine.key_builder.i_to_outer_key(600), ); let shard = Shard::new( base_shard.engine_id, &base_shard.properties.to_pb(base_shard.id), base_shard.ver, range.clone(), base_shard.get_data().inner_key_off, engine.opts.clone(), &engine.master_key, engine.encryption_key_manager.clone(), ); (shard, range) } fn new_table_with_inner_keys(engine: &TestEngine, id: u64, keys: &[&[u8]]) -> SsTable { let mut builder = table::sstable::builder::Builder::new( id, engine.opts.table_builder_options.block_size, engine.opts.table_builder_options.compression_tps[0], engine.opts.table_builder_options.compression_lvl, ChecksumType::Crc32, None, ); for (version, key) in keys.iter().enumerate() { let value = table::Value::new_with_meta_version(0, version as u64 + 1, 0, b"value"); builder .add(InnerKey::from_inner_buf(key), &value, None) .unwrap(); } let mut data = Vec::new(); builder.finish(0, &mut data).unwrap(); SsTable::new( Arc::new(InMemFile::new(id, Bytes::from(data))), BlockCache::None, EncryptionState::Disabled, ) .unwrap() } #[test] fn test_approx_split_rejects_inconsistent_keyspace_metadata() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(true, 256, ""); let (mut shard, mut range) = new_approx_split_test_shard(&engine); assert!(range.keyspace_id > 0); range.outer_start = Bytes::new(); shard.range = range.clone(); let mut saved_vals: Vec>> = Vec::new(); let mut write_cf = ShardCfBuilder::new(WRITE_CF); write_cf.add_table( new_table(&engine, 10, 300, 400, 110, false, &mut saved_vals), 1, ); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(range); builder.set_cfs([ write_cf.build(), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data_opt(builder.build(), false); assert!(shard.approx_split(&[], &[]).unwrap().is_none()); } #[test] fn test_approx_split_rejects_multi_keyspace_range() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(true, 256, ""); let base_shard = engine.get_shard(1).unwrap(); let keyspace_1_prefix = ApiV2::get_txn_keyspace_prefix(1); let keyspace_2_prefix = ApiV2::get_txn_keyspace_prefix(2); let mut outer_start = keyspace_1_prefix.clone(); outer_start.extend_from_slice(b"a"); let mut outer_end = keyspace_2_prefix.clone(); outer_end.extend_from_slice(b"z"); let range = ShardRange::new(&outer_start, &outer_end); assert_eq!(range.keyspace_id, 1); assert_eq!(range.get_end_keyspace_id(), 2); // Across keyspaces, the outer prefixes determine the order while the // stored inner keys retain only the suffix. Use the opposite suffix order // to exercise the ambiguity that makes this range unsafe to split. let keyspace_1_z = [keyspace_1_prefix.as_slice(), b"z"].concat(); let keyspace_2_a = [keyspace_2_prefix.as_slice(), b"a"].concat(); assert!(keyspace_1_z < keyspace_2_a); assert!(InnerKey::from_outer_key(&keyspace_1_z) > InnerKey::from_outer_key(&keyspace_2_a)); let shard = Shard::new( base_shard.engine_id, &base_shard.properties.to_pb(base_shard.id), base_shard.ver, range, base_shard.get_data().inner_key_off, engine.opts.clone(), &engine.master_key, engine.encryption_key_manager.clone(), ); let mut write_cf = ShardCfBuilder::new(WRITE_CF); write_cf.add_table(new_table_with_inner_keys(&engine, 41, &[b"z"]), 1); write_cf.add_table(new_table_with_inner_keys(&engine, 42, &[b"a"]), 1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([ write_cf.build(), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data_opt(builder.build(), false); assert!(shard.approx_split(&[], &[]).unwrap().is_none()); } #[rstest] #[case::enable_key_off(true)] #[case::disable_key_off(false)] fn test_approx_split_prefers_sst_start_then_uses_largest_table_blocks( #[case] enable_inner_key_off: bool, ) { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(enable_inner_key_off, 256, ""); let (shard, range) = new_approx_split_test_shard(&engine); let outer_key = |i| engine.key_builder.i_to_outer_key(i); let mut saved_vals: Vec>> = Vec::new(); let mut write_cf = ShardCfBuilder::new(WRITE_CF); for (id, start, end) in [(11, 20, 300), (12, 300, 400), (13, 400, 550)] { write_cf.add_table( new_table(&engine, id, start, end, id + 100, false, &mut saved_vals), 1, ); } let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(range.clone()); builder.set_cfs([ write_cf.build(), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data_opt(builder.build(), false); match shard .approx_split(&outer_key(100), &outer_key(500)) .unwrap() .unwrap() { ApproxSplitKey::Sst(key) => assert_eq!( key.into_encoded(), txn_types::Key::from_raw(&outer_key(400)).into_encoded() ), ApproxSplitKey::Block(_) => panic!("SST start keys must be preferred"), } let large = new_table(&engine, 21, 20, 580, 121, false, &mut saved_vals); let small = new_table(&engine, 22, 50, 250, 122, false, &mut saved_vals); assert!(large.kv_size > small.kv_size); let effective_start = outer_key(100); let effective_end = outer_key(500); let effective_range = DataBound::new( InnerKey::from_outer_key(&effective_start), InnerKey::from_outer_end_key(&effective_end), false, ); let index = large.try_load_index(false).unwrap(); let first = index.seek_block(effective_range.lower_bound); let last = index.seek_block_bigger_or_equal(effective_range.upper_bound); assert!(first < last); let expected_inner_key = index.clone_block_key(first + (last - first) / 2); let expected_outer_key = range.to_outer_key(expected_inner_key.as_ref()); let mut write_cf = ShardCfBuilder::new(WRITE_CF); write_cf.add_table(small, 1); write_cf.add_table(large, 2); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([ write_cf.build(), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data_opt(builder.build(), false); match shard .approx_split(&outer_key(100), &outer_key(500)) .unwrap() .unwrap() { ApproxSplitKey::Block(key) => assert_eq!( key.into_encoded(), txn_types::Key::from_raw(&expected_outer_key).into_encoded() ), ApproxSplitKey::Sst(_) => panic!("no SST starts strictly inside the range"), } let mut unloaded = FileMeta::default(); unloaded.cf = WRITE_CF as i8; unloaded.level = 2; unloaded.smallest = Bytes::copy_from_slice(engine.key_builder.i_to_inner_key(80).as_ref().deref()); unloaded.biggest = Bytes::copy_from_slice(engine.key_builder.i_to_inner_key(520).as_ref().deref()); let mut unloaded_tables = HashMap::new(); unloaded_tables.insert(99, unloaded); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_unloaded_tbls(unloaded_tables); shard.set_data_opt(builder.build(), false); assert!( shard .approx_split(&outer_key(100), &outer_key(500)) .unwrap() .is_none() ); let unloaded_l0 = FileMeta::from_l0_table(&crate::util::new_l0_create_pb( 100, engine.key_builder.i_to_inner_key(80).to_vec(), engine.key_builder.i_to_inner_key(520).to_vec(), 1, )); assert_eq!(unloaded_l0.cf, -1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_unloaded_tbls(HashMap::from([(100, unloaded_l0)])); shard.set_data_opt(builder.build(), false); assert!( shard .approx_split(&outer_key(100), &outer_key(500)) .unwrap() .is_none() ); } #[test] fn test_approx_split_uses_loaded_l0() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(true, 256, ""); let (shard, range) = new_approx_split_test_shard(&engine); let outer_key = |i| engine.key_builder.i_to_outer_key(i); let mut begin = [0; NUM_CFS]; let mut end = [0; NUM_CFS]; begin[WRITE_CF] = 300; end[WRITE_CF] = 550; let file = new_l0table_file(&engine, 31, begin, end, 131, [false; NUM_CFS]); let l0 = L0Table::new(file, BlockCache::None, false, EncryptionState::Disabled) .unwrap() .unwrap(); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(range); builder.set_l0_tbls(vec![l0]); shard.set_data_opt(builder.build(), false); match shard .approx_split(&outer_key(100), &outer_key(500)) .unwrap() .unwrap() { ApproxSplitKey::Sst(key) => assert_eq!( key.into_encoded(), txn_types::Key::from_raw(&outer_key(300)).into_encoded() ), ApproxSplitKey::Block(_) => panic!("loaded L0 start key must participate"), } } #[test] fn test_approx_split_uses_ia_sst_index() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(true, 256, ""); let (shard, range) = new_approx_split_test_shard(&engine); let outer_key = |i| engine.key_builder.i_to_outer_key(i); let mut saved_vals: Vec>> = Vec::new(); let local = new_table(&engine, 32, 20, 580, 132, false, &mut saved_vals); let manager = engine.new_ia_manager( crate::ia::util::IaManagerOptionsBuilder::default() .build() .unwrap(), ); let ia = test_ia_file::convert_local_sst_to_ia(&local, manager); assert!(!ia.is_sync()); let mut write_cf = ShardCfBuilder::new(WRITE_CF); write_cf.add_table(ia, 1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_range(range); builder.set_cfs([ write_cf.build(), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data_opt(builder.build(), false); match shard .approx_split(&outer_key(100), &outer_key(500)) .unwrap() .unwrap() { ApproxSplitKey::Block(key) => { let key = key.into_encoded(); assert!(txn_types::Key::from_raw(&outer_key(100)).into_encoded() < key); assert!(key < txn_types::Key::from_raw(&outer_key(500)).into_encoded()); } ApproxSplitKey::Sst(_) => panic!("IA SST start is outside the requested range"), } } #[test] fn test_refresh_stats() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine_opt(true, DEF_BLOCK_SIZE, ""); let shard = engine.get_shard(1).unwrap(); let mut saved_vals: Vec>> = Vec::new(); let mut write_cf_builder = ShardCfBuilder::new(WRITE_CF); write_cf_builder.add_table( new_table(&engine, 11, 100, 200, 101, true, &mut saved_vals), 3, ); write_cf_builder.add_table( new_table(&engine, 12, 50, 150, 102, false, &mut saved_vals), 2, ); write_cf_builder.add_table( new_table(&engine, 13, 70, 90, 103, true, &mut saved_vals), 1, ); let mut lock_cf_builder = ShardCfBuilder::new(LOCK_CF); lock_cf_builder.add_table( new_table(&engine, 21, 50, 150, 1001, false, &mut saved_vals), 2, ); lock_cf_builder.add_table( new_table(&engine, 22, 70, 90, 1002, true, &mut saved_vals), 1, ); let mut extra_cf_builder = ShardCfBuilder::new(EXTRA_CF); extra_cf_builder.add_table( new_table(&engine, 31, 50, 150, 150, false, &mut saved_vals), 1, ); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([ write_cf_builder.build(), lock_cf_builder.build(), extra_cf_builder.build(), ]); shard.set_data(builder.build()); shard.refresh_states(); // size per entry: key 9, value 18 assert_eq!(shard.get_estimated_entries(), 440); // 100 + 100 + 20 + 100 + 20 + 100 assert_eq!(load_u64(&shard.sst_max_ts), 150); // max(WRITE_CF, EXTRA_CF). TODO: test with mem tables. assert_eq!(shard.get_max_ts(), 150); // max(WRITE_CF, EXTRA_CF) assert_eq!(shard.get_estimated_kv_size(), 3780); // 100*27 + 120*9 assert_eq!(load_u64(&shard.lv2plus_max_ts), 102); // max(WRITE_CF & Level 2+) assert_eq!(load_u64(&shard.lv2plus_tombs), 100); // WRITE_CF & Level 2+ only assert_eq!(load_u64(&shard.lv2plus_entries_write_cf), 200); // WRITE_CF & Level 2+ only } #[test] fn test_refresh_stats_with_memtables() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let mut saved_vals: Vec>> = Vec::new(); let mut write_cf_builder = ShardCfBuilder::new(WRITE_CF); write_cf_builder.add_table( new_table(&engine, 41, 100, 200, 150, true, &mut saved_vals), 1, ); let extra_cf_builder = ShardCfBuilder::new(EXTRA_CF).build(); let mem_tbl = CfTable::new(); put_write_cf_memtable_record(&mem_tbl, b"write-max", 50, 180); put_extra_cf_memtable_record(&mem_tbl, b"extra-max", 60, 220); put_lock_cf_memtable_record(&mem_tbl, b"lock-ignored", 999); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(vec![mem_tbl]); builder.set_cfs([ write_cf_builder.build(), ShardCf::new(LOCK_CF), extra_cf_builder, ]); shard.set_data(builder.build()); shard.refresh_states(); assert_eq!(load_u64(&shard.sst_max_ts), 150); assert_eq!(load_u64(&shard.max_ts), 220); assert_eq!(shard.get_max_ts(), 220); } #[test] fn test_l0table_ignore_lock() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let file = new_l0table_file( &engine, 1, [0, 0, 0], [0, 100, 0], 100, [false, false, false], ); let l0table = L0Table::new( file.clone(), BlockCache::None, false, EncryptionState::Disabled, ) .unwrap(); assert!(l0table.is_some()); // l0table would be none when `write_cf_only` is true and WRITE_CF is empty. // See https://github.com/tidbcloud/cloud-storage-engine/issues/1026. let l0table_write_cf_only = L0Table::new(file, BlockCache::None, true, EncryptionState::Disabled).unwrap(); assert!(l0table_write_cf_only.is_none()); } #[test] fn test_build_pruned_extra_cf_iters_never_prunes_writable_memtable() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let raw_key = b"extra-prune-key"; // The writable memtable (index 0) has data_max_ts below the threshold, // but it must NOT be pruned because data_max_ts may lag behind actual // entry versions during concurrent writes. let writable = CfTable::new(); put_extra_cf_memtable_record(&writable, raw_key, 20, 20); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(vec![writable]); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data(builder.build()); let snap = shard.new_snap_access(); let mut stats = read::TsFilterStats::default(); let mut iter = snap.build_pruned_extra_cf_iter(raw_key, 30, true, &mut stats); assert_eq!(stats.candidate_memtables, 1); assert_eq!(stats.pruned_memtables_by_max_ts, 0); assert_eq!(stats.surviving_memtables, 1); iter.seek(InnerKey::from_inner_buf(raw_key)); assert!(iter.valid()); } #[test] fn test_build_pruned_extra_cf_iters_uses_per_cf_memtable_max_ts() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let raw_key = b"extra-prune-key"; // Use a frozen memtable (index 1) so the pruning check is exercised. // WRITE_CF data_max_ts (500) is above threshold, but EXTRA_CF (100) is // below; the pruning helper must use per-CF max_ts, not the cross-CF // aggregate. let writable = CfTable::new(); let frozen = CfTable::new(); put_extra_cf_memtable_record(&frozen, raw_key, 100, 100); put_write_cf_memtable_record(&frozen, b"write-only-key", 200, 500); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(vec![writable, frozen]); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), ShardCf::new(EXTRA_CF), ]); shard.set_data(builder.build()); let snap = shard.new_snap_access(); let mut stats = read::TsFilterStats::default(); let mut iter = snap.build_pruned_extra_cf_iter(raw_key, 200, true, &mut stats); assert_eq!(stats.candidate_memtables, 1); assert_eq!(stats.pruned_memtables_by_max_ts, 1); assert_eq!(stats.surviving_memtables, 0); iter.seek(InnerKey::from_inner_buf(raw_key)); assert!(!iter.valid()); } #[test] fn test_build_pruned_extra_cf_iters_prunes_old_sources() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let raw_key = b"extra-prune-key"; let writable = CfTable::new(); put_extra_cf_memtable_record(&writable, raw_key, 400, 400); let immutable = CfTable::new(); put_extra_cf_memtable_record(&immutable, raw_key, 200, 200); let l0_tbl = new_extra_cf_l0_table(&engine, 7001, raw_key, &[(150, 150)]); let stable_overlap = new_extra_cf_table(&engine, 7002, raw_key, &[(100, 100)]); let stable_other = new_extra_cf_table(&engine, 7003, b"other-extra-key", &[(900, 900)]); let mut extra_cf_builder = ShardCfBuilder::new(EXTRA_CF); extra_cf_builder.add_table(stable_overlap, 1); extra_cf_builder.add_table(stable_other, 1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(vec![writable, immutable]); builder.set_l0_tbls(vec![l0_tbl]); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), extra_cf_builder.build(), ]); shard.set_data(builder.build()); let snap = shard.new_snap_access(); let mut stats = read::TsFilterStats::default(); let mut iter = snap.build_pruned_extra_cf_iter(raw_key, 250, true, &mut stats); assert_eq!(stats.candidate_memtables, 2); assert_eq!(stats.pruned_memtables_by_max_ts, 1); assert_eq!(stats.surviving_memtables, 1); assert_eq!(stats.candidate_l0s, 1); assert_eq!(stats.pruned_l0s_by_max_ts, 1); assert_eq!(stats.surviving_l0s, 0); assert_eq!(stats.candidate_stables, 1); assert_eq!(stats.pruned_stables_by_max_ts, 1); assert_eq!(stats.surviving_stables, 0); iter.seek(InnerKey::from_inner_buf(raw_key)); assert!(iter.valid()); let um = parse_test_user_meta(iter.value().user_meta(), iter.key().deref()); assert_eq!(um.start_ts, 400); assert_eq!(um.commit_ts, 400); iter.next(); assert!(!iter.valid() || !iter.key().starts_with(raw_key)); } #[test] fn test_build_pruned_extra_cf_iters_prunes_sources_when_max_ts_equals_threshold() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let raw_key = b"extra-prune-key"; let writable = CfTable::new(); let frozen = CfTable::new(); put_extra_cf_memtable_record(&frozen, raw_key, 30, 30); let l0_tbl = new_extra_cf_l0_table(&engine, 7101, raw_key, &[(30, 30)]); let stable_tbl = new_extra_cf_table(&engine, 7102, raw_key, &[(30, 30)]); let stable_other = new_extra_cf_table(&engine, 7103, b"other-extra-key", &[(30, 30)]); let mut extra_cf_builder = ShardCfBuilder::new(EXTRA_CF); extra_cf_builder.add_table(stable_tbl, 1); extra_cf_builder.add_table(stable_other, 1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(vec![writable, frozen]); builder.set_l0_tbls(vec![l0_tbl]); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), extra_cf_builder.build(), ]); shard.set_data(builder.build()); let snap = shard.new_snap_access(); let (stats, metas) = collect_pruned_extra_cf_user_metas(&snap, raw_key, 30); assert_eq!(stats.candidate_memtables, 1); assert_eq!(stats.pruned_memtables_by_max_ts, 1); assert_eq!(stats.surviving_memtables, 0); assert_eq!(stats.candidate_l0s, 1); assert_eq!(stats.pruned_l0s_by_max_ts, 1); assert_eq!(stats.surviving_l0s, 0); assert_eq!(stats.candidate_stables, 1); assert_eq!(stats.pruned_stables_by_max_ts, 1); assert_eq!(stats.surviving_stables, 0); assert!(metas.is_empty()); } #[test] fn test_build_pruned_extra_cf_iters_skips_non_overlapping_l0_from_candidates() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let raw_key = b"extra-prune-key"; let writable = CfTable::new(); put_extra_cf_memtable_record(&writable, raw_key, 251, 251); let l0_tbl = new_extra_cf_l0_table(&engine, 7351, b"other-extra-key", &[(999, 999)]); let stable_tbl = new_extra_cf_table(&engine, 7352, raw_key, &[(249, 249)]); let mut extra_cf_builder = ShardCfBuilder::new(EXTRA_CF); extra_cf_builder.add_table(stable_tbl, 1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(vec![writable]); builder.set_l0_tbls(vec![l0_tbl]); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), extra_cf_builder.build(), ]); shard.set_data(builder.build()); let snap = shard.new_snap_access(); let (stats, metas) = collect_pruned_extra_cf_user_metas(&snap, raw_key, 250); assert_eq!(stats.candidate_memtables, 1); assert_eq!(stats.pruned_memtables_by_max_ts, 0); assert_eq!(stats.surviving_memtables, 1); assert_eq!(stats.candidate_l0s, 0); assert_eq!(stats.pruned_l0s_by_max_ts, 0); assert_eq!(stats.surviving_l0s, 0); assert_eq!(stats.candidate_stables, 1); assert_eq!(stats.pruned_stables_by_max_ts, 1); assert_eq!(stats.surviving_stables, 0); assert_eq!(metas.len(), 1); assert_eq!(metas[0].start_ts, 251); assert_eq!(metas[0].commit_ts, 251); } #[test] fn test_build_pruned_extra_cf_iters_keeps_rollback_only_source_with_inflated_max_ts() { ::test_util::init_log_for_test(); let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let raw_key = b"extra-prune-key"; let stable_tbl = new_extra_cf_table(&engine, 7401, raw_key, &[(500, 0)]); let mut extra_cf_builder = ShardCfBuilder::new(EXTRA_CF); extra_cf_builder.add_table(stable_tbl, 1); let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), extra_cf_builder.build(), ]); shard.set_data(builder.build()); let snap = shard.new_snap_access(); let (stats, metas) = collect_pruned_extra_cf_user_metas(&snap, raw_key, 250); assert_eq!(stats.candidate_stables, 1); assert_eq!(stats.pruned_stables_by_max_ts, 0); assert_eq!(stats.surviving_stables, 1); assert_eq!(metas.len(), 1); assert_eq!(metas[0].start_ts, 500); assert_eq!(metas[0].commit_ts, 0); assert!(metas[0].is_rollback()); } #[test] fn test_build_pruned_extra_cf_iters_matches_full_scan_randomized() { ::test_util::init_log_for_test(); let duration_override = extra_cf_fuzz_duration_from_env(EXTRA_CF_KVENGINE_FUZZ_DURATION_ENV); let root_seed = rand::thread_rng().next_u64(); println!( "extra cf kvengine differential fuzz start: root_seed={root_seed:#x} duration_override={duration_override:?} default_rounds={} cases_per_round={}", EXTRA_CF_KVENGINE_DEFAULT_ROUNDS, EXTRA_CF_KVENGINE_CASES_PER_ROUND, ); let start = Instant::now(); let mut next_report = EXTRA_CF_KVENGINE_REPORT_INTERVAL; let mut completed_rounds = 0; let mut round_seed_rng = StdRng::seed_from_u64(root_seed); while should_continue_extra_cf_fuzz( duration_override, start, completed_rounds, EXTRA_CF_KVENGINE_DEFAULT_ROUNDS, ) { let round_seed = round_seed_rng.next_u64(); run_randomized_pruned_extra_cf_diff(round_seed, EXTRA_CF_KVENGINE_CASES_PER_ROUND); completed_rounds += 1; let elapsed = start.saturating_elapsed(); while duration_override.is_some() && elapsed >= next_report { println!( "extra cf kvengine differential fuzz progress: root_seed={root_seed:#x} elapsed={elapsed:?} rounds={completed_rounds} cases={}", completed_rounds * EXTRA_CF_KVENGINE_CASES_PER_ROUND, ); next_report += EXTRA_CF_KVENGINE_REPORT_INTERVAL; } } println!( "extra cf kvengine differential fuzz done: root_seed={root_seed:#x} elapsed={:?} rounds={completed_rounds} cases={}", start.saturating_elapsed(), completed_rounds * EXTRA_CF_KVENGINE_CASES_PER_ROUND, ); } fn collect_pruned_extra_cf_user_metas( snap: &SnapAccess, raw_key: &[u8], threshold: u64, ) -> (read::TsFilterStats, Vec) { let mut stats = read::TsFilterStats::default(); let mut iter = snap.build_pruned_extra_cf_iter(raw_key, threshold, true, &mut stats); iter.seek(InnerKey::from_inner_buf(raw_key)); if !iter.valid() { return (stats, Vec::new()); } let mut metas = Vec::new(); while iter.valid() && iter.key().starts_with(raw_key) { if !iter.value().is_deleted() { metas.push(parse_test_user_meta( iter.value().user_meta(), iter.key().deref(), )); } iter.next(); } (stats, metas) } fn run_randomized_pruned_extra_cf_diff(seed: u64, cases: usize) { let (engine, _) = new_test_engine(); let shard = engine.get_shard(1).unwrap(); let target_key = b"extra-prune-key".to_vec(); let other_keys = [ b"alpha-other-key".to_vec(), b"middle-other-key".to_vec(), b"omega-other-key".to_vec(), ]; let all_keys = { let mut keys = vec![target_key.clone()]; keys.extend(other_keys.iter().cloned()); keys }; let mut rng = StdRng::seed_from_u64(seed); let mut next_file_id = 80_000_u64; for case_idx in 0..cases { let mut used_start_ts = HashSet::new(); let mut all_entries = Vec::new(); let mut mem_tbls = Vec::new(); for _ in 0..rng.gen_range(0..=3) { let entries = gen_random_extra_cf_entries(&mut rng, &all_keys, &mut used_start_ts); if entries.is_empty() { continue; } let mem_tbl = CfTable::new(); for (raw_key, start_ts, commit_ts) in &entries { put_extra_cf_memtable_record(&mem_tbl, raw_key, *start_ts, *commit_ts); } all_entries.extend(entries); mem_tbls.push(mem_tbl); } if mem_tbls.is_empty() { mem_tbls.push(CfTable::new()); } let mut l0_tbls = Vec::new(); for _ in 0..rng.gen_range(0..=3) { let entries = gen_random_extra_cf_entries(&mut rng, &all_keys, &mut used_start_ts); if entries.is_empty() { continue; } let l0_tbl = new_extra_cf_l0_table_entries(&engine, next_file_id, &entries); next_file_id += 1; all_entries.extend(entries); l0_tbls.push(l0_tbl); } let stable_entries = if rng.gen_bool(0.75) { gen_random_extra_cf_entries(&mut rng, &all_keys, &mut used_start_ts) } else { Vec::new() }; if !stable_entries.is_empty() { all_entries.extend(stable_entries.clone()); } let mut extra_cf_builder = ShardCfBuilder::new(EXTRA_CF); if !stable_entries.is_empty() { let stable_tbl = new_extra_cf_table_entries(&engine, next_file_id, &stable_entries); next_file_id += 1; extra_cf_builder.add_table(stable_tbl, 1); } let mut builder = ShardDataBuilder::new(shard.get_data()); builder.set_mem_tbls(mem_tbls); builder.set_l0_tbls(l0_tbls); builder.set_cfs([ ShardCf::new(WRITE_CF), ShardCf::new(LOCK_CF), extra_cf_builder.build(), ]); shard.set_data(builder.build()); let threshold = choose_random_threshold(&mut rng, &all_entries, &target_key); let snap = shard.new_snap_access(); let full = collect_full_extra_cf_conflicts(&snap, &target_key, threshold); let pruned = collect_pruned_extra_cf_conflicts(&snap, &target_key, threshold); assert_eq!( pruned, full, "seed={seed} case={case_idx} threshold={threshold} entries={:?}", debug_extra_cf_entries(&all_entries) ); } } fn should_continue_extra_cf_fuzz( duration_override: Option, start: Instant, completed_rounds: usize, default_rounds: usize, ) -> bool { match duration_override { Some(duration) => completed_rounds == 0 || start.saturating_elapsed() < duration, None => completed_rounds < default_rounds, } } fn extra_cf_fuzz_duration_from_env(env_var: &str) -> Option { match std::env::var(env_var) { Ok(value) => Some(Duration::from( value .parse::() .unwrap_or_else(|err| panic!("invalid {}={value:?}: {}", env_var, err)), )), Err(std::env::VarError::NotPresent) => None, Err(err) => panic!("failed to read {}: {}", env_var, err), } } fn gen_random_extra_cf_entries( rng: &mut StdRng, keys: &[Vec], used_start_ts: &mut HashSet, ) -> Vec<(Vec, u64, u64)> { let mut entries = Vec::new(); for key in keys { let record_count = rng.gen_range(0..=3); for _ in 0..record_count { let start_ts = loop { let candidate = rng.gen_range(1..=2_000); if used_start_ts.insert(candidate) { break candidate; } }; let commit_ts = if rng.gen_bool(0.35) { 0 } else { start_ts + rng.gen_range(1..=200) }; entries.push((key.clone(), start_ts, commit_ts)); } } entries } fn choose_random_threshold( rng: &mut StdRng, entries: &[(Vec, u64, u64)], target_key: &[u8], ) -> u64 { let mut versions: Vec = entries .iter() .filter(|(raw_key, ..)| raw_key.as_slice() == target_key) .map(|(_, start_ts, commit_ts)| { if *commit_ts == 0 { *start_ts } else { *commit_ts } }) .collect(); versions.sort_unstable(); versions.dedup(); let mut candidates = vec![0, 1, 50, 250, 500, 1_000, 2_000]; if let Some(min) = versions.first().copied() { candidates.push(min); candidates.push(min.saturating_sub(1)); candidates.push(min.saturating_add(1)); } if let Some(max) = versions.last().copied() { candidates.push(max); candidates.push(max.saturating_add(1)); } for version in versions.choose_multiple(rng, versions.len().min(4)) { candidates.push(*version); } candidates.sort_unstable(); candidates.dedup(); candidates[rng.gen_range(0..candidates.len())] } fn collect_full_extra_cf_conflicts( snap: &SnapAccess, raw_key: &[u8], threshold: u64, ) -> Vec<(u64, u64)> { let mut iters: Vec> = Vec::new(); snap.iter_tables(EXTRA_CF, |table_ref| match table_ref { read::TableRef::Memtable(mem_tbl) => { iters.push(mem_tbl.get_cf(EXTRA_CF).new_iterator(false)) } read::TableRef::L0(l0_tbl) => { if let Some(tbl) = l0_tbl.get_cf(EXTRA_CF) { iters.push(tbl.new_iterator(false, true)); } } read::TableRef::SsTable(tbl) => iters.push(tbl.new_iterator(false, true)), }); if iters.is_empty() { return Vec::new(); } let mut iter = table::new_merge_iterator(iters, false); iter.seek(InnerKey::from_inner_buf(raw_key)); let mut conflicts = Vec::new(); while iter.valid() && iter.key().starts_with(raw_key) { if !iter.value().is_deleted() { let user_meta = parse_test_user_meta(iter.value().user_meta(), iter.key().deref()); if !user_meta.is_rollback() && user_meta.commit_ts > threshold { conflicts.push((user_meta.start_ts, user_meta.commit_ts)); } } iter.next(); } conflicts } fn collect_pruned_extra_cf_conflicts( snap: &SnapAccess, raw_key: &[u8], threshold: u64, ) -> Vec<(u64, u64)> { let mut stats = read::TsFilterStats::default(); let mut iter = snap.build_pruned_extra_cf_iter(raw_key, threshold, true, &mut stats); iter.seek(InnerKey::from_inner_buf(raw_key)); if !iter.valid() { return Vec::new(); } let mut conflicts = Vec::new(); while iter.valid() && iter.key().starts_with(raw_key) { if !iter.value().is_deleted() { let user_meta = parse_test_user_meta(iter.value().user_meta(), iter.key().deref()); if !user_meta.is_rollback() && user_meta.commit_ts > threshold { conflicts.push((user_meta.start_ts, user_meta.commit_ts)); } } iter.next(); } conflicts } fn debug_extra_cf_entries(entries: &[(Vec, u64, u64)]) -> Vec<(String, u64, u64)> { let mut debug_entries: Vec<(String, u64, u64)> = entries .iter() .map(|(raw_key, start_ts, commit_ts)| { ( String::from_utf8_lossy(raw_key).into_owned(), *start_ts, *commit_ts, ) }) .collect(); debug_entries.sort(); debug_entries } fn parse_test_user_meta(user_meta: &[u8], key: impl AsRef<[u8]>) -> UserMeta { let key = key.as_ref(); assert!( !user_meta.is_empty(), "empty user_meta for key {:?}", String::from_utf8_lossy(key) ); assert_eq!( user_meta[0], USER_META_FORMAT_V1, "unexpected user_meta format for key {:?}: {:?}", String::from_utf8_lossy(key), user_meta ); UserMeta::from_slice(user_meta) } fn print_locks(snap: &SnapAccess, all_versions: bool, read_ts: Option) { let mut it = snap.new_iterator(LOCK_CF, false, all_versions, read_ts, false); it.rewind(); while it.valid() { debug!( "key: {:?}, version: {:?}, meta: {:?}", tikv_util::escape(it.key()), it.version(), it.meta() ); it.next(); } } fn keyspace_prefix(keyspace_id: u32) -> Vec { let mut prefix = keyspace_id.to_be_bytes().to_vec(); prefix[0] = b'x'; prefix } #[inline] fn encode_i64_to_comparable_u64(v: i64) -> u64 { (v as u64) ^ 1 << 63_u64 } fn prepare_table_region( engine: &TestEngine, apply_tx: &mpsc::Sender, keyspace_id: u32, table_id: i64, ) -> u64 { // Split keyspace let keyspace_1_prefix = keyspace_prefix(keyspace_id); let mut table_1_prefix = keyspace_1_prefix.clone(); table_1_prefix.extend_from_slice(b"t"); table_1_prefix.extend_from_slice( encode_i64_to_comparable_u64(table_id) .to_be_bytes() .as_slice(), ); let mut table_2_prefix = keyspace_1_prefix; table_2_prefix.extend_from_slice(b"t"); table_2_prefix.extend_from_slice( encode_i64_to_comparable_u64(table_id + 1) .to_be_bytes() .as_slice(), ); let split_keys = vec![ keyspace_prefix(keyspace_id), table_1_prefix, table_2_prefix, keyspace_prefix(keyspace_id + 1), ]; let mut splitter = Splitter::new(split_keys, IdVer::new(1, 1), 1, apply_tx.clone()); let handle = thread::spawn(move || { splitter.run(); }); let ok = try_wait(|| engine.shards.len() == 5, 2); assert!(ok, "split failed"); handle.join().unwrap(); 4 } fn i_to_key(i: i32, min_blob_size: u32) -> String { if min_blob_size > 0 { // 3 -> strlen("key") format!("key{:0>1$}", i, min_blob_size as usize - 3) } else { format!("key{:0>1$}", i, 6) } } #[maybe_async::both] async fn new_table_ext( engine: &TestEngine, id: u64, begin: usize, end: usize, version: F, del: bool, saved_vals: &mut Vec>>, ) -> SsTable where F: Fn(usize) -> u64, { let block_size = engine.opts.table_builder_options.block_size; let comp_tp = engine.opts.table_builder_options.compression_tps[0]; let comp_lvl = engine.opts.table_builder_options.compression_lvl; let mut builder = table::sstable::builder::Builder::new( id, block_size, comp_tp, comp_lvl, ChecksumType::Crc32, None, ); for i in begin..end { let version = version(i); let key = engine.key_builder.i_to_inner_key(i); let val = if del { table::Value::new_with_meta_version(BIT_DELETE, version, 0, &[]) } else { let val = Rc::new(key.as_ref().repeat(2)); // Save the val with Rc, make sure the val_rc stay valid during the iterator // lifecycle. saved_vals.push(val.clone()); table::Value::new_with_meta_version(0, version, 0, &val) }; builder.add(key.as_ref(), &val, None).unwrap(); } let mut data_buf = Vec::new(); builder.finish(0, &mut data_buf).unwrap(); let data = Bytes::from(data_buf); let opts = dfs::Options::default(); dfs_create_table(engine.fs.as_ref(), id, data.clone(), opts) .await .unwrap(); let file = InMemFile::new(id, data).await; SsTable::new(Arc::new(file), BlockCache::None, EncryptionState::Disabled).unwrap() } #[maybe_async::both] async fn new_table( engine: &TestEngine, id: u64, begin: usize, end: usize, version: u64, del: bool, saved_vals: &mut Vec>>, ) -> SsTable { new_table_ext(engine, id, begin, end, move |_| version, del, saved_vals).await } fn new_extra_cf_table( engine: &TestEngine, id: u64, raw_key: &[u8], records: &[(u64, u64)], ) -> SsTable { let entries: Vec<(Vec, u64, u64)> = records .iter() .map(|(start_ts, commit_ts)| (raw_key.to_vec(), *start_ts, *commit_ts)) .collect(); new_extra_cf_table_entries(engine, id, &entries) } fn new_extra_cf_table_entries( engine: &TestEngine, id: u64, entries: &[(Vec, u64, u64)], ) -> SsTable { let block_size = engine.opts.table_builder_options.block_size; let comp_tp = engine.opts.table_builder_options.compression_tps[0]; let comp_lvl = engine.opts.table_builder_options.compression_lvl; let mut builder = table::sstable::builder::Builder::new( id, block_size, comp_tp, comp_lvl, ChecksumType::Crc32, None, ); let mut rows: Vec<(Vec, u64, u8, Vec)> = entries .iter() .map(|(raw_key, start_ts, commit_ts)| { let key = encode_extra_txn_status_key(raw_key, *start_ts); let version = if *commit_ts == 0 { *start_ts } else { *commit_ts }; let user_meta = UserMeta::new(*start_ts, *commit_ts).to_array(); let mut buf = Vec::with_capacity(user_meta.len() + 1); buf.extend_from_slice(&user_meta); buf.push(0); (key.to_vec(), version, user_meta.len() as u8, buf) }) .collect(); rows.sort_by(|(left, ..), (right, ..)| left.cmp(right)); for (key, version, user_meta_len, buf) in rows { let val = table::Value::new_with_meta_version(0, version, user_meta_len, &buf); builder .add(InnerKey::from_inner_buf(key.as_slice()), &val, None) .unwrap(); } let mut data_buf = Vec::new(); builder.finish(0, &mut data_buf).unwrap(); let data = Bytes::from(data_buf); let opts = dfs::Options::default(); dfs_create_table(engine.fs.as_ref(), id, data.clone(), opts).unwrap(); let file = InMemFile::new(id, data); SsTable::new(Arc::new(file), BlockCache::None, EncryptionState::Disabled).unwrap() } fn new_extra_cf_l0_table( engine: &TestEngine, id: u64, raw_key: &[u8], records: &[(u64, u64)], ) -> L0Table { let entries: Vec<(Vec, u64, u64)> = records .iter() .map(|(start_ts, commit_ts)| (raw_key.to_vec(), *start_ts, *commit_ts)) .collect(); new_extra_cf_l0_table_entries(engine, id, &entries) } fn new_extra_cf_l0_table_entries( engine: &TestEngine, id: u64, entries: &[(Vec, u64, u64)], ) -> L0Table { let block_size = engine.opts.table_builder_options.block_size; let fs = engine.fs.clone(); let mut builder = L0Builder::new(id, block_size, 1.into(), ChecksumType::Crc32, None); let mut rows: Vec<(Vec, u64, u8, Vec)> = entries .iter() .map(|(raw_key, start_ts, commit_ts)| { let key = encode_extra_txn_status_key(raw_key, *start_ts); let version = if *commit_ts == 0 { *start_ts } else { *commit_ts }; let user_meta = UserMeta::new(*start_ts, *commit_ts).to_array(); let mut buf = Vec::with_capacity(user_meta.len() + 1); buf.extend_from_slice(&user_meta); buf.push(0); (key.to_vec(), version, user_meta.len() as u8, buf) }) .collect(); rows.sort_by(|(left, ..), (right, ..)| left.cmp(right)); for (key, version, user_meta_len, buf) in rows { let val = table::Value::new_with_meta_version(0, version, user_meta_len, &buf); builder .add( EXTRA_CF, InnerKey::from_inner_buf(key.as_slice()), &val, None, ) .unwrap(); } let (_, data) = builder.finish().unwrap(); let opts = dfs::Options::default(); let runtime = fs.get_runtime(); runtime.block_on(fs.create(id, data.clone(), opts)).unwrap(); let file = Arc::new(InMemFile::new(id, data)); L0Table::new(file, BlockCache::None, false, EncryptionState::Disabled) .unwrap() .unwrap() } fn put_write_cf_memtable_record(mem_tbl: &CfTable, raw_key: &[u8], start_ts: u64, commit_ts: u64) { let user_meta = UserMeta::new(start_ts, commit_ts).to_array(); let mut wb = crate::table::memtable::WriteBatch::new(); wb.put( InnerKey::from_outer_key(raw_key), 0, &user_meta, commit_ts, b"value", ); mem_tbl.get_cf(WRITE_CF).put_batch(&mut wb, None, WRITE_CF); } fn put_lock_cf_memtable_record(mem_tbl: &CfTable, raw_key: &[u8], version: u64) { let mut wb = crate::table::memtable::WriteBatch::new(); wb.put(InnerKey::from_outer_key(raw_key), 0, &[], version, b"lock"); mem_tbl.get_cf(LOCK_CF).put_batch(&mut wb, None, LOCK_CF); } fn put_extra_cf_memtable_record(mem_tbl: &CfTable, raw_key: &[u8], start_ts: u64, commit_ts: u64) { let key = encode_extra_txn_status_key(raw_key, start_ts); let version = if commit_ts == 0 { start_ts } else { commit_ts }; let user_meta = UserMeta::new(start_ts, commit_ts).to_array(); let mut wb = crate::table::memtable::WriteBatch::new(); wb.put(InnerKey::from_inner_buf(&key), 0, &user_meta, version, &[0]); mem_tbl.get_cf(EXTRA_CF).put_batch(&mut wb, None, EXTRA_CF); } fn dfs_create_table( fs: &dyn dfs::Dfs, file_id: u64, data: Bytes, opts: dfs::Options, ) -> dfs::Result<()> { let runtime = fs.get_runtime(); runtime.block_on(fs.create(file_id, data, opts)) } async fn dfs_create_table_async( dfs: &dyn dfs::Dfs, file_id: u64, data: Bytes, opts: dfs::Options, ) -> dfs::Result<()> { dfs.create(file_id, data, opts).await } fn new_l0table_file( engine: &TestEngine, id: u64, begin: [usize; NUM_CFS], end: [usize; NUM_CFS], version: u64, del: [bool; NUM_CFS], ) -> Arc { let block_size = engine.opts.table_builder_options.block_size; let fs = engine.fs.clone(); let mut builder = L0Builder::new(id, block_size, version.into(), ChecksumType::Crc32, None); for cf in 0..NUM_CFS { for i in begin[cf]..end[cf] { let key = engine.key_builder.i_to_inner_key(i); if del[cf] { let val = table::Value::new_with_meta_version(BIT_DELETE, version, 0, &[]); builder.add(cf, key.as_ref(), &val, None).unwrap(); } else { let val_buf = key.as_ref().repeat(2); let val = table::Value::new_with_meta_version(0, version, 0, &val_buf); builder.add(cf, key.as_ref(), &val, None).unwrap(); } } } let (_, data) = builder.finish().unwrap(); let opts = dfs::Options::default(); let runtime = fs.get_runtime(); runtime.block_on(fs.create(id, data.clone(), opts)).unwrap(); Arc::new(InMemFile::new(id, data)) } fn load_data( begin: usize, end: usize, version: u64, tx: mpsc::Sender, min_blob_size: u32, ) { let mut wb = WriteBatch::new(1); for i in begin..end { let key = i_to_key(i as i32, min_blob_size); for cf in 0..3 { let val = key.repeat(cf + 2); let version = if cf == 1 { 0 } else { version }; wb.put(cf, key.as_bytes(), val.as_bytes(), 0, &[], version); } if i % 100 == 99 { info!("load data {}:{}", i - 99, i); write_data(wb, &tx); wb = WriteBatch::new(1); thread::sleep(Duration::from_millis(10)); } } if wb.num_entries() > 0 { write_data(wb, &tx); } } fn load_data_ext( engine: &TestEngine, begin: [usize; NUM_CFS], end: [usize; NUM_CFS], version: u64, del: [bool; NUM_CFS], tx: &mpsc::Sender, ) -> u64 { let mut wb = WriteBatch::new(1); for cf in 0..NUM_CFS { for i in begin[cf]..end[cf] { let key = engine.key_builder.i_to_outer_key(i); let version = if cf == 1 { 0 } else { version }; if del[cf] { wb.put(cf, &key, &[], BIT_DELETE, &[], version); } else { let val = key.repeat(2); wb.put(cf, &key, &val, 0, &[], version); } } } if wb.num_entries() > 0 { write_data(wb, tx) } else { 0 } } fn wait_for_compact_to_level_1_plus(en: &Engine, i: usize, timeout: Duration) { let key = i_to_key(i as i32, en.opts.blob_table_build_options.min_blob_size); let ok = try_wait( || { let shard = get_shard_for_key(key.as_bytes(), en); let stats = en.get_shard_stat(shard.id); stats.compaction_score < 1.0 && stats.cfs[WRITE_CF] .levels .iter() .any(|lv| lv.num_tables > 0) }, timeout.as_secs() as usize, ); assert!(ok, "wait for compact to level 1+ timeout"); } fn switch_mem_table(tx: &mpsc::Sender) { let mut wb = WriteBatch::new(1); wb.set_switch_mem_table(); write_data(wb, tx); } fn write_data(wb: WriteBatch, applier_tx: &mpsc::Sender) -> u64 /* write_sequence */ { let (result_tx, result_rx) = mpsc::bounded(1); let task = ApplyTask::new_wb(wb, result_tx); if let Err(err) = applier_tx.send(task) { panic!("{:?}", err); } result_rx.recv().unwrap().unwrap() } #[maybe_async::both] async fn check_get( begin: usize, end: usize, version: u64, cfs: &[usize], en: &Engine, exist: bool, check_version: Option, min_blob_size: u32, ) { for i in begin..end { let key = i_to_key(i as i32, min_blob_size); let shard = get_shard_for_key(key.as_bytes(), en); let snap = SnapAccess::new(&shard); for &cf in cfs { let version = if cf == 1 { 0 } else { version }; let item = snap.get(cf, key.as_bytes(), version).await; if item.is_valid() { if !exist { if item.is_deleted() { continue; } let shard_stats = shard.get_stats(); panic!( "got key {}, shard {}:{}, cf {}, stats {:?}", key, shard.id, shard.ver, cf, shard_stats, ); } assert_eq!(item.get_value(), key.repeat(cf + 2).as_bytes()); if cf != 1 { if let Some(check_version) = check_version { assert_eq!(item.version, check_version); } } } else if exist { let shard_stats = shard.get_stats(); panic!( "failed to get key {}, shard {}, stats {:?}", key, shard.tag(), shard_stats, ); } } } } #[maybe_async::both] async fn check_iterator(begin: usize, end: usize, en: &Engine) { for cf in 0..3 { let mut i = begin; // let ids = vec![2, 3, 4, 5, 1]; let ids = vec![1]; for id in ids { let shard = en.get_shard(id).unwrap(); let snap = SnapAccess::new(&shard); let mut iter = snap.new_iterator(cf, false, false, None, true).await; iter.seek(shard.outer_start.chunk()).await; while iter.valid() { if iter.key.chunk() >= shard.outer_end.chunk() { break; } let key = i_to_key(i as i32, en.opts.blob_table_build_options.min_blob_size); assert_eq!(iter.key(), key.as_bytes()); assert_eq!(iter.val(), key.repeat(cf + 2).as_bytes()); i += 1; iter.next().await; } } assert_eq!(i, end); } } fn get_shard_for_key(key: &[u8], en: &Engine) -> Arc { for id in 1_u64..=5 { if let Some(shard) = en.get_shard(id) { if shard .data_bound() .overlap_key(InnerKey::from_inner_buf(key)) { return shard; } } } en.get_shard(1).unwrap() } #[must_use] fn try_wait(f: F, seconds: usize) -> bool where F: Fn() -> bool, { let begin = Instant::now_coarse(); let timeout = Duration::from_secs(seconds as u64); while begin.saturating_elapsed() < timeout { if f() { return true; } thread::sleep(Duration::from_millis(100)) } false } pub(crate) fn generate_encryption_key() -> EncryptionKey { let master_key_plain_text = thread_rng().gen::<[u8; 32]>().to_vec(); let master_key = MasterKey::new(&master_key_plain_text); master_key.generate_encryption_key() }