@@ -176,6 +176,28 @@ public void failTest(
176176 killNode );
177177 }
178178
179+ public void failAndRollbackTest (
180+ final int dataReplicateFactor ,
181+ final int schemaReplicationFactor ,
182+ final int configNodeNum ,
183+ final int dataNodeNum ,
184+ KeySetView <String , Boolean > killConfigNodeKeywords ,
185+ KeySetView <String , Boolean > killDataNodeKeywords ,
186+ KillNode killNode )
187+ throws Exception {
188+ generalTestWithAllOptions (
189+ dataReplicateFactor ,
190+ schemaReplicationFactor ,
191+ configNodeNum ,
192+ dataNodeNum ,
193+ killConfigNodeKeywords ,
194+ killDataNodeKeywords ,
195+ actionOfKillNode ,
196+ false ,
197+ killNode ,
198+ true );
199+ }
200+
179201 public void killClusterTest (
180202 KeySetView <String , Boolean > configNodeKeywords , boolean expectMigrateSuccess )
181203 throws Exception {
@@ -204,6 +226,31 @@ public void generalTestWithAllOptions(
204226 final boolean expectMigrateSuccess ,
205227 KillNode killNode )
206228 throws Exception {
229+ generalTestWithAllOptions (
230+ dataReplicateFactor ,
231+ schemaReplicationFactor ,
232+ configNodeNum ,
233+ dataNodeNum ,
234+ configNodeKeywords ,
235+ dataNodeKeywords ,
236+ actionWhenDetectKeyWords ,
237+ expectMigrateSuccess ,
238+ killNode ,
239+ false );
240+ }
241+
242+ private void generalTestWithAllOptions (
243+ final int dataReplicateFactor ,
244+ final int schemaReplicationFactor ,
245+ final int configNodeNum ,
246+ final int dataNodeNum ,
247+ KeySetView <String , Boolean > configNodeKeywords ,
248+ KeySetView <String , Boolean > dataNodeKeywords ,
249+ Consumer <KillPointContext > actionWhenDetectKeyWords ,
250+ final boolean expectMigrateSuccess ,
251+ KillNode killNode ,
252+ boolean expectRollbackWhenFail )
253+ throws Exception {
207254 // prepare env
208255 EnvFactory .getEnv ()
209256 .getConfig ()
@@ -266,26 +313,42 @@ public void generalTestWithAllOptions(
266313 statement .execute (buildRegionMigrateCommand (selectedRegion , originalDataNode , destDataNode ));
267314
268315 boolean success = false ;
269- Predicate <TShowRegionResp > migrateRegionPredicate =
270- tShowRegionResp -> {
271- Map <Integer , Set <Integer >> newRegionMap =
272- getRegionMap (tShowRegionResp .getRegionInfoList ());
273- Set <Integer > dataNodes = newRegionMap .get (selectedRegion );
274- return !dataNodes .contains (originalDataNode ) && dataNodes .contains (destDataNode );
275- };
276- try {
277- awaitUntilSuccess (
278- client ,
279- selectedRegion ,
280- migrateRegionPredicate ,
281- Optional .of (destDataNode ),
282- Optional .of (originalDataNode ));
283- success = true ;
284- } catch (ConditionTimeoutException e ) {
285- if (expectMigrateSuccess ) {
286- LOGGER .error ("Region migrate failed" , e );
316+ if (expectRollbackWhenFail ) {
317+ awaitKillPointsTriggered (configNodeKeywords );
318+ awaitKillPointsTriggered (dataNodeKeywords );
319+ try {
320+ awaitUntilSuccess (
321+ client ,
322+ selectedRegion ,
323+ rollbackPredicate (selectedRegion , regionMap .get (selectedRegion ), destDataNode ),
324+ Optional .empty (),
325+ Optional .of (destDataNode ));
326+ } catch (ConditionTimeoutException e ) {
327+ LOGGER .error ("Region migrate did not roll back" , e );
287328 Assert .fail ();
288329 }
330+ } else {
331+ Predicate <TShowRegionResp > migrateRegionPredicate =
332+ tShowRegionResp -> {
333+ Map <Integer , Set <Integer >> newRegionMap =
334+ getRegionMap (tShowRegionResp .getRegionInfoList ());
335+ Set <Integer > dataNodes = newRegionMap .get (selectedRegion );
336+ return !dataNodes .contains (originalDataNode ) && dataNodes .contains (destDataNode );
337+ };
338+ try {
339+ awaitUntilSuccess (
340+ client ,
341+ selectedRegion ,
342+ migrateRegionPredicate ,
343+ Optional .of (destDataNode ),
344+ Optional .of (originalDataNode ));
345+ success = true ;
346+ } catch (ConditionTimeoutException e ) {
347+ if (expectMigrateSuccess ) {
348+ LOGGER .error ("Region migrate failed" , e );
349+ Assert .fail ();
350+ }
351+ }
289352 }
290353 if (!expectMigrateSuccess && success ) {
291354 LOGGER .error ("Region migrate succeeded unexpectedly" );
@@ -311,6 +374,28 @@ public void generalTestWithAllOptions(
311374 }
312375 }
313376
377+ private static Predicate <TShowRegionResp > rollbackPredicate (
378+ int selectedRegion , Set <Integer > originalDataNodes , int destDataNode ) {
379+ return tShowRegionResp -> {
380+ List <TRegionInfo > selectedDataRegionInfos =
381+ tShowRegionResp .getRegionInfoList ().stream ()
382+ .filter (
383+ regionInfo ->
384+ regionInfo .getConsensusGroupId ().getType () == TConsensusGroupType .DataRegion
385+ && regionInfo .getConsensusGroupId ().getId () == selectedRegion )
386+ .collect (Collectors .toList ());
387+ Set <Integer > dataNodes =
388+ selectedDataRegionInfos .stream ()
389+ .map (TRegionInfo ::getDataNodeId )
390+ .collect (Collectors .toSet ());
391+ return dataNodes .equals (originalDataNodes )
392+ && !dataNodes .contains (destDataNode )
393+ && selectedDataRegionInfos .stream ()
394+ .allMatch (
395+ regionInfo -> RegionStatus .Running .getStatus ().equals (regionInfo .getStatus ()));
396+ };
397+ }
398+
314399 public static Set <Integer > getAllDataNodes (Statement statement ) throws Exception {
315400 ResultSet result = statement .executeQuery (SHOW_DATANODES );
316401 Set <Integer > allDataNodeId = new HashSet <>();
@@ -428,6 +513,13 @@ void checkKillPointsAllTriggered(KeySetView<String, Boolean> killPoints) {
428513 }
429514 }
430515
516+ private static void awaitKillPointsTriggered (KeySetView <String , Boolean > killPoints ) {
517+ if (killPoints .isEmpty ()) {
518+ return ;
519+ }
520+ Awaitility .await ().atMost (2 , TimeUnit .MINUTES ).until (killPoints ::isEmpty );
521+ }
522+
431523 private static String buildRegionMigrateCommand (int who , int from , int to ) {
432524 String result = String .format (REGION_MIGRATE_COMMAND_FORMAT , who , from , to );
433525 LOGGER .info (result );
0 commit comments