11mod ctl;
22
33use futures_timer:: Delay ;
4- use log:: info;
4+ use log:: { info, warn } ;
55use std:: { env, time:: Duration } ;
66use tikv_client:: { ColumnFamily , Key , RawClient , Result , TransactionClient } ;
77
88const ENV_PD_ADDRS : & str = "PD_ADDRS" ;
99const ENV_ENABLE_MULIT_REGION : & str = "MULTI_REGION" ;
10+ const REGION_SPLIT_TIME_LIMIT : Duration = Duration :: from_secs ( 15 ) ;
1011
1112// Delete all entries in TiKV to leave a clean space for following tests.
1213pub async fn clear_tikv ( ) {
@@ -24,6 +25,11 @@ pub async fn clear_tikv() {
2425// To test with multiple regions, prewrite some data. Tests that hope to test
2526// with multiple regions should use keys in the corresponding ranges.
2627pub async fn init ( ) -> Result < ( ) > {
28+ // ignore SetLoggerError
29+ let _ = simple_logger:: SimpleLogger :: new ( )
30+ . with_level ( log:: LevelFilter :: Warn )
31+ . init ( ) ;
32+
2733 if enable_multi_region ( ) {
2834 // 1000 keys: 0..1000
2935 let keys_1 = std:: iter:: successors ( Some ( 0u32 ) , |x| Some ( x + 1 ) )
@@ -36,14 +42,14 @@ pub async fn init() -> Result<()> {
3642 . take ( count as usize - 1 )
3743 . map ( |x| x. to_be_bytes ( ) . to_vec ( ) ) ;
3844
39- ensure_region_splitted ( keys_1. chain ( keys_2) , 100 ) . await ?;
45+ ensure_region_split ( keys_1. chain ( keys_2) , 100 ) . await ?;
4046 }
4147
4248 clear_tikv ( ) . await ;
4349 Ok ( ( ) )
4450}
4551
46- async fn ensure_region_splitted (
52+ async fn ensure_region_split (
4753 keys : impl IntoIterator < Item = impl Into < Key > > ,
4854 region_count : u32 ,
4955) -> Result < ( ) > {
@@ -52,7 +58,7 @@ async fn ensure_region_splitted(
5258 }
5359
5460 // 1. write plenty transactional keys
55- // 2. wait until regions splitted
61+ // 2. wait until regions split
5662
5763 let client = TransactionClient :: new ( pd_addrs ( ) ) . await ?;
5864 let mut txn = client. begin_optimistic ( ) . await ?;
@@ -65,11 +71,16 @@ async fn ensure_region_splitted(
6571 txn. commit ( ) . await ?;
6672
6773 info ! ( "splitting regions..." ) ;
74+ let start_time = std:: time:: Instant :: now ( ) ;
6875 loop {
6976 if ctl:: get_region_count ( ) . await ? as u32 >= region_count {
7077 break ;
7178 }
72- Delay :: new ( Duration :: from_secs ( 1 ) ) . await ;
79+ if start_time. elapsed ( ) > REGION_SPLIT_TIME_LIMIT {
80+ warn ! ( "Stop splitting regions: time limit exceeded" ) ;
81+ break ;
82+ }
83+ Delay :: new ( Duration :: from_millis ( 200 ) ) . await ;
7384 }
7485
7586 Ok ( ( ) )
0 commit comments