HDFS-16336. De-flake TestRollingUpgrade#testRollback (#3686)

Signed-off-by: Akira Ajisaka <aajisaka@apache.org>

Author: virajjasani

hadoop-hdfs-project/hadoop-hdfs/src/test/java/org/apache/hadoop/hdfs/TestRollingUpgrade.java

@@ -22,6 +22,7 @@
 import java.lang.management.ManagementFactory;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeoutException;
 
 import javax.management.AttributeNotFoundException;
 import javax.management.InstanceNotFoundException;
@@ -325,7 +326,7 @@ public void testRollback() throws Exception {
       out.write(data, 0, data.length);
       out.close();
 
-      checkMxBeanIsNull();
+      waitForNullMxBean();
       startRollingUpgrade(foo, bar, file, data, cluster);
       checkMxBean();
       cluster.getFileSystem().rollEdits();
@@ -356,6 +357,18 @@ public void testRollback() throws Exception {
     }
   }
 
+  private void waitForNullMxBean() throws TimeoutException, InterruptedException {
+    GenericTestUtils.waitFor(() -> {
+      try {
+        checkMxBeanIsNull();
+        return true;
+      } catch (Throwable t) {
+        LOG.error("Something went wrong.", t);
+        return false;
+      }
+    }, 100, 8000, "RollingUpgradeStatus is already set");
+  }
+
   private static void startRollingUpgrade(Path foo, Path bar,
       Path file, byte[] data,
       MiniDFSCluster cluster) throws IOException {