use std::{sync::Arc, time::Duration}; use anyhow::{Ok, Result}; use arc_swap::ArcSwap; use log::info; use scopeguard::defer; use storage_backend::storage::create_backend_storage_for_test; use tici::{ HandleRange, RangeBound, ShardBuildProgress, ShardInfo, ShardManifest, SplittableFrag, config::{make_default_test_meta_config, make_default_test_worker_config}, connection_pool::TiDBConnectionPool, meta::{ Bootstrap, MetaServiceCore, MetaStore, MetaStoreImplTiDB, metrics::{INDEX_COUNT_GAUGE_VEC, SHARD_COUNT_GAUGE_VEC}, writer::WorkerNodeMgr, }, worker::mocks::{make_frag_meta, make_frag_meta_by_range}, }; use tici_common::{ codec::{table_end_key, table_raw_key, table_start_key}, index::IndexKey, }; use tici_grpc::{AppendBaseFragMetaRequest, CreateIndexRequestExt, ErrorCode}; use tici_service_discovery::election::leadership::MockLeadership; use tici_test_utils::test_utils::{test_suite::E2ETestSuite, *}; use tokio::sync::RwLock; #[test] fn test_meta_server_bootstrap() -> Result<()> { fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap(); defer!({ fail::remove("worker_node_client_fail_fast"); }); let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { info!("Start bootstrap shard test..."); let config = make_default_test_meta_config(true); let tidb_pools = TiDBConnectionPool::new(&config.tidb_server); let bootstrap = Bootstrap::new(tidb_pools.clone()); bootstrap.clear_all_meta(false).await.unwrap(); assert!(!bootstrap.is_bootstrapped().await.unwrap()); bootstrap.ensure_bootstrapped().await.unwrap(); assert!(bootstrap.is_bootstrapped().await.unwrap()); let backend_storage = create_backend_storage_for_test(&config.s3).await.unwrap(); let meta_service = MetaServiceCore::new( ArcSwap::new(Arc::new(config.clone())), tidb_pools, MockLeadership::default(), backend_storage, ); meta_service.start().await.unwrap(); let (table_info, index_info, _partition_table_info) = { let req = workload_ci::make_create_index_req(&IndexKey::new_nullspace(1, 2)); req.get_grpc_table_info_and_index_info().unwrap() }; let _ = meta_service.bootstrap_index(table_info, index_info).await; assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 1 ); WorkerNodeMgr::send_create_worker_node( meta_service.worker_node_mgr.get_tx(), "AAA".to_string(), None, meta_service.meta_store.clone(), ) .await .unwrap(); assert_eq!( meta_service .shard_task_manager .bootstrap_shard( &IndexKey::new_nullspace(1, 2), &HandleRange::new(vec![b'A'], vec![b'B']), 11, ShardManifest::new_for_first(), "AAA".to_string(), ShardBuildProgress::default(), ) .await .unwrap() .read() .await .shard_id, 1 ); assert_eq!( SHARD_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 1 ); assert!( meta_service .shard_task_manager .bootstrap_shard( &IndexKey::new_nullspace(1, 2), &HandleRange::new(vec![b'A'], vec![b'B']), 11, ShardManifest::new_for_first(), "AAA".to_string(), ShardBuildProgress::default(), ) .await .is_err() ); meta_service.stop().await.unwrap(); assert!( Bootstrap::new_with_config(&config) .is_bootstrapped() .await .unwrap() ); let tidb_pools = TiDBConnectionPool::new(&config.tidb_server); let meta_store2 = MetaStoreImplTiDB::new(tidb_pools, MockLeadership::default()); let v = meta_store2.restore_shards().await.unwrap(); assert_eq!(v.len(), 1); assert_eq!(v[0].table_id(), 1); assert_eq!(v[0].shard_id, 1); assert_eq!(v[0].writer, "AAA".to_string()); }); Ok(()) } #[test] fn test_create_index_retry_bootstraps_missing_initial_shard() -> Result<()> { fail::cfg("worker_node_client_fail_fast", "return(2)").unwrap(); defer!({ fail::remove("worker_node_client_fail_fast"); }); let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let config = make_default_test_meta_config(true); let tidb_pools = TiDBConnectionPool::new(&config.tidb_server); let bootstrap = Bootstrap::new(tidb_pools.clone()); bootstrap.clear_all_meta(false).await.unwrap(); bootstrap.ensure_bootstrapped().await.unwrap(); let backend_storage = create_backend_storage_for_test(&config.s3).await.unwrap(); let storage = backend_storage.clone(); let meta_service = MetaServiceCore::new( ArcSwap::new(Arc::new(config.clone())), tidb_pools, MockLeadership::default(), backend_storage, ); meta_service.start().await.unwrap(); let table_id = rand::random_range(10_000_i64..=i64::MAX / 2); let index_key = IndexKey::new_nullspace(table_id, 3); let request = workload_ci::make_create_index_req(&index_key); let cdc_prefix = tici_common::s3_keys::cdc_path_prefix(&config.s3.prefix, table_id); let first_cdc_file = format!("{cdc_prefix}/465312161095221285/CDC00000000000000000001.json"); let initial_cdc_file = format!("{cdc_prefix}/465312161095221285/CDC00000000000000000002.json"); let cdc_index_file = format!("{cdc_prefix}/465312161095221285/meta/CDC.index"); storage .put_data(&config.s3.bucket, &first_cdc_file, b"{}\n".to_vec()) .await .unwrap(); storage .put_data(&config.s3.bucket, &initial_cdc_file, b"{}\n".to_vec()) .await .unwrap(); storage .put_data( &config.s3.bucket, &cdc_index_file, b"CDC00000000000000000002.json\n".to_vec(), ) .await .unwrap(); let first = meta_service.create_index(&request).await.unwrap(); assert_eq!(first.status, ErrorCode::UnknownError as i32); assert_eq!( meta_service .shard_task_manager .get_index(&index_key) .await .unwrap() .shards_len() .await, 0 ); assert!( meta_service .meta_store .restore_shards() .await .unwrap() .is_empty() ); WorkerNodeMgr::send_create_worker_node( meta_service.worker_node_mgr.get_tx(), "AAA".to_string(), None, meta_service.meta_store.clone(), ) .await .unwrap(); let retry = meta_service.create_index(&request).await.unwrap(); assert_eq!(retry.status, ErrorCode::Success as i32); let index = meta_service .shard_task_manager .get_index(&index_key) .await .unwrap(); assert_eq!(index.shards_len().await, 1); let shard = index.select_first_shard().await.unwrap(); let shard = shard.read().await; assert_eq!( shard.range, HandleRange::new( table_start_key(&index_key).to_vec(), table_end_key(&index_key).to_vec(), ) ); assert_eq!(shard.writer, "AAA"); assert_eq!(shard.progress.cdc_s3_last_file, initial_cdc_file); drop(shard); let later_cdc_file = format!("{cdc_prefix}/465312161095221285/CDC00000000000000000003.json"); storage .put_data(&config.s3.bucket, &later_cdc_file, b"{}\n".to_vec()) .await .unwrap(); let third = meta_service.create_index(&request).await.unwrap(); assert_eq!(third.status, ErrorCode::Success as i32); assert_eq!(index.shards_len().await, 1); let shard = index.select_first_shard().await.unwrap(); let shard = shard.read().await; assert_eq!(shard.progress.cdc_s3_last_file, initial_cdc_file); drop(shard); assert_eq!( meta_service .meta_store .restore_shards() .await .unwrap() .len(), 1 ); meta_service.stop().await.unwrap(); }); Ok(()) } fn test_create_shard(drop_add_shard: bool, restart_meta_after_add: bool) -> Result<()> { let rt: tokio::runtime::Runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; fail::cfg("worker_service_mock_meta_service_url", "return").unwrap(); if drop_add_shard { fail::cfg("worker_service_mock_drop_add_shard", "return").unwrap(); } defer!({ fail::remove("worker_service_mock_meta_service_url"); fail::remove("worker_service_mock_drop_add_shard"); }); info!( "Start create shard test drop_add_shard={} restart_meta_after_add={} ...", drop_add_shard, restart_meta_after_add ); rt.block_on(async { let mut test_suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); if !drop_add_shard { test_suite .get_worker_node_handle("127.0.0.1:8510") .shard_write_manager .as_mock() .check_to_add_shard .store(true, std::sync::atomic::Ordering::Relaxed); } let index_key = IndexKey::new_nullspace(1, 1001); test_suite.must_create_index(&index_key).await.unwrap(); test_suite.check_meta_subscriber().await; if restart_meta_after_add { info!("restart meta service"); test_suite .restart_meta(make_default_test_meta_config(true)) .await .unwrap(); } tokio::time::sleep(std::time::Duration::from_secs(1)).await; test_suite.check_meta_subscriber().await; test_suite.validate_add_index_result(&index_key).await; let shard = test_suite.get_first_shard().await; let shard_id = shard.shard_id; let base_frag = make_frag_meta(&index_key, &mock_frag_path("AAA")); test_suite .must_append_base_frag_impl( &test_suite.get_one_worker_node_addr(), shard_id, &index_key, vec![base_frag.clone()], (0, 0), ) .await .unwrap(); test_suite .must_append_base_frag_impl( &test_suite.get_one_worker_node_addr(), shard_id, &index_key, vec![base_frag.clone()], (0, 1), ) .await .unwrap(); test_suite .must_append_delta_by_shard( &shard, "CCC", 15, AppendExtraParam { do_manifest_check: true, upload_to_s3: false, ..Default::default() }, ) .await; test_suite .must_append_delta_by_shard( &shard, "CCC", 15, AppendExtraParam { do_manifest_check: true, upload_to_s3: false, dedup: (1, 0, 0), ..Default::default() }, ) .await; let shard = test_suite.get_first_shard().await; assert_eq!(shard.manifest.len(), 2); test_suite.must_meta_subscriber_consistent().await; assert_eq!( test_suite .get_meta_service() .shard_task_manager .get_index(&index_key) .await .unwrap() .get_all_referenced_size() .await, 2 ); // Split with error index_id. It's OK because index_id is not encoded. let index_key_error = IndexKey::new_nullspace(1, 0); test_suite .must_manual_split_shard( &index_key, RangeBound::make_start(table_start_key(&index_key_error).to_vec()), vec![ HandleRange::new( table_start_key(&index_key_error).to_vec(), table_raw_key(&index_key_error, b"A"), ), HandleRange::new( table_raw_key(&index_key_error, b"A"), table_end_key(&index_key_error).to_vec(), ), ], 0, ) .await .unwrap(); assert_eq!( test_suite .get_meta_service() .shard_task_manager .get_index(&index_key) .await .unwrap() .get_all_referenced_size() .await, 2 ); test_suite.check_meta_subscriber().await; test_suite .must_create_index(&IndexKey::new_nullspace(3, 1001)) .await .unwrap(); test_suite.stop_worker().await.unwrap(); // Check the shards info with restarting meta service let shards_before = test_suite.collect_shards_from_restore().await; test_suite .restart_meta(make_default_test_meta_config(true)) .await .unwrap(); let shards_after = test_suite.collect_shards_from_restore().await; assert_eq!(shards_before, shards_after); test_suite.check_meta_subscriber().await; test_suite.stop_meta().await.unwrap(); }); Ok(()) } #[test] fn test_multiple_indexes() -> Result<()> { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; fail::cfg("worker_service_mock_meta_service_url", "return").unwrap(); defer!({ fail::remove("worker_service_mock_meta_service_url"); }); rt.block_on(async { let mut test_suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(true), ) .await .unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 0 ); let index_key1 = IndexKey::new_nullspace(1, 1001); test_suite.must_create_index(&index_key1).await.unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 1 ); let index_key2 = IndexKey::new_nullspace(1, 1002); test_suite.must_create_index(&index_key2).await.unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 2 ); let index_key3 = IndexKey::new_nullspace(2, 1001); test_suite.must_create_index(&index_key3).await.unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 3 ); let index = test_suite.get_index(&index_key1).await; let shard = index.select_first_shard().await.unwrap(); let shard_unwrap = shard.read().await.clone(); // fragment 0 of index i test_suite .must_append_delta_by_shard( &shard_unwrap, &format!("A{}_{}", 1001, 0), 1, AppendExtraParam::default(), ) .await; let shards = test_suite.collect_shards_from_restore().await; assert_eq!(shards.len(), 3); let mut req = tici_grpc::DebugGetShardInfoRequest::default(); for shard in shards.iter() { req.shard_ids.push(shard.shard_id); } let resp = test_suite .get_meta_internal_client() .debug_get_shard_info(req) .await .unwrap() .into_inner(); for shard in shards.iter() { let reported: Vec<_> = resp .shards .iter() .filter(|x| x.shard_info.as_ref().unwrap().shard_id == shard.shard_id) .collect(); assert_eq!(reported.len(), 1); assert_eq!(reported[0].frag_metas, shard.manifest.cloned_fragments()); assert_eq!( reported[0].shard_info.as_ref().unwrap(), &shard.get_header() ); assert_eq!(reported[0].s3_bucket, shard.frag_upload_location.bucket); } test_suite.must_drop_index(&index_key1).await.unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 2 ); assert_eq!( SHARD_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 2 ); test_suite.must_drop_index(&index_key2).await.unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 1 ); assert_eq!( SHARD_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 1 ); test_suite.must_drop_index(&index_key3).await.unwrap(); assert_eq!( INDEX_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 0 ); assert_eq!( SHARD_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), 0 ); test_suite.stop_worker().await.unwrap(); assert_eq!(test_suite.collect_shards_from_restore().await.len(), 0); assert_eq!(test_suite.collect_indexes_from_restore().await.len(), 0); test_suite.stop_meta().await.unwrap(); }); Ok(()) } #[test] fn test_mixed_workloads() -> Result<()> { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; fail::cfg("worker_service_mock_meta_service_url", "return").unwrap(); defer!({ fail::remove("worker_service_mock_meta_service_url"); }); let mut cfg = make_default_test_meta_config(true); cfg.shard.auto_compaction = false; cfg.shard.enable_split = false; cfg.shard.enable_patrol(); rt.block_on(async { let test_suite = Arc::new(RwLock::new( E2ETestSuite::new_with_mock_cfg(cfg, make_default_test_worker_config(true)) .await .unwrap(), )); let boot = async move |suite: Arc>, index_id: i64| { let index_key = IndexKey::new_nullspace(1, index_id); suite .write() .await .must_create_index(&IndexKey::new_nullspace(1, index_id)) .await .unwrap(); tokio::time::sleep(std::time::Duration::from_millis(100)).await; let index = suite.read().await.get_index(&index_key).await; let shard = index.select_first_shard().await.unwrap(); let shard_unwrap = shard.read().await.clone(); let shard_id = shard_unwrap.shard_id; // fragment 0 of index i suite .read() .await .must_append_delta_by_shard( &shard_unwrap, &format!("A{}_{}", index_id, 0), 1, AppendExtraParam { do_manifest_check: true, upload_to_s3: false, ..Default::default() }, ) .await; shard_id }; let shard_id = boot(test_suite.clone(), 1001).await; info!("Booted index {}", 1001); let suite = test_suite.clone(); let handle1 = tokio::spawn(async move { for i in 1002..1100 { boot(suite.clone(), i).await; info!("Booted index {}", i); } }); let suite = test_suite.clone(); let handle2 = tokio::spawn(async move { for i in 0..100 { let index_key = IndexKey::new_nullspace(1, 1001); tokio::time::sleep(std::time::Duration::from_millis(100)).await; let input = SplittableFrag::new( make_frag_meta(&index_key, &mock_frag_path(format!("A{}_{}", 1001, i))), shard_id, 0, ); let output = make_frag_meta(&index_key, &mock_frag_path(format!("A{}_{}", 1001, i + 1))); let state = suite .read() .await .do_compact_scheduled(&index_key, shard_id, vec![input], false) .await .unwrap(); suite .read() .await .do_compact_finished(shard_id, state.unwrap(), vec![output], true) .await .unwrap(); info!("Compacted index {}", i); } }); handle1.await.unwrap(); handle2.await.unwrap(); test_suite.write().await.stop_worker().await.unwrap(); assert_eq!( SHARD_COUNT_GAUGE_VEC .get_metric_with_label_values(&["total"]) .iter() .next() .unwrap() .get(), test_suite .read() .await .collect_shards_from_restore() .await .len() as i64 ); test_suite.read().await.validate_size_histo().await; test_suite.write().await.stop_meta().await.unwrap(); }); Ok(()) } #[test] fn test_create_shard_f_f() -> Result<()> { test_create_shard(false, false).unwrap(); Ok(()) } #[test] fn test_create_shard_t_t() -> Result<()> { test_create_shard(true, true).unwrap(); Ok(()) } #[test] fn test_fts_index_find_overlap() -> Result<()> { fn make_key_range(index_key: &IndexKey, left: &[u8], right: &[u8]) -> tici_grpc::KeyRange { tici_grpc::KeyRange { start_key: table_raw_key(index_key, left), end_key: table_raw_key(index_key, right), } } // helper to make handle range from index key and left/right suffix fn make_handle_range(index_key: &IndexKey, left: &[u8], right: &[u8]) -> HandleRange { HandleRange::new( table_raw_key(index_key, left), table_raw_key(index_key, right), ) } // helper to assert shard fragments and base slice counts async fn assert_shard_range( shards: &[Arc>], expected_ranges: &[HandleRange], stage: &str, ) { info!("Check shards {} ...", stage); assert_eq!(shards.len(), expected_ranges.len()); for (i, shard_arc) in shards.iter().enumerate() { let shard = shard_arc.read().await.clone(); assert_eq!(shard.range, expected_ranges[i]); } } let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1001); suite.must_create_index(&index_key).await.unwrap(); let index = suite.get_index(&index_key).await; assert_eq!(index.shards_len().await, 1); // clean the bootstrapped shard index.remove_all_shards().await.unwrap(); assert_eq!(index.shards_len().await, 0); // mock some in-memory shards for testing { let range = make_handle_range(&index_key, b"a", b"d"); let shard = Arc::new(RwLock::new(ShardInfo::new_for_test(range.clone()))); index .add_shard_for_test(range.sortable_start(), shard.clone()) .await; } { let range = make_handle_range(&index_key, b"d", b"h"); let shard = Arc::new(RwLock::new(ShardInfo::new_for_test(range.clone()))); index .add_shard_for_test(range.sortable_start(), shard.clone()) .await; } { let range = make_handle_range(&index_key, b"h", b"m"); let shard = Arc::new(RwLock::new(ShardInfo::new_for_test(range.clone()))); index .add_shard_for_test(range.sortable_start(), shard.clone()) .await; } assert_eq!(index.shards_len().await, 3); // find_overlapping_shards should return all 3 shards when range is empty let shards = index.find_overlapping_shards(&[], false).await; assert_shard_range( &shards, &[ make_handle_range(&index_key, b"a", b"d"), make_handle_range(&index_key, b"d", b"h"), make_handle_range(&index_key, b"h", b"m"), ], "range is empty, include_end_key=false", ) .await; let shards = index.find_overlapping_shards(&[], true).await; assert_shard_range( &shards, &[ make_handle_range(&index_key, b"a", b"d"), make_handle_range(&index_key, b"d", b"h"), make_handle_range(&index_key, b"h", b"m"), ], "range is empty, include_end_key=true", ) .await; // respect end_key inclusion let shards = index .find_overlapping_shards(&[make_key_range(&index_key, b"a", b"d")], false) .await; assert_shard_range( &shards, &[make_handle_range(&index_key, b"a", b"d")], "[a,d) include_end_key=false", ) .await; let shards = index .find_overlapping_shards(&[make_key_range(&index_key, b"a", b"d")], true) .await; assert_shard_range( &shards, &[ make_handle_range(&index_key, b"a", b"d"), make_handle_range(&index_key, b"d", b"h"), ], "[a,d] include_end_key=true", ) .await; }); Ok(()) } #[test] fn test_append_to_base_overlap() -> Result<()> { // helper to make handle range from index key and left/right suffix fn make_handle_range(index_key: &IndexKey, left: &[u8], right: &[u8]) -> HandleRange { HandleRange::new( table_raw_key(index_key, left), table_raw_key(index_key, right), ) } // helper to assert shard fragments and base slice counts fn assert_shard_frag( shards: &[ShardInfo], expected_ranges: &[HandleRange], expected_counts: &[[usize; 2]; 2], stage: &str, ) { info!("Check shards {} ...", stage); assert_eq!(shards.len(), 2); let left_shard = &shards[0]; assert_eq!(left_shard.range, expected_ranges[0]); assert_eq!( left_shard.manifest.get_fragments().len(), expected_counts[0][0] ); assert_eq!( left_shard.manifest.base_slice().len(), expected_counts[0][1] ); let right_shard = &shards[1]; assert_eq!(right_shard.range, expected_ranges[1]); assert_eq!( right_shard.manifest.get_fragments().len(), expected_counts[1][0] ); assert_eq!( right_shard.manifest.base_slice().len(), expected_counts[1][1] ); } fn assert_shard_base_paths( shards: &[ShardInfo], expected_ranges: &[HandleRange], expected_paths: &[Vec], stage: &str, ) { info!("Check shards {} base fragments ...", stage); assert_eq!(shards.len(), expected_ranges.len()); assert_eq!(shards.len(), expected_paths.len()); for (i, shard) in shards.iter().enumerate() { assert_eq!(shard.range, expected_ranges[i]); let mut actual = shard .manifest .base_slice() .iter() .map(|frag| frag.frag_path.clone()) .collect::>(); actual.sort(); let mut expected = expected_paths[i].clone(); expected.sort(); assert_eq!(actual, expected); } } let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1001); suite.must_create_index(&index_key).await.unwrap(); { // Split the init shard by "D" info!("Split shard at D..."); let index = suite.get_index(&index_key).await; let shard_arc = index.select_first_shard().await.unwrap(); let shard = shard_arc.read().await.clone(); suite .must_split_shard_simple(&index_key, shard.range.clone(), Some(b"D")) .await .unwrap(); } let tbl_start_key = table_start_key(&index_key); let tbl_end_key = table_end_key(&index_key); let expected_split_key = table_raw_key(&index_key, b"D"); let expected_shard_ranges = vec![ HandleRange::new(tbl_start_key.clone(), expected_split_key.clone()), HandleRange::new(expected_split_key.clone(), tbl_end_key.clone()), ]; // Ensure the splitted shards are empty and splitted at "D" assert_shard_frag( &suite.collect_shards_from_index(&index_key).await, &expected_shard_ranges, &[[0, 0], [0, 0]], "after split at D", ); { // Append base fragments that cover [A, D] info!("Append base fragment covering [A, D]..."); let addr = suite.get_one_worker_node_addr(); let frag_metas = vec![make_frag_meta_by_range( &mock_frag_path("frag_1"), &make_handle_range(&index_key, b"A", b"D"), )]; suite .do_append_base_frag(&addr, &index_key, frag_metas) .await .unwrap(); } // Should append the fragment to left shard and right shard separately assert_shard_frag( &suite.collect_shards_from_index(&index_key).await, &expected_shard_ranges, &[[1, 1], [1, 1]], "after append [A, D]", ); { // Append base fragments that cover [A, B] info!("Append base fragment covering [A, B]..."); let addr = suite.get_one_worker_node_addr(); let frag_metas = vec![make_frag_meta_by_range( &mock_frag_path("frag_2"), &make_handle_range(&index_key, b"A", b"B"), )]; suite .do_append_base_frag(&addr, &index_key, frag_metas) .await .unwrap(); } // Should only append the fragment [A, B] to the left shard assert_shard_frag( &suite.collect_shards_from_index(&index_key).await, &expected_shard_ranges, &[[2, 2], [1, 1]], "after append [A, B]", ); { // Append base fragments that cover [D, E] info!("Append base fragment covering [D, E]..."); let addr = suite.get_one_worker_node_addr(); let frag_metas = vec![make_frag_meta_by_range( &mock_frag_path("frag_3"), &make_handle_range(&index_key, b"D", b"E"), )]; suite .do_append_base_frag(&addr, &index_key, frag_metas) .await .unwrap(); } // Should only append the fragment [D, E] to the right shard assert_shard_frag( &suite.collect_shards_from_index(&index_key).await, &expected_shard_ranges, &[[2, 2], [2, 2]], "after append [D, E]", ); { // Append base fragments with different ranges in a single request info!("Append base fragments covering [A, B], [A, D], [D, E]..."); let addr = suite.get_one_worker_node_addr(); let frag_metas = vec![ make_frag_meta_by_range( &mock_frag_path("frag_4"), &make_handle_range(&index_key, b"A", b"B"), ), make_frag_meta_by_range( &mock_frag_path("frag_5"), &make_handle_range(&index_key, b"A", b"D"), ), make_frag_meta_by_range( &mock_frag_path("frag_6"), &make_handle_range(&index_key, b"D", b"E"), ), ]; suite .do_append_base_frag(&addr, &index_key, frag_metas) .await .unwrap(); } // Should append overlapping fragments to both shards and exclusive ones only to // their shard let shards = suite.collect_shards_from_index(&index_key).await; assert_shard_frag( &shards, &expected_shard_ranges, &[[4, 4], [4, 4]], "after append mixed ranges", ); assert_shard_base_paths( &shards, &expected_shard_ranges, &[ mock_frag_paths(&["frag_1", "frag_2", "frag_4", "frag_5"]), mock_frag_paths(&["frag_1", "frag_3", "frag_5", "frag_6"]), ], "after append mixed ranges", ); }); Ok(()) } #[test] fn test_append_base_split_race() -> Result<()> { fn make_handle_range(index_key: &IndexKey, left: &[u8], right: &[u8]) -> HandleRange { HandleRange::new( table_raw_key(index_key, left), table_raw_key(index_key, right), ) } let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1001); suite.must_create_index(&index_key).await.unwrap(); let index = suite.get_index(&index_key).await; let shard_arc = index.select_first_shard().await.unwrap(); let shard = shard_arc.read().await.clone(); fail::cfg("append_base_frag_meta_before_get_shard", "pause").unwrap(); defer!({ fail::remove("append_base_frag_meta_before_get_shard"); }); let mut client = suite.get_meta_internal_client(); let frag_metas = vec![make_frag_meta_by_range( &mock_frag_path("race_frag"), &make_handle_range(&index_key, b"A", b"D"), )]; let append_req = AppendBaseFragMetaRequest { keyspace_id: index_key.keyspace_id, table_id: index_key.table_id, index_id: index_key.index_id, frag_metas, }; let append_handle = tokio::spawn(async move { client.append_base_frag_meta(append_req).await }); tokio::time::sleep(Duration::from_millis(200)).await; assert!( !append_handle.is_finished(), "append request finished before split" ); suite .must_split_shard_simple(&index_key, shard.range.clone(), Some(b"D")) .await .unwrap(); fail::remove("append_base_frag_meta_before_get_shard"); let resp = append_handle.await.unwrap().unwrap().into_inner(); assert_eq!(resp.status, ErrorCode::UnknownError as i32); let shards = suite.collect_shards_from_index(&index_key).await; assert_eq!(shards.len(), 2); for shard in shards { assert_eq!(shard.manifest.base_slice().len(), 0); } }); Ok(()) } #[test] fn test_append_delta_merge_race() -> Result<()> { fn make_handle_range(index_key: &IndexKey, left: &[u8], right: &[u8]) -> HandleRange { HandleRange::new( table_raw_key(index_key, left), table_raw_key(index_key, right), ) } let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1002); suite.must_create_index(&index_key).await.unwrap(); let index = suite.get_index(&index_key).await; let shard_arc = index.select_first_shard().await.unwrap(); let shard = shard_arc.read().await.clone(); suite .must_split_shard_simple(&index_key, shard.range.clone(), Some(b"D")) .await .unwrap(); let shards = index.clone_shards().await; assert_eq!(shards.len(), 2); let first = shards[0].read().await.clone(); let second = shards[1].read().await.clone(); let (left, _right) = if first.range.sortable_start() <= second.range.sortable_start() { (first, second) } else { (second, first) }; let (target, frag_range) = (left, (b"A".as_slice(), b"C".as_slice())); fail::cfg("append_delta_frag_meta_before_get_shard", "pause").unwrap(); defer!({ fail::remove("append_delta_frag_meta_before_get_shard"); }); let mut client = suite.get_meta_internal_client(); let append_req = tici_grpc::AppendDeltaFragMetaRequest { keyspace_id: index_key.keyspace_id, table_id: index_key.table_id, index_id: index_key.index_id, shard_id: target.shard_id, worker_node_addr: target.writer.clone(), frag_metas: vec![make_frag_meta_by_range( &mock_frag_path("race_delta"), &make_handle_range(&index_key, frag_range.0, frag_range.1), )], progress: Some(tici_grpc::ShardProgress { last_cdc_file: create_cdc_file_name(1), }), }; let append_handle = tokio::spawn(async move { client.append_delta_frag_meta(append_req).await }); tokio::time::sleep(Duration::from_millis(200)).await; assert!( !append_handle.is_finished(), "append delta request finished before merge" ); suite .must_merge_shard(shards.as_slice(), String::new()) .await; fail::remove("append_delta_frag_meta_before_get_shard"); let resp = append_handle.await.unwrap().unwrap().into_inner(); assert_eq!(resp.status, ErrorCode::ShardNotFound as i32); let merged = index.select_first_shard().await.unwrap(); let merged = merged.read().await; assert_eq!(index.shards_len().await, 1); assert_eq!(merged.manifest.get_fragments().len(), 0); }); Ok(()) } #[test] fn test_finished_append_base_rejects_append_and_survives_restart() -> Result<()> { fn make_handle_range(index_key: &IndexKey, left: &[u8], right: &[u8]) -> HandleRange { HandleRange::new( table_raw_key(index_key, left), table_raw_key(index_key, right), ) } let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let mut suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1002); suite.must_create_index(&index_key).await.unwrap(); let index = suite.get_index(&index_key).await; assert!(!index.is_finished_append_base()); assert!(index.finish_append_base().await.unwrap()); assert!(!index.finish_append_base().await.unwrap()); let restored_indexes = suite.collect_indexes_from_restore().await; assert!( restored_indexes .iter() .any(|(index_info, _, _, _, properties)| { index_info.keyspace_id == index_key.keyspace_id && index_info.table_id == index_key.table_id && index_info.index_id == index_key.index_id && properties.finished_append_base }) ); suite .restart_meta(make_default_test_meta_config(false)) .await .unwrap(); let index = suite.get_index(&index_key).await; assert!(index.is_finished_append_base()); let mut client = suite.get_meta_internal_client(); let append_req = AppendBaseFragMetaRequest { keyspace_id: index_key.keyspace_id, table_id: index_key.table_id, index_id: index_key.index_id, frag_metas: vec![make_frag_meta_by_range( &mock_frag_path("finished_base_frag"), &make_handle_range(&index_key, b"A", b"D"), )], }; let resp = client .append_base_frag_meta(append_req) .await .unwrap() .into_inner(); assert_eq!(resp.status, ErrorCode::InvalidOp as i32); }); Ok(()) } #[test] fn test_finish_append_base_switches_base_to_delta_and_survives_restart() -> Result<()> { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let mut suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1003); suite.must_create_index(&index_key).await.unwrap(); let shard = suite.get_first_shard().await; let shard_id = shard.shard_id; let worker_addr = suite.get_one_worker_node_addr(); suite .must_append_base_frag( &worker_addr, shard_id, &index_key, vec![make_frag_meta(&index_key, &mock_frag_path("base_0"))], ) .await .unwrap(); suite .must_append_base_frag( &worker_addr, shard_id, &index_key, vec![make_frag_meta(&index_key, &mock_frag_path("base_1"))], ) .await .unwrap(); suite .must_append_delta_by_shard(&shard, "delta_0", 1, AppendExtraParam::default()) .await; let before_switch = suite .get_shard_from_index(&index_key, shard_id) .await .unwrap() .read() .await .manifest .get_fragments() .iter() .map(|frag| frag.frag_path.clone()) .collect::>(); let index = suite.get_index(&index_key).await; assert!(index.finish_append_base().await.unwrap()); assert!(!index.finish_append_base().await.unwrap()); assert!(index.is_finished_append_base()); let shard_after_switch = suite .get_shard_from_index(&index_key, shard_id) .await .unwrap(); let shard_after_switch = shard_after_switch.read().await; assert!(shard_after_switch.manifest.base_slice().is_empty()); assert_eq!( shard_after_switch.manifest.delta_slice().len(), before_switch.len() ); assert_eq!( shard_after_switch .manifest .get_fragments() .iter() .map(|frag| frag.frag_path.clone()) .collect::>(), before_switch, ); drop(shard_after_switch); let restored_indexes = suite.collect_indexes_from_restore().await; assert!( restored_indexes .iter() .any(|(index_info, _, _, _, properties)| { index_info.keyspace_id == index_key.keyspace_id && index_info.table_id == index_key.table_id && index_info.index_id == index_key.index_id && properties.finished_append_base }) ); let restored_shard = suite.get_shard_from_restore(shard_id).await.unwrap(); assert!(restored_shard.manifest.base_slice().is_empty()); assert_eq!( restored_shard .manifest .get_fragments() .iter() .map(|frag| frag.frag_path.clone()) .collect::>(), before_switch, ); suite .restart_meta(make_default_test_meta_config(false)) .await .unwrap(); let index = suite.get_index(&index_key).await; assert!(index.is_finished_append_base()); let restored_shard = suite.get_shard_from_restore(shard_id).await.unwrap(); assert!(restored_shard.manifest.base_slice().is_empty()); assert_eq!( restored_shard .manifest .get_fragments() .iter() .map(|frag| frag.frag_path.clone()) .collect::>(), before_switch, ); }); Ok(()) } #[test] fn test_finish_append_base_allows_ongoing_compaction() -> Result<()> { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; fail::cfg("worker_service_mock_meta_service_url", "return").unwrap(); fail::cfg("mock_etcd_leader_term", "return").unwrap(); defer!({ fail::remove("worker_service_mock_meta_service_url"); fail::remove("mock_etcd_leader_term"); }); rt.block_on(async { let suite = E2ETestSuite::new_with_mock_cfg( make_default_test_meta_config(false), make_default_test_worker_config(false), ) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1004); suite.must_create_index(&index_key).await.unwrap(); let shard = suite.get_first_shard().await; let shard_id = shard.shard_id; let worker_addr = suite.get_one_worker_node_addr(); suite .must_append_base_frag( &worker_addr, shard_id, &index_key, vec![make_frag_meta(&index_key, &mock_frag_path("base_0"))], ) .await .unwrap(); suite .must_append_delta_by_shard(&shard, "delta_0", 1, AppendExtraParam::default()) .await; let input = suite .get_shard_from_index(&index_key, shard_id) .await .unwrap() .read() .await .manifest .delta_slice() .to_vec(); let compaction_state = suite .do_compact_scheduled(&index_key, shard_id, input, true) .await .unwrap() .unwrap(); let index = suite.get_index(&index_key).await; assert!(index.finish_append_base().await.unwrap()); let shard = suite .get_shard_from_index(&index_key, shard_id) .await .unwrap(); assert_eq!(shard.read().await.progress.compaction_state.len(), 1); assert!(shard.read().await.manifest.base_slice().is_empty()); suite .do_compact_finished( shard_id, compaction_state, vec![make_frag_meta( &index_key, &mock_frag_path("compacted_after_switch"), )], true, ) .await .unwrap(); let restored = suite.get_shard_from_restore(shard_id).await.unwrap(); assert!(restored.manifest.base_slice().is_empty()); assert_eq!(restored.manifest.delta_slice().len(), 2); assert!( restored .manifest .delta_slice() .iter() .any(|frag| frag.frag_path == mock_frag_path("compacted_after_switch")) ); }); Ok(()) } #[test] fn test_handle_compaction_switches_base_to_delta_when_finished_append_base_leftovers_exist() -> Result<()> { let rt = tokio::runtime::Builder::new_multi_thread() .worker_threads(4) .enable_all() .build()?; rt.block_on(async { let mut meta_cfg = make_default_test_meta_config(false); meta_cfg.shard.uni_compaction.min_compaction_frags = 100; let mut suite = E2ETestSuite::new_with_mock_cfg(meta_cfg, make_default_test_worker_config(false)) .await .unwrap(); let index_key = IndexKey::new_nullspace(1, 1005); suite.must_create_index(&index_key).await.unwrap(); let shard = suite.get_first_shard().await; let shard_id = shard.shard_id; let worker_addr = suite.get_one_worker_node_addr(); suite .must_append_base_frag( &worker_addr, shard_id, &index_key, vec![make_frag_meta(&index_key, &mock_frag_path("base_0"))], ) .await .unwrap(); suite .must_append_base_frag( &worker_addr, shard_id, &index_key, vec![make_frag_meta(&index_key, &mock_frag_path("base_1"))], ) .await .unwrap(); suite .must_append_delta_by_shard(&shard, "delta_0", 1, AppendExtraParam::default()) .await; suite .must_append_delta_by_shard(&shard, "delta_1", 2, AppendExtraParam::default()) .await; suite .must_append_delta_by_shard(&shard, "delta_2", 3, AppendExtraParam::default()) .await; let index = suite.get_index(&index_key).await; assert!(index.finish_append_base().await.unwrap()); let shard_ref = suite .get_shard_from_index(&index_key, shard_id) .await .unwrap(); // We forcibly add some base fragments, which is not permitted in formal logic. { let mut shard_guard = shard_ref.write().await; let frags = shard_guard.manifest.cloned_fragments(); let mut manifest = ShardManifest::new_for_first(); manifest.append_base_fragments( frags[..2].to_vec(), shard_guard.shard_id, shard_guard.seq, ); manifest.append_delta_fragments( frags[2..].to_vec(), shard_guard.shard_id, shard_guard.seq, ); shard_guard.manifest = manifest; } // handle_compaction may first switch leftover base fragments into delta when // finished_append_base is already set, so the helper's "append adds exactly one // delta fragment" count check no longer holds for this append. suite .must_append_delta_by_shard( &shard, "delta_3", 4, AppendExtraParam { do_manifest_check: false, ..Default::default() }, ) .await; let shard_after = suite .get_shard_from_index(&index_key, shard_id) .await .unwrap(); let shard_after = shard_after.read().await; assert!(shard_after.manifest.base_slice().is_empty()); assert_eq!( shard_after .manifest .get_fragments() .iter() .map(|frag| frag.frag_path.clone()) .collect::>(), mock_frag_paths(&[ "base_1", "base_0", "delta_0", "delta_1", "delta_2", "delta_3" ]) ); suite.stop().await.unwrap(); }); Ok(()) }