Add timeouts to presto-router TestHealthChecks

Author: ZacBlanco

presto-router/src/test/java/com/facebook/presto/router/TestHealthChecks.java

@@ -21,6 +21,7 @@
 import com.facebook.airlift.node.testing.TestingNodeModule;
 import com.facebook.presto.ClientRequestFilterModule;
 import com.facebook.presto.router.cluster.ClusterManager;
+import com.facebook.presto.router.cluster.RemoteClusterInfo;
 import com.facebook.presto.router.cluster.RequestInfo;
 import com.facebook.presto.server.MockHttpServletRequest;
 import com.facebook.presto.server.security.ServerSecurityModule;
@@ -38,11 +39,19 @@
 
 import java.io.File;
 import java.net.URI;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.function.Supplier;
 
 import static com.facebook.presto.router.TestingRouterUtil.getConfigFile;
+import static com.google.common.base.Preconditions.checkArgument;
+import static java.lang.Thread.sleep;
+import static java.time.Instant.now;
+import static java.util.concurrent.TimeUnit.MINUTES;
 import static org.testng.Assert.assertFalse;
 import static org.testng.Assert.assertTrue;
 
@@ -99,9 +108,28 @@ public void tearDownServer()
         }
     }
 
+    static void waitUntil(Supplier<Boolean> condition, int value, TimeUnit unit)
+            throws TimeoutException, InterruptedException
+    {
+        checkArgument(value > 0, "timeout value must be greater than 0");
+        Instant start = now();
+        long timeoutMillis = unit.toMillis(value);
+        long sleepMillis = Math.min(timeoutMillis / 10, 50);
+        Instant deadline = start.plusMillis(timeoutMillis);
+        while (true) {
+            if (condition.get()) {
+                return;
+            }
+            if (now().isAfter(deadline)) {
+                throw new TimeoutException();
+            }
+            sleep(sleepMillis);
+        }
+    }
+
     @Test
     public void testHealthChecks()
-            throws InterruptedException
+            throws InterruptedException, TimeoutException
     {
         TestingPrestoServer server0 = prestoServers.get(0);
         TestingPrestoServer server1 = prestoServers.get(1);
@@ -113,25 +141,22 @@ public void testHealthChecks()
         assertTrue(healthyDestinations.contains(server2.getBaseUrl()));
 
         server0.stopResponding();
-        while (
-                clusterManager.getRemoteClusterInfos().get(server0.getBaseUrl()).isHealthy()
-                        || !clusterManager.getRemoteClusterInfos().get(server1.getBaseUrl()).isHealthy()
-                        || !clusterManager.getRemoteClusterInfos().get(server2.getBaseUrl()).isHealthy()) {
-            Thread.sleep(10);
-        }
+        waitUntil(() ->
+                !clusterManager.getRemoteClusterInfos().get(server0.getBaseUrl()).isHealthy()
+                        && clusterManager.getRemoteClusterInfos().get(server1.getBaseUrl()).isHealthy()
+                        && clusterManager.getRemoteClusterInfos().get(server2.getBaseUrl()).isHealthy(), 2, MINUTES);
 
         healthyDestinations = getDestinations(3);
         assertFalse(healthyDestinations.contains(server0.getBaseUrl()));
         assertTrue(healthyDestinations.contains(server1.getBaseUrl()));
         assertTrue(healthyDestinations.contains(server2.getBaseUrl()));
 
         server0.startResponding();
-        while (
-                !clusterManager.getRemoteClusterInfos().get(server0.getBaseUrl()).isHealthy()
-                        || !clusterManager.getRemoteClusterInfos().get(server1.getBaseUrl()).isHealthy()
-                        || !clusterManager.getRemoteClusterInfos().get(server2.getBaseUrl()).isHealthy()) {
-            Thread.sleep(10);
-        }
+        waitUntil(() -> prestoServers.stream()
+                        .map(TestingPrestoServer::getBaseUrl)
+                        .map(clusterManager.getRemoteClusterInfos()::get)
+                        .allMatch(RemoteClusterInfo::isHealthy),
+                2, MINUTES);
 
         healthyDestinations = getDestinations(3);
         assertTrue(healthyDestinations.contains(server0.getBaseUrl()));