br: wait tiflash replicas ready && fix unstable test

br/pkg/restore/client.go

@@ -1824,11 +1824,28 @@
 				zap.Stringer("table", tbl.OldTable.Info.Name),
 				zap.Stringer("db", tbl.OldTable.DB.Name))
 			for {
-				progress, err := infosync.CalculateTiFlashProgress(tbl.Table.ID, tbl.Table.TiFlashReplica.Count, tiFlashStores)
-				if err != nil {
-					log.Warn("failed to get tiflash replica progress, wait for next retry", zap.Error(err))
-					time.Sleep(time.Second)
-					continue
+				var progress float64
+				if pi := tbl.Table.GetPartitionInfo(); pi != nil && len(pi.Definitions) > 0 {
+					for _, p := range pi.Definitions {
+						progressOfPartition, err := infosync.MustGetTiFlashProgress(p.ID, tbl.Table.TiFlashReplica.Count, &tiFlashStores)
+						if err != nil {
+							log.Warn("failed to get progress for tiflash partition replica, retry it",
+								zap.Int64("tableID", tbl.Table.ID), zap.Int64("partitionID", p.ID), zap.Error(err))
+							time.Sleep(time.Second)
+							continue
+						}
+						progress += progressOfPartition
+					}
+					progress = progress / float64(len(pi.Definitions))
+				} else {
+					var err error
+					progress, err = infosync.MustGetTiFlashProgress(tbl.Table.ID, tbl.Table.TiFlashReplica.Count, &tiFlashStores)
+					if err != nil {
+						log.Warn("failed to get progress for tiflash replica, retry it",
+							zap.Int64("tableID", tbl.Table.ID), zap.Error(err))
+						time.Sleep(time.Second)
+						continue
+					}
 				}
 				// check until progress is 1
 				if progress == 1 {

br/tests/br_tiflash/run.sh

@@ -23,16 +23,29 @@ run_sql "CREATE DATABASE $DB"
 
 run_sql "CREATE TABLE $DB.kv(k varchar(256) primary key, v int)"
 
+run_sql "CREATE TABLE $DB.partition_kv(\
+ k INT, \
+ v INT, \
+ PRIMARY KEY(k) CLUSTERED \
+) PARTITION BY RANGE(k) (\
+  PARTITION p0 VALUES LESS THAN (200), \
+  PARTITION p1 VALUES LESS THAN (400), \
+  PARTITION p2 VALUES LESS THAN MAXVALUE)"
+
 stmt="INSERT INTO $DB.kv(k, v) VALUES ('1-record', 1)"
+parition_stmt="INSERT INTO $DB.partition_kv(k, v) VALUES (1, 1)"
 for i in $(seq 2 $RECORD_COUNT); do
     stmt="$stmt,('$i-record', $i)"
+    parition_stmt="$parition_stmt,($i, $i)"
 done
 run_sql "$stmt"
+run_sql "$parition_stmt"
 
 if ! run_sql "ALTER TABLE $DB.kv SET TIFLASH REPLICA 1"; then
   # 10s should be enough for tiflash-proxy get started
   sleep 10
   run_sql "ALTER TABLE $DB.kv SET TIFLASH REPLICA 1"
+  run_sql "ALTER TABLE $DB.partition_kv SET TIFLASH REPLICA 1"
 fi
 
 
@@ -54,6 +67,8 @@ run_sql "DROP DATABASE $DB"
 run_br restore full -s "local://$TEST_DIR/$DB" --pd $PD_ADDR --wait-tiflash-ready=true
 
 # check TiFlash sync
+echo "wait 3 seconds for tiflash tick puller triggered"
+sleep 3
 if ! [ $(run_sql "select * from information_schema.tiflash_replica" | grep "PROGRESS" | sed "s/[^0-9]//g") -eq 1 ]; then
     echo "restore didn't wait tiflash synced after set --wait-tiflash-ready=true."
     exit 1

br/tests/run.sh

@@ -21,6 +21,11 @@ export TEST_DIR=/tmp/backup_restore_test
 export COV_DIR="/tmp/group_cover"
 source $CUR/_utils/run_services
 
+# Create COV_DIR if not exists
+if [ -d "$COV_DIR" ]; then
+   mkdir -p $COV_DIR
+fi
+
 # Reset TEST_DIR
 rm -rf $TEST_DIR && mkdir -p $TEST_DIR