HBASE-18068 Fix flaky test TestAsyncSnapshotAdminApi

- internalRestoreSnapshot() returns future which completes by just getting proc_id from master. Changed it to wait for the procedure to complete.
- Refactor TestAsyncSnapshotAdminApi: Add cleanup() which deletes all tables and snapshots after every test run. Simplifies individual tests.

Change-Id: Idc30fb699db32d58fd0f60da220953a430f1d3cc

hbase-client/src/main/java/org/apache/hadoop/hbase/client/AsyncHBaseAdmin.java

@@ -90,7 +90,6 @@ import org.apache.hadoop.hbase.shaded.protobuf.generated.AdminProtos.GetOnlineRe
 import org.apache.hadoop.hbase.shaded.protobuf.generated.AdminProtos.GetOnlineRegionResponse;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.AdminProtos.SplitRegionRequest;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.AdminProtos.SplitRegionResponse;
-import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.ProcedureDescription;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.TableSchema;
 import org.apache.hadoop.hbase.shaded.protobuf.generated.MasterProtos.AbortProcedureRequest;
@@ -1528,7 +1527,8 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
     return future;
   }
 
-  public CompletableFuture<Void> addReplicationPeer(String peerId, ReplicationPeerConfig peerConfig) {
+  public CompletableFuture<Void> addReplicationPeer(String peerId,
+      ReplicationPeerConfig peerConfig) {
     return this
         .<Void> newMasterCaller()
         .action(
@@ -1798,11 +1798,11 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
           future.completeExceptionally(err);
         } else {
           // Step.2 Restore snapshot
-          internalRestoreSnapshot(snapshotName, tableName).whenComplete((ret2, err2) -> {
+          internalRestoreSnapshot(snapshotName, tableName).whenComplete((void2, err2) -> {
             if (err2 != null) {
               // Step.3.a Something went wrong during the restore and try to rollback.
               internalRestoreSnapshot(failSafeSnapshotSnapshotName, tableName)
-                  .whenComplete((ret3, err3) -> {
+                  .whenComplete((void3, err3) -> {
                     if (err3 != null) {
                       future.completeExceptionally(err3);
                     } else {
@@ -1864,13 +1864,8 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
           future.completeExceptionally(err2);
         } else if (!exists) {
           // if table does not exist, then just clone snapshot into new table.
-          internalRestoreSnapshot(snapshotName, finalTableName).whenComplete((ret, err3) -> {
+          completeConditionalOnFuture(future,
-            if (err3 != null) {
+              internalRestoreSnapshot(snapshotName, finalTableName));
-              future.completeExceptionally(err3);
-            } else {
-              future.complete(ret);
-            }
-          });
         } else {
           isTableDisabled(finalTableName).whenComplete((disabled, err4) -> {
             if (err4 != null) {
@@ -1878,14 +1873,8 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
             } else if (!disabled) {
               future.completeExceptionally(new TableNotDisabledException(finalTableName));
             } else {
-              restoreSnapshotWithFailSafe(snapshotName, finalTableName, takeFailSafeSnapshot)
+              completeConditionalOnFuture(future,
-                  .whenComplete((ret, err5) -> {
+                  restoreSnapshotWithFailSafe(snapshotName, finalTableName, takeFailSafeSnapshot));
-                    if (err5 != null) {
-                      future.completeExceptionally(err5);
-                    } else {
-                      future.complete(ret);
-                    }
-                  });
             }
           });
         }
@@ -1894,6 +1883,17 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
     return future;
   }
 
+  private <T> void completeConditionalOnFuture(CompletableFuture<T> dependentFuture,
+      CompletableFuture<T> parentFuture) {
+    parentFuture.whenComplete((res, err) -> {
+      if (err != null) {
+        dependentFuture.completeExceptionally(err);
+      } else {
+        dependentFuture.complete(res);
+      }
+    });
+  }
+
   @Override
   public CompletableFuture<Void> cloneSnapshot(String snapshotName, TableName tableName) {
     CompletableFuture<Void> future = new CompletableFuture<>();
@@ -1903,13 +1903,7 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
       } else if (exists) {
         future.completeExceptionally(new TableExistsException(tableName));
       } else {
-        internalRestoreSnapshot(snapshotName, tableName).whenComplete((ret, err2) -> {
+        completeConditionalOnFuture(future, internalRestoreSnapshot(snapshotName, tableName));
-          if (err2 != null) {
-            future.completeExceptionally(err2);
-          } else {
-            future.complete(ret);
-          }
-        });
       }
     });
     return future;
@@ -1924,13 +1918,15 @@ public class AsyncHBaseAdmin implements AsyncAdmin {
     } catch (IllegalArgumentException e) {
       return failedFuture(e);
     }
-    return this.<Void> newMasterCaller()
+    return waitProcedureResult(
+        this.<Long> newMasterCaller()
             .action((controller, stub) -> this
-            .<RestoreSnapshotRequest, RestoreSnapshotResponse, Void> call(controller, stub,
+                .<RestoreSnapshotRequest, RestoreSnapshotResponse, Long> call(controller, stub,
                     RestoreSnapshotRequest.newBuilder().setSnapshot(snapshot)
                         .setNonceGroup(ng.getNonceGroup()).setNonce(ng.newNonce()).build(),
-              (s, c, req, done) -> s.restoreSnapshot(c, req, done), resp -> null))
+                    (s, c, req, done) -> s.restoreSnapshot(c, req, done),
-        .call();
+                    (resp) -> resp.getProcId()))
+            .call());
   }
 
   @Override

hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAsyncSnapshotAdminApi.java

@@ -22,9 +22,13 @@ import org.apache.hadoop.hbase.TableName;
 import org.apache.hadoop.hbase.testclassification.ClientTests;
 import org.apache.hadoop.hbase.testclassification.MediumTests;
 import org.apache.hadoop.hbase.util.Bytes;
+import org.junit.After;
 import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.rules.TestName;
 
 import java.io.IOException;
 import java.util.Collections;
@@ -34,14 +38,30 @@ import java.util.regex.Pattern;
 @Category({ MediumTests.class, ClientTests.class })
 public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
 
-  @Test
-  public void testTakeSnapshot() throws Exception {
   String snapshotName1 = "snapshotName1";
   String snapshotName2 = "snapshotName2";
-    TableName tableName = TableName.valueOf("testTakeSnapshot");
+  String snapshotName3 = "snapshotName3";
+
+  @Rule
+  public TestName testName = new TestName();
+  TableName tableName;
+
+  @Before
+  public void setup() {
+    tableName = TableName.valueOf(testName.getMethodName());
+  }
+
+  @After
+  public void cleanup() throws Exception {
+    admin.deleteSnapshots(".*").get();
+    admin.disableTables(".*").get();
+    admin.deleteTables(".*").get();
+  }
+
+  @Test
+  public void testTakeSnapshot() throws Exception {
     Admin syncAdmin = TEST_UTIL.getAdmin();
 
-    try {
     Table table = TEST_UTIL.createTable(tableName, Bytes.toBytes("f1"));
     for (int i = 0; i < 3000; i++) {
       table.put(new Put(Bytes.toBytes(i)).addColumn(Bytes.toBytes("f1"), Bytes.toBytes("cq"),
@@ -65,21 +85,13 @@ public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
     Assert.assertEquals(snapshotName2, snapshots.get(1).getName());
     Assert.assertEquals(tableName, snapshots.get(1).getTableName());
     Assert.assertEquals(SnapshotType.FLUSH, snapshots.get(1).getType());
-    } finally {
-      syncAdmin.deleteSnapshot(snapshotName1);
-      syncAdmin.deleteSnapshot(snapshotName2);
-      TEST_UTIL.deleteTable(tableName);
-    }
   }
 
   @Test
   public void testCloneSnapshot() throws Exception {
-    String snapshotName1 = "snapshotName1";
-    TableName tableName = TableName.valueOf("testCloneSnapshot");
     TableName tableName2 = TableName.valueOf("testCloneSnapshot2");
     Admin syncAdmin = TEST_UTIL.getAdmin();
 
-    try {
     Table table = TEST_UTIL.createTable(tableName, Bytes.toBytes("f1"));
     for (int i = 0; i < 3000; i++) {
       table.put(new Put(Bytes.toBytes(i)).addColumn(Bytes.toBytes("f1"), Bytes.toBytes("cq"),
@@ -106,10 +118,6 @@ public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
     Assert.assertTrue(!syncAdmin.tableExists(tableName2));
     admin.cloneSnapshot(snapshotName1, tableName2).get();
     syncAdmin.tableExists(tableName2);
-    } finally {
-      syncAdmin.deleteSnapshot(snapshotName1);
-      TEST_UTIL.deleteTable(tableName);
-    }
   }
 
   private void assertResult(TableName tableName, int expectedRowCount) throws IOException {
@@ -131,12 +139,6 @@ public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
 
   @Test
   public void testRestoreSnapshot() throws Exception {
-    String snapshotName1 = "snapshotName1";
-    String snapshotName2 = "snapshotName2";
-    TableName tableName = TableName.valueOf("testRestoreSnapshot");
-    Admin syncAdmin = TEST_UTIL.getAdmin();
-
-    try {
     Table table = TEST_UTIL.createTable(tableName, Bytes.toBytes("f1"));
     for (int i = 0; i < 3000; i++) {
       table.put(new Put(Bytes.toBytes(i)).addColumn(Bytes.toBytes("f1"), Bytes.toBytes("cq"),
@@ -157,22 +159,10 @@ public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
     admin.restoreSnapshot(snapshotName2, false).get();
     admin.enableTable(tableName).get();
     assertResult(tableName, 3000);
-    } finally {
-      syncAdmin.deleteSnapshot(snapshotName1);
-      syncAdmin.deleteSnapshot(snapshotName2);
-      TEST_UTIL.deleteTable(tableName);
-    }
   }
 
   @Test
   public void testListSnapshots() throws Exception {
-    String snapshotName1 = "snapshotName1";
-    String snapshotName2 = "snapshotName2";
-    String snapshotName3 = "snapshotName3";
-    TableName tableName = TableName.valueOf("testListSnapshots");
-    Admin syncAdmin = TEST_UTIL.getAdmin();
-
-    try {
     Table table = TEST_UTIL.createTable(tableName, Bytes.toBytes("f1"));
     for (int i = 0; i < 3000; i++) {
       table.put(new Put(Bytes.toBytes(i)).addColumn(Bytes.toBytes("f1"), Bytes.toBytes("cq"),
@@ -192,23 +182,10 @@ public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
     Assert.assertEquals(admin.listTableSnapshots("testListSnapshots", "s(.*)").get().size(), 3);
     Assert.assertEquals(admin.listTableSnapshots("fakeTableName", "snap(.*)").get().size(), 0);
     Assert.assertEquals(admin.listTableSnapshots("test(.*)", "snap(.*)[1|3]").get().size(), 2);
-
-    } finally {
-      syncAdmin.deleteSnapshot(snapshotName1);
-      syncAdmin.deleteSnapshot(snapshotName2);
-      syncAdmin.deleteSnapshot(snapshotName3);
-      TEST_UTIL.deleteTable(tableName);
-    }
   }
 
   @Test
   public void testDeleteSnapshots() throws Exception {
-    String snapshotName1 = "snapshotName1";
-    String snapshotName2 = "snapshotName2";
-    String snapshotName3 = "snapshotName3";
-    TableName tableName = TableName.valueOf("testDeleteSnapshots");
-
-    try {
     Table table = TEST_UTIL.createTable(tableName, Bytes.toBytes("f1"));
     for (int i = 0; i < 3000; i++) {
       table.put(new Put(Bytes.toBytes(i)).addColumn(Bytes.toBytes("f1"), Bytes.toBytes("cq"),
@@ -238,8 +215,5 @@ public class TestAsyncSnapshotAdminApi extends TestAsyncAdminBase {
 
     admin.deleteTableSnapshots("(.*)", "(.*)3").get();
     Assert.assertEquals(admin.listSnapshots().get().size(), 0);
-    } finally {
-      TEST_UTIL.deleteTable(tableName);
-    }
   }
 }