// Copyright 2023 PingCAP, Inc. use std::{sync::Arc, time::Duration}; use api_version::ApiV2; use bytes::Bytes; use cloud_encryption::KeyspaceEncryptionConfig; use kvengine::{ KvEnginePerKeyspaceConfig, WRITE_CF, dfs::{DFSConfig, Dfs, FileType, Options, S3Fs, new_dfs_from_config}, table::{ file::InMemFile, schema_file::{SchemaFile, build_schema_file, schema_from_table_info}, }, }; use kvproto::{kvrpcpb, metapb}; use native_br::{ backup::{self, BackupType}, common::now, metrics::NATIVE_BR_RFENGINE_WAL_EPOCH_OVERWRITTEN_ERROR, packing::{ CopyPackedRun, MigratePackEnv, NoopReporter, PackContext, PackEnv, UnpackRun, offline_pd::OfflinePd, }, restore::RestoreConfig, restore_keyspace, step, }; use pd_client::{PdClient, pd_control::KeyspaceMeta}; use rand::prelude::*; use rstest::{fixture, rstest}; use schema::{generate_storage_class_schema_data_for_test, schema::StorageClass}; use security::SecurityConfig; use tempfile::TempDir; use test_cloud_server::{ ServerCluster, ServerClusterBuilder, TikvWorkerOptions, client::{ ClusterClient, ClusterClientOptions, CommitAction, MutateOptions, PrewriteExt, RequestOptions, TxnMutations, TxnWriteMethod, }, keyspace::CreateKeyspaceOptions, must_wait, oss::{ObjectStorageService, prepare_dfs}, try_wait, util::{ Mutation, broadcast_schema_file_request, broadcast_schema_file_request_and_check, get_table_split_keys, request_major_compaction, }, }; use test_pd_client::PdWrapper; use tikv::config::TikvConfig; use tikv_util::{ codec::bytes::encode_bytes, config::{ReadableDuration, ReadableSize}, debug, info, store::new_learner_peer, time, warn, }; use tokio::runtime::{Handle, Runtime}; use txn_types::{LockType, WriteType}; use crate::{ alloc_node_id_vec, create_keyspace_meta, native_backup::{ DummyStepReporter, gen_keyspace_key, generate_backup_name, i_to_key, i_to_val, wait_all_dfs_workers_healthy, }, new_security_config, transaction::helper::{ get_mvcc_info, must_unlocked, must_written_with_properties, new_lock_mutation, new_shared_lock_mutation, }, wait_for_keyspace_stats, }; const BASIC_DATA_COUNT: usize = 10; const NODES_COUNT: usize = 5; const KEYSPACE_COUNT: usize = 3; const DEFAULT_LOOP_COUNT: usize = 1; const DEFAULT_TARGET_REGIONS: usize = 2; // The numbers have no special meaning, just to make them different. // Require to be larger than 72, see `i_to_val`. const BASIC_DATA_LEN: usize = 80; const LONG_DATA_LEN: usize = 8 * 1024; const IMPORT_DATA_LEN: usize = 82; const PRE_BACKUP_DATA_LEN: usize = 84; const POST_BACKUP_DATA_LEN: usize = 86; const PRE_PITR_DATA_LEN: usize = 88; const FINAL_DATA_LEN: usize = 90; const TEST_CMEK_ID: &str = "random"; const TEST_CMEK_VENDOR: &str = "test"; fn enable_test_cmek(pd_client: Arc, keyspace_id: u32) { pd_client .set_keyspace_encryption( keyspace_id, KeyspaceEncryptionConfig { enabled: true, cmek_id: Some(TEST_CMEK_ID.to_owned()), vendor: Some(TEST_CMEK_VENDOR.to_owned()), ..Default::default() }, ) .unwrap(); } fn assert_test_cmek(pd_client: &dyn PdClient, keyspace_id: u32) { let cfg = pd_client.get_keyspace_encryption(keyspace_id).unwrap(); assert!(cfg.enabled); assert_eq!(cfg.cmek_id.as_deref(), Some(TEST_CMEK_ID)); assert_eq!(cfg.vendor.as_deref(), Some(TEST_CMEK_VENDOR)); } fn put_raw_kv_pairs( client: &mut ClusterClient, pairs: impl IntoIterator, Vec)>, ) { let mutations = pairs .into_iter() .map(|(key, value)| Mutation { key: Bytes::from(key), value: Bytes::from(value), op: kvrpcpb::Op::Put, ..Default::default() }) .collect(); client.kv_mutate(mutations).unwrap(); } #[test] fn test_restore_keyspace() { let loop_count = std::env::var("LOOP") .unwrap_or_default() .parse::() .unwrap_or(DEFAULT_LOOP_COUNT); let target_regions = std::env::var("TARGET_REGIONS") .unwrap_or_default() .parse::() .unwrap_or(DEFAULT_TARGET_REGIONS); test_restore_keyspace_opt(TestRestoreKeyspaceOptions { loop_count, target_regions, }); } struct TestRestoreKeyspaceOptions { loop_count: usize, target_regions: usize, } impl Default for TestRestoreKeyspaceOptions { fn default() -> Self { TestRestoreKeyspaceOptions { loop_count: DEFAULT_LOOP_COUNT, target_regions: DEFAULT_TARGET_REGIONS, } } } fn test_restore_keyspace_opt(options: TestRestoreKeyspaceOptions) { // keyspace_id, data_count, shuffle_regions, has_learner, loop_count let cases = [ (0, 30, None, false, 1), ( 1, 30, Some(options.target_regions), false, options.loop_count, ), (1, 30, None, true, 1), // Don't shuffle regions for stability. ]; let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let mut cluster = ServerCluster::new( alloc_node_id_vec(NODES_COUNT), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); // Set small mem-table size to make data reach L1 and generate over bound // shards. conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(1); conf.coprocessor.region_split_size = ReadableSize::kb(128); // kv_opts.base_size = 8kb conf.coprocessor.region_bucket_size = ReadableSize::kb(64); conf.rfengine.target_file_size = ReadableSize::mb(1); conf.rfengine.lightweight_backup = true; conf.rfengine.wal_chunk_target_file_size = ReadableSize::kb(128); conf.memory.enable_heap_profiling = false; }, ); cluster.wait_region_replicated(&[], 3); let pd_client = cluster.get_pd_client(); let mut client = cluster.new_client(); // Split keyspaces. client.split_keyspaces(0..KEYSPACE_COUNT as u32); cluster.wait_pd_region_count(KEYSPACE_COUNT + 2); // Import basic data. for keyspace_id in 0..KEYSPACE_COUNT as u32 { client.put_kv( 0..BASIC_DATA_COUNT, gen_keyspace_key(keyspace_id), i_to_val(BASIC_DATA_LEN), ); } let runtime = Runtime::new().unwrap(); // Run cases. for (case_idx, &(keyspace_id, data_count, shuffle_regions, has_learner, loop_count)) in cases.iter().enumerate() { step!("case group: {case_idx}"); if has_learner { pd_client.disable_default_operator(); // To prevent PD removing learners due to exceed `max-replicas`. let (start, end) = ApiV2::get_keyspace_range_by_id(keyspace_id); add_learners(&mut cluster, &start, &end); check_learners(&cluster, &start, &end, NODES_COUNT - 3, 10); step!("add learner done"); } else { pd_client.enable_default_operator(); // Remove learners by peer count check. } for loop_idx in 0..loop_count { test_restore_keyspace_impl( case_idx, loop_idx, &mut cluster, &dfs_config, keyspace_id, data_count, shuffle_regions, has_learner, &runtime, ); } } cluster.stop(); oss.shutdown(); } fn test_restore_keyspace_impl( case_idx: usize, loop_idx: usize, cluster: &mut ServerCluster, dfs_config: &DFSConfig, keyspace_id: u32, data_count: usize, shuffle_regions: Option, has_learner: bool, runtime: &Runtime, ) { step!("case: {}:{}", case_idx, loop_idx); let mut rng = rand::thread_rng(); let mut client = cluster.new_client(); let i_to_key = gen_keyspace_key(keyspace_id); let (keyspace_start, keyspace_end) = ApiV2::get_keyspace_range_by_id(keyspace_id); let backup_config = backup::BackupConfig { dfs: dfs_config.clone(), backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; let restore_config = RestoreConfig::default_for_test(); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); // Import data. let commit_action = if rng.gen_ratio(2, 3) { CommitAction::AsyncCommitSecondaryKeys(Duration::ZERO) } else { CommitAction::AsyncCommitSecondaryKeys(Duration::MAX) }; client .try_put_kv( 0..data_count, &i_to_key, i_to_val(IMPORT_DATA_LEN), MutateOptions { commit_action, ..Default::default() }, ) .unwrap(); let origin_ref_store = client.dump_ref_store(); // Prepare for PiTR. let truncate_ts = if rng.gen_ratio(1, 2) { let ts = client.get_ts(); client.put_kv( 0..(data_count / 2).max(1), &i_to_key, i_to_val(PRE_PITR_DATA_LEN), ); step!("another writes for pitr done"); client.verify_data_with_ref_store(); Some(ts.into_inner()) } else { None }; // Perform snapshot backup. let snapshot_backup_name = generate_backup_name(); let backup_meta = { let backup_ts = client.get_ts().into_inner(); // Put more data before backup to verify truncate ts take affect. client.put_kv( 0..(data_count / 2).max(1), &i_to_key, i_to_val(PRE_BACKUP_DATA_LEN), ); step!("another writes done"); client.verify_data_with_ref_store(); assert!( client .verify_data_with_given_ref_store( &origin_ref_store, None, &RequestOptions::default() ) .is_err(), "case: {}:{}", case_idx, loop_idx ); step!("verify before backup ok"); let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config.clone(), BackupType::Lightweight, snapshot_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); info!("backup_cluster result: {:?}", backup_meta); client.verify_data_with_ref_store(); step!("backup done"); backup_meta }; // Shuffle regions. if let (Some(shuffle_regions), 0) = (shuffle_regions, loop_idx) { let keyspace_prefix = ApiV2::get_keyspace_prefix_by_id(keyspace_id); let region_count = client.pd_client().get_regions_number(); let mut i = 0; while i < shuffle_regions { let split_key = rng.r#gen::() % data_count; if let Err(e) = client.try_split(&i_to_key(split_key), 1) { warn!("try split error: {:?}", e); continue; } i += 1; } cluster.wait_pd_region_min_count(region_count + shuffle_regions); let split_region_count = client.pd_client().get_regions_number(); for _ in 0..(shuffle_regions / 2) { let source_key = rng.r#gen::() % data_count; if let Err(e) = client.try_merge_adjacent_region( &i_to_key(source_key), Some(&keyspace_prefix), Duration::from_secs(3), ) { warn!("try_merge_adjacent_region fail: {:?}", e); } } let merge_region_count = client.pd_client().get_regions_number(); step!( "shuffle regions done, regions {} -> {} -> {}", region_count, split_region_count, merge_region_count ); } // Put another data. client.put_kv(0..data_count, &i_to_key, i_to_val(POST_BACKUP_DATA_LEN)); step!("another writes done"); client.verify_data_with_ref_store(); assert!( client .verify_data_with_given_ref_store(&origin_ref_store, None, &RequestOptions::default()) .is_err(), "case: {}:{}", case_idx, loop_idx, ); step!("verify ok"); // Test for "PiTR on restored data". For the following scenario: // T0: put data0. // T1: perform backup. // T2: put data1. // T3: restore to backup. Now data0 should be in store. // T4: PiTR to T2. Now data1 should be in store. let truncate_ts_pitr = client.get_ts().into_inner(); let ref_store_pitr = client.dump_ref_store(); // Instant backup before restore let instant_backup_name = { let backup_ts = client.get_ts().into_inner(); let instant_backup_name = generate_backup_name(); let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, instant_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, Some(backup_meta), ) .expect("backup::backup_cluster"); info!("backup_cluster result (instant): {:?}", backup_meta); step!("backup for pitr done"); instant_backup_name }; let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); // Restore keyspace. restore_keyspace::restore_keyspace( keyspace_id, keyspace_id, &snapshot_backup_name, None, s3fs.clone(), restore_config.clone(), cluster.get_pd_client(), None, runtime, truncate_ts, None, reporter.clone(), object_cache.clone(), limiter.clone(), ) .unwrap(); step!("restore done"); // Verify restored data. { let mut verify_options: Vec<(Option<(&[u8], &[u8])>, RequestOptions)> = vec![(None, RequestOptions::default())]; if has_learner { verify_options.push(( Some((&keyspace_start, &keyspace_end)), RequestOptions::learner(), )); } let mut verify_res = vec![]; for (range, opt) in verify_options { verify_res.push( client .verify_data_with_given_ref_store(&origin_ref_store, range, &opt) .unwrap(), ); } step!( "verify restore data done{}, kv count: {:?}", if has_learner { " (with learner)" } else { "" }, verify_res, ); } // PiTR on restored data. { restore_keyspace::restore_keyspace( keyspace_id, keyspace_id, &instant_backup_name, None, s3fs, restore_config, cluster.get_pd_client(), None, runtime, Some(truncate_ts_pitr), None, reporter, object_cache, limiter, ) .unwrap(); step!("restore (pitr on restored data) done"); let verify_res = client .verify_data_with_given_ref_store( &ref_store_pitr, Some((&keyspace_start, &keyspace_end)), &RequestOptions::default(), ) .unwrap(); step!( "verify (pitr on restored data) done, kv count: {}", verify_res.0 ); } // Write more data to verify the sequences. for id in 0..KEYSPACE_COUNT as u32 { let count = if id == keyspace_id { data_count } else { BASIC_DATA_COUNT }; client.put_kv(0..count, gen_keyspace_key(id), i_to_val(FINAL_DATA_LEN)); } client.verify_data_with_ref_store(); step!("verify more writes done"); } /// This test is to verify that `restore_keyspace` can resolve locks. /// /// Resolving locks is necessary when a backup is performed after the primary /// key is committed but the secondary keys are not in a transaction. /// /// See https://github.com/tidbcloud/cloud-storage-engine/issues/1134. #[rstest] #[case::sync(false)] #[case::async_commit(true)] fn test_restore_keyspace_with_resolve_locks(#[case] async_commit: bool) { let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let _enter = runtime.enter(); let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None); let mut cluster = ServerClusterBuilder::new( alloc_node_id_vec(NODES_COUNT), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); // Set small mem-table size to make data reach L0+. conf.rocksdb.writecf.write_buffer_size = ReadableSize(128); conf.coprocessor.region_split_size = ReadableSize::kb(128); // kv_opts.base_size = 8kb conf.coprocessor.region_bucket_size = ReadableSize::kb(64); conf.rfengine.target_file_size = ReadableSize::mb(1); conf.rfengine.lightweight_backup = true; conf.rfengine.wal_chunk_target_file_size = ReadableSize::kb(128); }, ) .pd(pd_wrapper) .build(); let keyspace_id = runtime.block_on( cluster.create_keyspace(&CreateKeyspaceOptions::default(), Duration::from_secs(10)), ); cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default()); cluster.wait_region_replicated(&[], 3); let pd_client = cluster.get_pd_client(); let mut client = cluster.new_client_opt(ClusterClientOptions { txn_file_max_chunk_size: Some(1024), ..Default::default() }); client.split_keyspace(keyspace_id); // Import data let i_to_key = gen_keyspace_key(keyspace_id); client.put_kv(0..200, &i_to_key, i_to_val(BASIC_DATA_LEN)); // Leave one shared-lock transaction in prewrite state. Native BR restore // should drop the shared-lock wrapper instead of failing on it. let shared_prewrite_key = i_to_key(351); { let shared_primary = i_to_key(350); let shared_key = shared_prewrite_key.clone(); let start_ts = client.get_ts(); let for_update_ts = client.get_ts(); client .kv_pessimistic_lock( shared_primary.clone().into(), vec![shared_primary.clone().into()], start_ts, 3000, for_update_ts, None, ) .unwrap(); client .kv_shared_pessimistic_lock( shared_primary.clone().into(), vec![shared_key.clone().into()], start_ts, 3000, for_update_ts, None, ) .unwrap(); let mut txn = client.begin_transaction(Some(start_ts)); client .kv_prewrite_ext( shared_primary.clone().into(), None, TxnMutations::from_normal(vec![ new_lock_mutation(shared_primary), new_shared_lock_mutation(shared_key), ]), &mut txn, PrewriteExt { for_update_ts, ..Default::default() }, ) .unwrap(); } if async_commit { client.set_async_commit(); } // Delete data without committing secondary keys, to simulate the case that if // we don't resolve locks during restoration, the secondary keys will be rolled // back unexpectedly after deletion of primary key is compacted. for range in [0..25, 25..50, 50..100] { let commit_action = if async_commit { CommitAction::AsyncCommit(Duration::MAX) } else { CommitAction::AsyncCommitSecondaryKeys(Duration::MAX) }; client .try_del_kv( range, &i_to_key, MutateOptions { commit_action, ..Default::default() }, ) .unwrap(); } // Wait for data of WRITE_CF being compacted to L1+, and trigger // `kvengine::Engine::load_unloaded_tables`. wait_for_keyspace_stats( &runtime, &cluster, &pd_client, keyspace_id, |stats| stats.cfs[WRITE_CF].levels.iter().any(|l| l.num_tables > 0), false, Duration::from_secs(10), ) .unwrap(); let truncate_ts_a = client.get_ts().into_inner(); let ref_store_a = client.dump_ref_store(); // Import data by async commit / txn file { for i in [125, 150, 220] { // Split region to make secondary batches of txn chunks. client.split(&i_to_key(i)); } let options = if async_commit { MutateOptions { commit_action: CommitAction::AsyncCommit(Duration::MAX), write_method: TxnWriteMethod::Normal, ..Default::default() } } else { MutateOptions { commit_action: CommitAction::AsyncCommitSecondaryKeys(Duration::MAX), write_method: TxnWriteMethod::FileBased, ..Default::default() } }; client.try_del_kv(100..200, &i_to_key, options).unwrap(); // Test rollback of txn file locks. if !async_commit { client .try_put_kv( 200..300, &i_to_key, i_to_val(IMPORT_DATA_LEN), MutateOptions { commit_action: CommitAction::NoCommit, write_method: TxnWriteMethod::FileBased, ..Default::default() }, ) .unwrap(); } } let origin_ref_store = client.dump_ref_store(); // Perform backup. let backup_ts = client.get_ts().into_inner(); let snapshot_backup_name = generate_backup_name(); { let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, snapshot_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); info!("backup_cluster result: {:?}", backup_meta); } let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); // Restore keyspace. { restore_keyspace::restore_keyspace( keyspace_id, keyspace_id, &snapshot_backup_name, None, s3fs.clone(), restore_config.clone(), cluster.get_pd_client(), None, &runtime, None, None, reporter.clone(), object_cache.clone(), limiter.clone(), ) .unwrap(); // Invoke major compaction to reproduce issue by compacting the deletion of // primary key. // If restore keyspace do not resolve locks, the secondary keys of "put_kv" will // not be compacted, and the following `wait_for_keyspace_stats` will fail. let ts = client.get_ts(); runtime .block_on(pd_client.advance_txn_safe_point(keyspace_id, ts)) .unwrap(); runtime .block_on(pd_client.advance_gc_safe_point(keyspace_id, ts)) .unwrap(); runtime .block_on(cluster.update_gc_states_immediately()) .unwrap(); request_major_compaction(&runtime, &pd_client, keyspace_id); wait_for_keyspace_stats( &runtime, &cluster, &pd_client, keyspace_id, |stats| stats.cfs[WRITE_CF].levels.iter().all(|l| l.num_tables == 0), false, Duration::from_secs(10), ) .unwrap(); // Verify restored data. let verified_cnt = client .verify_data_with_given_ref_store(&origin_ref_store, None, &RequestOptions::default()) .unwrap(); assert_eq!(verified_cnt, (0, 200)); must_unlocked(&mut client, &shared_prewrite_key); if let Some(info) = get_mvcc_info(&mut client, &shared_prewrite_key) { assert!( info.get_writes().is_empty(), "restored in-flight shared lock should not leave write CF records: {:?}", info.get_writes() ); } } // Restore keyspace with truncate ts (truncate_ts_a). // To test for resolving locks after truncate ts. { restore_keyspace::restore_keyspace( keyspace_id, keyspace_id, &snapshot_backup_name, None, s3fs, restore_config, cluster.get_pd_client(), None, &runtime, Some(truncate_ts_a), None, reporter, object_cache, limiter, ) .unwrap(); // Verify restored data. let verified_cnt = client .verify_data_with_given_ref_store(&ref_store_a, None, &RequestOptions::default()) .unwrap(); assert_eq!(verified_cnt, (100, 100)); } cluster.stop(); oss.shutdown(); } #[test] fn test_restore_keyspace_discards_committed_shared_secondary_lock_metadata() { let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_committed_shared_lock_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let _enter = runtime.enter(); let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None); let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(128); conf.coprocessor.region_split_size = ReadableSize::kb(128); conf.coprocessor.region_bucket_size = ReadableSize::kb(64); conf.rfengine.target_file_size = ReadableSize::mb(1); conf.rfengine.lightweight_backup = true; conf.rfengine.wal_chunk_target_file_size = ReadableSize::kb(128); }) .pd(pd_wrapper) .build(); cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default()); cluster.wait_region_replicated(&[], 3); let keyspace_id = runtime.block_on( cluster.create_keyspace(&CreateKeyspaceOptions::default(), Duration::from_secs(10)), ); let mut client = cluster.new_client(); client.split_keyspace(keyspace_id); let i_to_key = gen_keyspace_key(keyspace_id); let primary_key = i_to_key(10); let shared_secondary_key = i_to_key(20); let primary_value = b"restore-shared-primary".to_vec(); // Build the exact leftover state this compatibility test cares about: // the transaction has a normal primary key and a shared-lock-only // secondary key. The shared secondary first goes through shared // pessimistic locking so prewrite updates an existing SharedLocks wrapper // instead of writing a plain lock value. let start_ts = client.get_ts(); let for_update_ts = client.get_ts(); client .kv_pessimistic_lock( primary_key.clone().into(), vec![primary_key.clone().into()], start_ts, 20_000, for_update_ts, None, ) .expect("pessimistic lock primary"); client .kv_shared_pessimistic_lock( primary_key.clone().into(), vec![shared_secondary_key.clone().into()], start_ts, 20_000, for_update_ts, None, ) .expect("shared pessimistic lock secondary"); // Prewrite both keys, then commit only the primary. This simulates the // online window where the primary is already committed but the shared // secondary wrapper is still left in LOCK_CF when native BR takes backup. let txn_muts = TxnMutations::from_normal(vec![ Mutation { op: kvproto::kvrpcpb::Op::Put, key: primary_key.clone().into(), value: primary_value.clone().into(), ..Default::default() }, new_shared_lock_mutation(shared_secondary_key.clone()), ]); let mut txn = client.begin_transaction(Some(start_ts)); client .kv_prewrite_ext( primary_key.clone().into(), None, txn_muts.clone(), &mut txn, PrewriteExt::with_for_update_ts(for_update_ts), ) .expect("prewrite put primary and shared secondary"); let commit_ts = client.get_ts(); client .kv_commit_ext( txn_muts, start_ts, commit_ts, CommitAction::AsyncCommitSecondaryKeys(Duration::MAX), ) .expect("commit primary only"); // Sanity-check the backup source state: user-visible primary data is // committed, while the shared secondary still exists only as lock metadata. must_unlocked(&mut client, &primary_key); must_written_with_properties( &mut client, &primary_key, start_ts, commit_ts, WriteType::Put, ); // After shared prewrite, the SharedLocks wrapper still owns this key in // LOCK_CF, but its holder has advanced from shared pessimistic lock to a // lock-only prewrite. The MVCC lock info therefore exposes `Op::Lock`, // so the helper should assert `LockType::Lock`, not `LockType::Shared`. client.must_locked_with_properties( &shared_secondary_key, &primary_key, start_ts, LockType::Lock, ); // Take a lightweight backup with the committed primary and leftover shared // secondary wrapper both present. Restore should preserve the visible // primary write and clean the shared-lock wrapper. let backup_name = generate_backup_name(); let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, backup_name.clone(), cluster.get_pd_client().as_ref(), client.get_ts().into_inner(), None, ) .expect("backup::backup_cluster"); restore_keyspace::restore_keyspace( keyspace_id, keyspace_id, &backup_name, None, s3fs, restore_config, cluster.get_pd_client(), None, &runtime, None, None, reporter, object_cache, limiter, ) .unwrap(); // Native BR intentionally uses visible-data-only semantics for shared-lock // wrappers: the primary user value and Put write must be restored, but the // shared-lock-only secondary should be unlocked without materializing a // committed WriteType::Lock record. let (restored_value, _) = client.must_get_key(&primary_key, time::Instant::now()); assert_eq!(restored_value, primary_value); must_unlocked(&mut client, &primary_key); must_written_with_properties( &mut client, &primary_key, start_ts, commit_ts, WriteType::Put, ); must_unlocked(&mut client, &shared_secondary_key); if let Some(info) = get_mvcc_info(&mut client, &shared_secondary_key) { assert!( info.get_writes().is_empty(), "restored committed shared secondary should not leave write CF records: {:?}", info.get_writes() ); } cluster.stop(); oss.shutdown(); } #[rstest] #[case::two_pc(false)] #[case::async_commit(true)] fn test_restore_keyspace_with_leftover_pessimistic_lock(#[case] async_commit: bool) { let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_pessimistic_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let _enter = runtime.enter(); let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None); let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rocksdb.writecf.write_buffer_size = ReadableSize::kb(128); conf.coprocessor.region_split_size = ReadableSize::kb(128); conf.coprocessor.region_bucket_size = ReadableSize::kb(64); conf.rfengine.target_file_size = ReadableSize::mb(1); conf.rfengine.lightweight_backup = true; conf.rfengine.wal_chunk_target_file_size = ReadableSize::kb(128); // This test verifies restoring leftover pessimistic locks that were // persisted in LOCK_CF. In-memory locks are intentionally cleared // on restore and do not produce the expected WriteType::Lock record. conf.pessimistic_txn.in_memory = false; }) .pd(pd_wrapper) .build(); cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default()); cluster.wait_region_replicated(&[], 3); let keyspace_id = runtime.block_on( cluster.create_keyspace(&CreateKeyspaceOptions::default(), Duration::from_secs(10)), ); let mut client = cluster.new_client(); if async_commit { client.set_async_commit(); } client.split_keyspace(keyspace_id); let i_to_key = gen_keyspace_key(keyspace_id); let primary_key = i_to_key(10); let leftover_key = i_to_key(200); client.try_split(&leftover_key, 5).unwrap(); cluster.wait_pd_region_min_count(4); let start_ts = client.get_ts(); let for_update_ts = client.get_ts(); let (_lock_results, key_errors) = client .kv_pessimistic_lock( primary_key.clone().into(), vec![primary_key.clone().into(), leftover_key.clone().into()], start_ts, 20_000, for_update_ts, None, ) .expect("pessimistic lock"); assert!( key_errors.is_empty(), "unexpected key_errors: {:?}", key_errors ); let mut txn = client.begin_transaction(Some(start_ts)); let mut options = MutateOptions::pessimistic(start_ts, for_update_ts); options.commit_action = if async_commit { CommitAction::AsyncCommit(Duration::ZERO) } else { CommitAction::SyncCommit }; let commit_ts = client .prewrite_pessimistic_and_commit( vec![Mutation { op: kvproto::kvrpcpb::Op::Put, key: primary_key.clone().into(), value: b"restore-lock".to_vec().into(), ..Default::default() }], vec![], options, PrewriteExt::with_for_update_ts(for_update_ts), &mut txn, ) .unwrap(); must_wait( || { client .kv_get_mvcc_by_key(&primary_key) .ok() .flatten() .is_some_and(|info| { !info.has_lock() && info.get_writes().iter().any(|write| { write.commit_ts == commit_ts.into_inner() && write.get_start_ts() == start_ts.into_inner() && write.get_type() == kvproto::kvrpcpb::Op::Put }) }) }, 10, || { format!( "primary key {} was not committed before backup", log_wrappers::Value::key(&primary_key) ) }, ); client.must_locked_with_properties( &leftover_key, &primary_key, start_ts, LockType::Pessimistic, ); must_written_with_properties( &mut client, &primary_key, start_ts, commit_ts, WriteType::Put, ); let backup_name = generate_backup_name(); let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, backup_name.clone(), cluster.get_pd_client().as_ref(), client.get_ts().into_inner(), None, ) .expect("backup::backup_cluster"); restore_keyspace::restore_keyspace( keyspace_id, keyspace_id, &backup_name, None, s3fs, restore_config, cluster.get_pd_client(), None, &runtime, None, None, reporter, object_cache, limiter, ) .unwrap(); must_unlocked(&mut client, &primary_key); must_unlocked(&mut client, &leftover_key); must_written_with_properties( &mut client, &primary_key, start_ts, commit_ts, WriteType::Put, ); must_written_with_properties( &mut client, &leftover_key, start_ts, commit_ts, WriteType::Lock, ); client.verify_data_with_ref_store(); cluster.stop(); oss.shutdown(); } #[test] fn test_restore_keyspace_restore_on_restore() { test_util::init_log_for_test(); const KEYSPACE_ID: u32 = 1; const SPLIT_KEY: usize = 100; const LEFT_KEY_A: usize = 0; const RIGHT_KEY_C: usize = 200; const RFENGINE_TARGET_FILE_SIZE: usize = 16 * 1024; let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rfengine.target_file_size = ReadableSize::kb(16); conf.rfengine.wal_chunk_target_file_size = ReadableSize::kb(128); conf.rfengine.lightweight_backup = true; conf.memory.enable_heap_profiling = false; }); cluster.wait_region_replicated(&[], 3); let mut client = cluster.new_client(); client.split_keyspace(KEYSPACE_ID); let i_to_key = gen_keyspace_key(KEYSPACE_ID); let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; // Write data a to the left region and create backup A. client.put_kv( LEFT_KEY_A..(LEFT_KEY_A + 1), &i_to_key, i_to_val(BASIC_DATA_LEN), ); client.verify_data_with_ref_store(); let ref_store_a = client.dump_ref_store(); wait_all_dfs_workers_healthy(&cluster, "dfs workers healthy before backup A"); let backup_a = generate_backup_name(); backup::backup_cluster_with_ts( backup_config.clone(), BackupType::Lightweight, backup_a.clone(), cluster.get_pd_client().as_ref(), client.get_ts().into_inner(), None, ) .expect("backup::backup_cluster for backup A"); let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); client.try_split(&i_to_key(SPLIT_KEY), 5).unwrap(); cluster.wait_pd_region_min_count(4); // Restore to backup A. restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &backup_a, None, s3fs.clone(), restore_config.clone(), cluster.get_pd_client(), None, &runtime, None, None, reporter.clone(), object_cache.clone(), limiter.clone(), ) .unwrap(); let verified_cnt = client .verify_data_with_given_ref_store(&ref_store_a, None, &RequestOptions::default()) .unwrap(); assert_eq!(verified_cnt, (1, 0)); // Write large data c to the right region to trigger backup snapshot. // Without the backup snapshot, the restore will be rejected because we can not // preprocess restore_shard. for k in RIGHT_KEY_C..(RIGHT_KEY_C + 16) { client.put_kv( std::iter::once(k), &i_to_key, i_to_val(RFENGINE_TARGET_FILE_SIZE), ); } let ref_store_c = client.dump_ref_store(); wait_all_dfs_workers_healthy(&cluster, "dfs workers healthy before backup B"); let backup_b = generate_backup_name(); backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, backup_b.clone(), cluster.get_pd_client().as_ref(), client.get_ts().into_inner(), None, ) .expect("backup::backup_cluster for backup B"); // Write more data to verify the restore below. client.put_kv( LEFT_KEY_A + 1..(LEFT_KEY_A + 2), &i_to_key, i_to_val(BASIC_DATA_LEN), ); let verified_cnt = client.verify_data_with_ref_store(); assert_eq!(verified_cnt, (18, 0)); // Restore to backup B and verify state c. // The left region just restore to previous restore. restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &backup_b, None, s3fs, restore_config, cluster.get_pd_client(), None, &runtime, None, None, reporter, object_cache, limiter, ) .unwrap(); let verified_cnt = client .verify_data_with_given_ref_store(&ref_store_c, None, &RequestOptions::default()) .unwrap(); assert_eq!(verified_cnt, (17, 0)); cluster.stop(); oss.shutdown(); } /// This test is to verify that `restore_keyspace` can handle the case that in /// lightweight backup, the last epoch has no chunk in DFS. /// /// See https://github.com/tidbcloud/cloud-storage-engine/issues/1319. #[test] fn test_restore_keyspace_with_no_chunk() { const KEYSPACE_ID: u32 = 1; let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let mut cluster = ServerCluster::new( alloc_node_id_vec(NODES_COUNT), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rfengine.target_file_size = ReadableSize::mb(512); // To prevent WAL to compact. conf.rfengine.wal_chunk_target_file_size = ReadableSize::mb(64); conf.rfengine.lightweight_backup = true; }, ); cluster.wait_region_replicated(&[], 3); let mut client = cluster.new_client(); client.split_keyspace(KEYSPACE_ID); // Import small data let i_to_key = gen_keyspace_key(KEYSPACE_ID); client.put_kv(0..1, &i_to_key, i_to_val(BASIC_DATA_LEN)); client.verify_data_with_ref_store(); // Perform backup. let snapshot_backup_name = generate_backup_name(); { let backup_ts = client.get_ts().into_inner(); let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, snapshot_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); info!("backup_cluster result: {:?}", backup_meta); } // Restore keyspace. let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &snapshot_backup_name, None, s3fs, restore_config, cluster.get_pd_client(), None, &runtime, None, None, reporter, object_cache, limiter, ) .unwrap(); // Verify restored data. client.verify_data_with_ref_store(); cluster.stop(); oss.shutdown(); } /// This test is to verify that `restore_keyspace` can properly handle when DFS /// is slow. /// /// See https://github.com/tidbcloud/cloud-storage-engine/issues/1977. #[test] #[ignore] fn test_restore_keyspace_with_slow_dfs() { const KEYSPACE_ID: u32 = 1; let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let mut cluster = ServerCluster::new( alloc_node_id_vec(NODES_COUNT), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rfengine.target_file_size = ReadableSize::kb(512); conf.rfengine.wal_chunk_target_file_size = ReadableSize::kb(32); conf.rfengine.lightweight_backup = true; }, ); cluster.wait_region_replicated(&[], 3); let mut client = cluster.new_client(); client.split_keyspace(KEYSPACE_ID); // Import data. let i_to_key = gen_keyspace_key(KEYSPACE_ID); for i in 0..100 { client.put_kv(i * 10..(i + 1) * 10, &i_to_key, i_to_val(BASIC_DATA_LEN)); } client.verify_data_with_ref_store(); // Workload to continuously generate WAL. let (tx, rx) = std::sync::mpsc::sync_channel(0); { let mut client = cluster.new_client(); std::thread::spawn(move || { while let Err(std::sync::mpsc::TryRecvError::Empty) = rx.try_recv() { let _ = client.try_put_kv( 10000..10050, &i_to_key, i_to_val(POST_BACKUP_DATA_LEN), MutateOptions::default(), ); let _ = client.try_del_kv(10000..10050, &i_to_key, MutateOptions::default()); } }); } // Delay "wal_chunks" but not snapshots. oss.set_put_delay("wal_chunks", Duration::from_secs(3)); // Perform backup. let snapshot_backup_name = generate_backup_name(); { let backup_ts = client.get_ts().into_inner(); let backup_config = backup::BackupConfig { dfs: dfs_config.clone(), backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, snapshot_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); info!("backup_cluster result: {:?}", backup_meta); } let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); NATIVE_BR_RFENGINE_WAL_EPOCH_OVERWRITTEN_ERROR.reset(); let mut ok = false; for _ in 0..30 { // Restore keyspace. restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &snapshot_backup_name, None, s3fs.clone(), restore_config.clone(), cluster.get_pd_client(), None, &runtime, None, None, reporter.clone(), object_cache.clone(), limiter.clone(), ) .unwrap(); // Note: the counter is not reliable enough that the error must happen, as there // are other concurrent test cases. But it's enough to verify the // process. if NATIVE_BR_RFENGINE_WAL_EPOCH_OVERWRITTEN_ERROR.get() > 0 { ok = true; break; } std::thread::sleep(Duration::from_secs(1)); } assert!(ok); tx.send(()).unwrap(); // Verify restored data. client.verify_data_with_ref_store(); cluster.stop(); oss.shutdown(); } #[rstest] #[case::backup_ts(false)] #[case::pitr(true)] fn test_restore_keyspace_with_schema(#[case] pitr: bool) { const KEYSPACE_ID: u32 = 1; const DB_ID: i64 = 100; const TABLE_ID: i64 = 200; let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); let mut cluster = ServerCluster::new( alloc_node_id_vec(NODES_COUNT), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rfengine.lightweight_backup = true; }, ); cluster.wait_region_replicated(&[], 3); let mut client = cluster.new_client(); client.split_keyspace(KEYSPACE_ID); // Import small data let i_to_key = gen_keyspace_key(KEYSPACE_ID); client.put_kv(0..1, &i_to_key, i_to_val(BASIC_DATA_LEN)); client.verify_data_with_ref_store(); // Persist TiDB meta KV and install the matching schema file at version 100. let mut table_infos = None; for schema_version in [99, 100] { let pairs = generate_storage_class_schema_data_for_test( KEYSPACE_ID, DB_ID, TABLE_ID, None, None, schema_version, StorageClass::Ia.into(), &mut table_infos, ) .unwrap(); put_raw_kv_pairs(&mut client, pairs); } let schemas_at_100 = table_infos .as_ref() .unwrap() .values() .map(schema_from_table_info) .collect::, _>>() .unwrap(); let broadcast_schema_update = |client: &ClusterClient, schema_file_data: Vec, wait_applied: bool| { let file_id = client.pd_client().alloc_id().unwrap(); let schema_file_data = Bytes::from(schema_file_data); runtime .block_on(s3fs.create( file_id, schema_file_data.clone(), Options::default().with_type(FileType::Schema), )) .unwrap(); let stores = client.pd_client().get_all_stores(true).unwrap(); if wait_applied { let schema_file = SchemaFile::open(Arc::new(InMemFile::new(file_id, schema_file_data))).unwrap(); runtime.block_on(broadcast_schema_file_request_and_check( &stores, KEYSPACE_ID, &schema_file, time::Duration::from_secs(10), )); } else { runtime.block_on(broadcast_schema_file_request( &stores, KEYSPACE_ID, file_id, time::Duration::from_secs(10), )); } }; broadcast_schema_update( &client, build_schema_file(KEYSPACE_ID, 100, schemas_at_100.clone(), 0), true, ); // Check the schema file already installed. must_wait( || { let data_stats = cluster.get_data_stats(); let mut schema_file_exists = false; data_stats.iter_shard_stats(|_, shard_stats| { if shard_stats.schema_version > 0 { schema_file_exists = true; return true; } false }); schema_file_exists }, 10, || "failed to install schema file".to_string(), ); // Stabilize schema-driven table splits before installing the next version. // A new shard rejects schema updates until its initial flush is applied. runtime .block_on(client.pd_client().split_regions_with_retry( get_table_split_keys(KEYSPACE_ID, &[TABLE_ID]), Duration::from_secs(10), )) .unwrap(); let (keyspace_start, keyspace_end) = ApiV2::get_keyspace_range_by_id(KEYSPACE_ID); let regions = runtime .block_on(client.pd_client().scan_regions( encode_bytes(&keyspace_start), encode_bytes(&keyspace_end), 100, )) .unwrap(); let initial_flush_finished = try_wait( || { regions.iter().all(|region_info| { let region = region_info.get_region(); region.peers.iter().all(|peer| { cluster .get_kvengine(cluster.get_server_node_id(peer.store_id).unwrap()) .get_shard(region.id) .is_some_and(|shard| shard.get_initial_flushed()) }) }) }, 10, ); assert!(initial_flush_finished, "initial flush should finish"); let truncate_ts = client.get_ts().into_inner(); let (source_schema_version, source_schemas) = if pitr { // Commit a later DDL and install its schema file. PiTR must rebuild // version 100 from backup meta KV instead of copying this file. let pairs = generate_storage_class_schema_data_for_test( KEYSPACE_ID, DB_ID, TABLE_ID, None, None, 101, StorageClass::Standard.into(), &mut table_infos, ) .unwrap(); put_raw_kv_pairs(&mut client, pairs); let schemas_at_101 = table_infos .as_ref() .unwrap() .values() .map(schema_from_table_info) .collect::, _>>() .unwrap(); broadcast_schema_update( &client, build_schema_file(KEYSPACE_ID, 101, schemas_at_101.clone(), 0), true, ); must_wait( || { let data_stats = cluster.get_data_stats(); let mut installed = false; data_stats.iter_shard_stats(|_, shard_stats| { installed = shard_stats.schema_version == 101; installed }); installed }, 10, || "failed to install schema file version 101".to_string(), ); (101, schemas_at_101) } else { (100, schemas_at_100) }; // Perform backup. let snapshot_backup_name = generate_backup_name(); { let backup_ts = client.get_ts().into_inner(); let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), skip_keyspace_meta: true, ..Default::default() }; let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, snapshot_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); info!("backup_cluster result: {:?}", backup_meta); } // Restore keyspace. let restore_config = RestoreConfig::default_for_test(); let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &snapshot_backup_name, None, s3fs.clone(), restore_config, cluster.get_pd_client(), None, &runtime, pitr.then_some(truncate_ts), None, reporter, object_cache, limiter, ) .unwrap(); let mut restore_version = 0; let mut restored_schema_version = 0; // Check if the schema file is update the restore_version after restoration. must_wait( || { let data_stats = cluster.get_data_stats(); let mut has_restore_version = false; data_stats.iter_shard_stats(|_, shard_stats| { if shard_stats.schema_restore_version > 0 { has_restore_version = true; restore_version = shard_stats.schema_restore_version; restored_schema_version = shard_stats.schema_version; return true; } false }); has_restore_version }, 10, || "failed to check schema file restore_version".to_string(), ); assert_eq!(restored_schema_version, 100); // Try to update schema with different restore_version. let schema_file_data = build_schema_file( KEYSPACE_ID, source_schema_version, source_schemas.clone(), restore_version + 1, ); broadcast_schema_update(&client, schema_file_data, false); // Enqueue a valid, newer schema after the stale update. Its application is // a positive fence: if the stale restore version were accepted first, this // update with the original restore version would be rejected. let fence_schema_version = source_schema_version + 1; broadcast_schema_update( &client, build_schema_file( KEYSPACE_ID, fence_schema_version, source_schemas, restore_version, ), true, ); must_wait( || { let data_stats = cluster.get_data_stats(); let mut fenced = false; data_stats.iter_shard_stats(|_, shard_stats| { if shard_stats.schema_restore_version > 0 { fenced = shard_stats.schema_restore_version == restore_version && shard_stats.schema_version == fence_schema_version; return !fenced; } false }); fenced }, 10, || "failed to fence stale schema restore version".to_string(), ); // Verify restored data. client.verify_data_with_ref_store(); cluster.stop(); oss.shutdown(); } #[derive(PartialEq, Debug)] enum RestoreResult { AlwaysSucceed, AlwaysFail, DependsOnTolerateErr, } // Arguments: (fail_on_backup, fail_on_restore, fail_on_same_node, expected) // fail_on_same_node: Whether backup and restore fail on the same node. #[rstest] #[case::fail_on_restore(false, true, false, RestoreResult::DependsOnTolerateErr)] #[case::both(true, true, false, RestoreResult::AlwaysFail)] #[case::fail_on_same_node(true, true, true, RestoreResult::AlwaysSucceed)] fn test_restore_keyspace_with_failed_store( #[case] fail_on_backup: bool, #[case] fail_on_restore: bool, #[case] fail_on_same_node: bool, #[case] expected: RestoreResult, ) { info!("test_restore_keyspace_with_failed_store"; "fail_on_backup" => fail_on_backup, "fail_on_restore" => fail_on_restore, "fail_on_same_node" => fail_on_same_node, "expected" => ?expected, ); const KEYSPACE_ID: u32 = 1; let (_temp_dir, mut oss, dfs_config) = prepare_dfs("test_restore_keyspace_"); let s3fs = new_dfs_from_config(dfs_config.clone()); let reporter = Arc::new(DummyStepReporter::default()); let runtime = Runtime::new().unwrap(); // Use number of nodes of 3 to make sure that every node has WAL chunks not // flushed yet. let mut cluster = ServerCluster::new(alloc_node_id_vec(3), |_, conf: &mut TikvConfig| { conf.dfs = dfs_config.clone(); conf.rfengine.target_file_size = ReadableSize::mb(512); // To prevent WAL to compact. conf.rfengine.wal_chunk_target_file_size = ReadableSize::mb(64); conf.rfengine.lightweight_backup = true; }); cluster.wait_region_replicated(&[], 3); let mut client = cluster.new_client(); client.split_keyspace(KEYSPACE_ID); // Import small data let i_to_key = gen_keyspace_key(KEYSPACE_ID); client.put_kv(0..1, &i_to_key, i_to_val(BASIC_DATA_LEN)); client.verify_data_with_ref_store(); let origin_ref_store = client.dump_ref_store(); // Perform backup. let snapshot_backup_name = generate_backup_name(); let backup_failed_node = { let backup_ts = client.get_ts().into_inner(); let backup_config = backup::BackupConfig { dfs: dfs_config, backup_delay: ReadableDuration::secs(1), tolerate_err: 1, skip_keyspace_meta: true, timeout: ReadableDuration::secs(3), ..Default::default() }; let failed_node_id = if fail_on_backup { Some(stop_node(&mut cluster, true, false, None)) } else { None }; let (_, backup_meta) = backup::backup_cluster_with_ts( backup_config, BackupType::Lightweight, snapshot_backup_name.clone(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); info!("backup_cluster result: {:?}", backup_meta); if let Some(node_id) = failed_node_id { cluster.start_node(node_id, |_, _| {}); } failed_node_id }; // Put another data. client.put_kv(1..2, &i_to_key, i_to_val(POST_BACKUP_DATA_LEN)); client.verify_data_with_ref_store(); if fail_on_restore { // Force stop one node without flush the last WAL chunk to DFS. let restore_failed_node = match backup_failed_node { Some(node_id) => { let node_id = if fail_on_same_node { node_id } else { let mut node_ids = cluster.get_nodes(); node_ids.retain(|&id| id != node_id); *node_ids.choose(&mut rand::thread_rng()).unwrap() }; Some(node_id) } _ => None, }; let store_tombstone = thread_rng().gen_bool(0.5); stop_node(&mut cluster, true, store_tombstone, restore_failed_node); } // Set small timeouts to make the test faster. let mut restore_config = RestoreConfig { tolerate_err: 0, timeout_fetch_wal: ReadableDuration::secs(1), timeout_restore_snapshot: ReadableDuration::secs(3), strict_tolerate: true, lower_memory: RestoreConfig::use_lower_memory(), ..Default::default() }; let object_cache = cluster.create_object_cache_randomly(); let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone()); let res = restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &snapshot_backup_name, None, s3fs.clone(), restore_config.clone(), cluster.get_pd_client(), None, &runtime, None, None, reporter.clone(), object_cache.clone(), limiter.clone(), ); assert_eq!( res.is_ok(), expected == RestoreResult::AlwaysSucceed, "res: {:?}, expected: {:?}", res, expected ); // Update tolerate error config and restore keyspace again. restore_config.tolerate_err = 1; // Restore keyspace. let res_tolerated = restore_keyspace::restore_keyspace( KEYSPACE_ID, KEYSPACE_ID, &snapshot_backup_name, None, s3fs, restore_config, cluster.get_pd_client(), None, &runtime, None, None, reporter, object_cache, limiter, ); assert_eq!( res_tolerated.is_err(), expected == RestoreResult::AlwaysFail, "res: {:?}, expected: {:?}", res, expected ); // Verify restored data. if res.is_err() && res_tolerated.is_err() { client.verify_data_with_ref_store(); } else { client .verify_data_with_given_ref_store(&origin_ref_store, None, &RequestOptions::default()) .unwrap(); } cluster.stop(); oss.shutdown(); } fn write_simple_data(cluster: &ServerCluster, keyspace: u32, expect_blob_file: bool) { let mut client = cluster.new_client_opt(ClusterClientOptions { with_lock_resolver: false, txn_file_max_chunk_size: Some(102400), ..Default::default() }); let i_to_key = gen_keyspace_key(keyspace); for i in 0..16 { client.put_kv(i..i + 1, &i_to_key, i_to_val(LONG_DATA_LEN)); } let region = cluster.get_pd_client().get_region(&i_to_key(0)).unwrap(); cluster.flush_memtable(region.id).unwrap(); cluster.wait_for_memtable_flushed(region.id, Duration::from_secs(10)); client.split(&i_to_key(256)); client.put_kv(16..512, &i_to_key, i_to_val(BASIC_DATA_LEN)); client.split(&i_to_key(512)); client.put_kv(512..1040, &i_to_key, i_to_val(BASIC_DATA_LEN)); let region = cluster.get_pd_client().get_region(&i_to_key(513)).unwrap(); cluster.flush_memtable(region.id).unwrap(); cluster.wait_for_memtable_flushed(region.id, Duration::from_secs(10)); let handle = Handle::current(); let mut muts = vec![]; for i in 1040..2020 { let mut m = Mutation::default(); m.set_key(i_to_key(i).to_vec()); m.set_value(i_to_val(BASIC_DATA_LEN)(0)); m.set_op(kvproto::kvrpcpb::Op::Put); muts.push(m); } let chunks = handle .block_on(TxnMutations::build( muts, vec![], TxnWriteMethod::FileBased, client.txn_file_helper(), )) .unwrap(); let start_ts = client.get_ts(); client .kv_prewrite(chunks.primary(), None, chunks.clone(), start_ts) .unwrap(); client .kv_commit(chunks, start_ts, client.get_ts(), false) .unwrap(); request_major_compaction(&runtime(), &cluster.get_pd_client(), keyspace); if expect_blob_file { must_wait( || { cluster.get_nodes().iter().any(|&node_id| { cluster .get_kvengine(node_id) .get_all_shard_stats() .iter() .any(|stats| { stats.active && stats.keyspace == keyspace && stats.blob_table_count > 0 }) }) }, 10, || "wait for packed-backup blob file".to_string(), ); } // Put something to memtable. client.put_kv(2020..2048, &i_to_key, i_to_val(BASIC_DATA_LEN)); } #[fixture] fn runtime() -> Runtime { tokio::runtime::Builder::new_multi_thread() .enable_all() .build() .unwrap() } struct TestDfs { _temp_dir: TempDir, oss: ObjectStorageService, cfg: DFSConfig, } impl TestDfs { fn s3fs(&self) -> S3Fs { self.prefixed_s3fs("") } fn prefixed_s3fs(&self, prefix: &str) -> S3Fs { S3Fs::new( format!("{}{}", self.cfg.prefix.clone(), prefix), self.cfg.s3_endpoint.clone(), self.cfg.s3_key_id.clone(), self.cfg.s3_secret_key.clone(), self.cfg.s3_region.clone(), self.cfg.s3_bucket.clone(), self.cfg.conn_options.clone(), self.cfg.read_only, ) } } #[fixture] fn dfs() -> TestDfs { let (temp_dir, oss, cfg) = prepare_dfs(&format!("test_{}", rand::random::())); TestDfs { _temp_dir: temp_dir, oss, cfg, } } #[rstest] #[case::lightweight_rewrite_keyspace_copy(BackupType::Lightweight, 2, true)] #[case::lightweight_rewrite_keyspace_no_copy(BackupType::Lightweight, 2, false)] #[case::full_same_keyspace_copy(BackupType::Full, 1, true)] #[case::full_rewrite_keyspace_no_copy(BackupType::Full, 2, false)] fn test_restore_packed_backup( #[case] backup_type: BackupType, #[case] target_keyspace: u32, #[case] copy_to_elsewhere: bool, runtime: Runtime, dfs: TestDfs, ) { run_test_restore_packed_backup( backup_type, target_keyspace, copy_to_elsewhere, false, runtime, dfs, ); } #[rstest] fn test_restore_packed_backup_with_cmek(runtime: Runtime, dfs: TestDfs) { run_test_restore_packed_backup(BackupType::Lightweight, 2, false, true, runtime, dfs); } fn run_test_restore_packed_backup( backup_type: BackupType, target_keyspace: u32, copy_to_elsewhere: bool, enable_cmek: bool, runtime: Runtime, mut dfs: TestDfs, ) { use uuid::Uuid; info!("Running case test_restore_packed_backup"; "backup_type" => ?backup_type, "target_keyspace" => target_keyspace, "copied_to_elsewhere" => copy_to_elsewhere, "enable_cmek" => enable_cmek, ); const KEYSPACE_ID: u32 = 1; let _g = runtime.enter(); // let mut s3fs = Arc::new(dfs.s3fs()); let tmpfs = Arc::new(dfs.prefixed_s3fs(&format!("/packed-backup/flush_temp/{}", Uuid::new_v4()))); let security_conf = if enable_cmek { new_security_config() } else { SecurityConfig::default() }; let pd = PdWrapper::new_test(1, &security_conf, None); let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf| { conf.dfs = dfs.cfg.clone(); conf.rfengine.target_file_size = ReadableSize::mb(512); // To prevent WAL to compact. conf.rfengine.wal_chunk_target_file_size = ReadableSize::mb(64); conf.rfengine.lightweight_backup = true; conf.raft_store.enable_inner_key_offset = true; conf.security = security_conf.clone(); conf.kvengine .per_keyspace_configs .push(KvEnginePerKeyspaceConfig { keyspace: KEYSPACE_ID, enable_blob: !enable_cmek, ..Default::default() }) }) .pd(pd) .build(); let mut opt = TikvWorkerOptions::default(); opt.kv_target_file_size = ReadableSize::mb(1); cluster.start_tikv_workers(alloc_node_id_vec(1), opt); cluster.wait_region_replicated(&[], 3); let mut client = cluster.new_client(); create_keyspace_meta(&cluster, KEYSPACE_ID, &runtime); if target_keyspace != KEYSPACE_ID { create_keyspace_meta(&cluster, target_keyspace, &runtime); } if enable_cmek { enable_test_cmek(cluster.get_pd_client(), KEYSPACE_ID); } client.split_keyspace(KEYSPACE_ID); client.split_keyspace(target_keyspace); write_simple_data(&cluster, KEYSPACE_ID, !enable_cmek); let backup_ts = client.get_ts().into_inner(); let backup_config = backup::BackupConfig { dfs: dfs.cfg.clone(), skip_keyspace_meta: true, ..Default::default() }; let (path, backup_meta) = backup::backup_cluster_with_ts( backup_config, backup_type, "".to_owned(), cluster.get_pd_client().as_ref(), backup_ts, None, ) .expect("backup::backup_cluster"); let mut ref_store = client.dump_ref_store(); client.del_kv(0..2048, i_to_key); let work_dir = TempDir::new().unwrap(); cluster.stop(); let pd_client = Arc::new( runtime .block_on(OfflinePd::new( backup_meta.cluster_id, backup_meta.keyspace_meta.clone(), )) .unwrap(), ) as Arc<_>; let mut keyspace_meta = KeyspaceMeta::default(); keyspace_meta.id = KEYSPACE_ID; keyspace_meta.name = format!("ks{KEYSPACE_ID}"); let mut ctx = PackContext::create_from_env( PackEnv { pd_client, dfs: s3fs.clone(), tmpfs: tmpfs.clone(), work_path: work_dir.path().to_owned(), restore_conf: Default::default(), _temp_dir: None, reporter: Arc::new(DummyStepReporter {}), }, backup_meta, None, keyspace_meta, ) .unwrap(); let (mut packed_path, _) = runtime.block_on(ctx.execute(&path)).unwrap(); // flush all shards for node in cluster.get_nodes() { futures::executor::block_on(cluster.get_rfengine(node).upload_wal_chunk()); } cluster.stop(); let new_files = ctx.overlay_fs().modified_files(); if new_files.is_empty() { warn!("no new file was flushed by this `pack` call. Maybe background flush triggered?"); } for (ty, id) in new_files { // Make sure the new flushed file not directly put to the upstream cluster's. let exists = runtime .block_on(s3fs.exists(id, Options::default().with_type(ty))) .unwrap(); assert!(!exists, "conflicting file found {}:{}", ty, id); } if copy_to_elsewhere { let s3fs_mid = Arc::new(dfs.prefixed_s3fs("/mid")); let copy_env = runtime .block_on(MigratePackEnv::load_exotic(s3fs_mid.clone(), &packed_path)) .unwrap(); let mut run = CopyPackedRun::new(copy_env, "/the-land-between", Arc::new(NoopReporter)); let copied = runtime.block_on(run.execute()).unwrap(); packed_path = copied.new_exotic_path; } // move to another cluster. let s3fs_mkii = dfs.prefixed_s3fs("/mkii"); cluster = ServerCluster::new( alloc_node_id_vec(NODES_COUNT), |_, conf: &mut TikvConfig| { conf.dfs = dfs.cfg.clone(); conf.dfs.prefix = format!("{}/mkii", conf.dfs.prefix); conf.rfengine.target_file_size = ReadableSize::mb(512); // To prevent WAL to compact. conf.rfengine.wal_chunk_target_file_size = ReadableSize::mb(64); conf.rfengine.lightweight_backup = true; conf.raft_store.enable_inner_key_offset = true; conf.security = security_conf.clone(); }, ); cluster.wait_region_replicated(&[], 3); client = cluster.new_client(); create_keyspace_meta(&cluster, KEYSPACE_ID, &runtime); if target_keyspace != KEYSPACE_ID { create_keyspace_meta(&cluster, target_keyspace, &runtime); } client.split_keyspace(KEYSPACE_ID); client.split_keyspace(target_keyspace); s3fs = Arc::new(s3fs_mkii); let env = runtime .block_on(MigratePackEnv::load_exotic(s3fs.clone(), &packed_path)) .unwrap(); let packed_path = runtime .block_on( UnpackRun::new(env, cluster.get_pure_pd_client(), Arc::new(NoopReporter)).execute(), ) .unwrap(); restore_keyspace::restore_keyspace( KEYSPACE_ID, target_keyspace, &packed_path, None, s3fs.clone(), RestoreConfig { tolerate_err: 0, timeout_restore_snapshot: ReadableDuration::secs(20), restore_packed_backup: true, ..Default::default() }, cluster.get_pd_client(), None, &runtime, None, None, Arc::new(DummyStepReporter::default()), None, None, ) .unwrap(); if target_keyspace != KEYSPACE_ID { ref_store.rewrite_keyspace_prefix(target_keyspace); } client .verify_data_with_given_ref_store(&ref_store, None, &Default::default()) .unwrap(); if enable_cmek { assert_test_cmek(cluster.get_pd_client().as_ref(), target_keyspace); } cluster.stop(); dfs.oss.shutdown() } // Stop random node when `node_id` is `None`. fn stop_node( cluster: &mut ServerCluster, force: bool, store_tombstone: bool, node_id: Option, ) -> u16 { let failed_node_id = node_id.unwrap_or_else(|| { let node_ids = cluster.get_nodes(); let mut rng = rand::thread_rng(); *node_ids.choose(&mut rng).unwrap() }); let store_id = cluster.get_store_id(failed_node_id); cluster.stop_node_force(failed_node_id, force); if store_tombstone { let mut store = cluster.get_pd_client().get_store(store_id).unwrap(); store.state = metapb::StoreState::Tombstone; cluster.get_pd_client().put_store(store).unwrap(); } info!("node stopped"; "node_id" => failed_node_id, "store_id" => store_id, "tombstone" => store_tombstone); failed_node_id } // Add learners to all stores with no peer of keyspace regions. // TODO: add to `test_cloud_server` for reuse. fn add_learners(cluster: &mut ServerCluster, start: &[u8], end: &[u8]) { let pd_client = cluster.get_pd_client(); let mut client = cluster.new_client(); let stores = cluster.get_stores(); client.clear_region_cache(); let mut next_key = start.to_owned(); while next_key.as_slice() < end { let region = client.get_region_by_key(&next_key); next_key = region.raw_end().to_owned(); let learner_stores = stores .iter() .filter(|&&store_id| region.peers.iter().all(|peer| peer.store_id != store_id)) .collect::>(); for &store_id in learner_stores { let learner_peer = new_learner_peer(store_id, pd_client.alloc_id().unwrap()); pd_client.must_add_peer(region.id, learner_peer); } } } fn check_learners_impl( cluster: &ServerCluster, start: &[u8], end: &[u8], min_count: usize, ) -> Vec { let mut client = cluster.new_client(); let mut res = vec![]; let mut next_key = start.to_owned(); while next_key.as_slice() < end { let region = client.get_region_by_key(&next_key); next_key = region.raw_end().to_owned(); let learner_cnt = region .peers .iter() .filter(|peer| peer.get_role() == metapb::PeerRole::Learner) .count(); debug!( "check region for learner: {:?}, count: {}", region, learner_cnt ); if learner_cnt < min_count { res.push(region); } } res } fn check_learners( cluster: &ServerCluster, start: &[u8], end: &[u8], min_count: usize, timeout_seconds: usize, ) { let ok = try_wait( || check_learners_impl(cluster, start, end, min_count).is_empty(), timeout_seconds, ); assert!( ok, "{:?}", check_learners_impl(cluster, start, end, min_count) ); }