Fix ClusterManager failing to poll

Previously, ClusterManager relies on CompletableFuture.supplyAsync
which uses the ForkJoin commonPool to execute tasks. On small
machines this pool has a limited number of threads and the previous
implementation used an infinite loop that would never break. This
caused a starvation of resources which could cause tests to hang and
due to the check never completing.

The new implementation uses a single-threaded scheduled executor to
prevent starvation and to properly setup/teardown resources. It also
improves the test flakiness by ensuring the file watch service starts
picking up changes to the config files first before moving on to
later test cases.

Author: ZacBlanco

presto-router/src/main/java/com/facebook/presto/router/RouterUtil.java

@@ -23,7 +23,7 @@
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.nio.file.Files;
-import java.nio.file.Paths;
+import java.nio.file.Path;
 import java.util.Optional;
 
 import static com.fasterxml.jackson.databind.DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES;
@@ -37,11 +37,11 @@ public final class RouterUtil
 
     private RouterUtil() {}
 
-    public static Optional<RouterSpec> parseRouterConfig(RouterConfig config)
+    public static Optional<RouterSpec> parseRouterConfig(Path configFile)
     {
         Optional<RouterSpec> routerSpec;
         try {
-            routerSpec = Optional.of(CODEC.fromJson(Files.readAllBytes(Paths.get(config.getConfigFile()))));
+            routerSpec = Optional.of(CODEC.fromJson(Files.readAllBytes(configFile)));
         }
         catch (IOException e) {
             throw new UncheckedIOException(e);

presto-router/src/main/java/com/facebook/presto/router/cluster/ClusterManager.java

@@ -13,7 +13,6 @@
  */
 package com.facebook.presto.router.cluster;
 
-import com.facebook.airlift.bootstrap.LifeCycleManager;
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.router.RouterConfig;
 import com.facebook.presto.router.scheduler.CustomSchedulerManager;
@@ -34,25 +33,25 @@
 import io.airlift.units.Duration;
 import org.weakref.jmx.Managed;
 
-import javax.annotation.PostConstruct;
+import javax.annotation.PreDestroy;
 import javax.inject.Inject;
 
-import java.io.File;
 import java.io.IOException;
 import java.net.URI;
+import java.nio.file.ClosedWatchServiceException;
 import java.nio.file.FileSystems;
 import java.nio.file.Path;
-import java.nio.file.StandardWatchEventKinds;
+import java.nio.file.Paths;
 import java.nio.file.WatchEvent;
 import java.nio.file.WatchKey;
 import java.nio.file.WatchService;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
@@ -67,15 +66,19 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableList.toImmutableList;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static java.nio.file.StandardWatchEventKinds.ENTRY_CREATE;
+import static java.nio.file.StandardWatchEventKinds.ENTRY_MODIFY;
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
+import static java.util.concurrent.TimeUnit.SECONDS;
 
 public class ClusterManager
+        implements AutoCloseable
 {
     private final AtomicReference<ClusterManagerConfig> currentConfig = new AtomicReference<>();
 
-    public final RouterConfig routerConfig;
+    private final Path configFile;
     private final Logger log = Logger.get(ClusterManager.class);
 
     // Cluster status
@@ -84,25 +87,72 @@ public class ClusterManager
 
     private final AtomicBoolean isWatchServiceStarted = new AtomicBoolean();
     private final RemoteInfoFactory remoteInfoFactory;
-    private final LifeCycleManager lifeCycleManager;
     private final HashMap<String, HashMap<URI, Integer>> serverWeights = new HashMap<>();
     private final CustomSchedulerManager schedulerManager;
+    private final ScheduledExecutorService scheduledExecutorService;
+    private final ScheduledFuture<?> configDetection;
+    private final WatchService watchService;
+    private final WatchKey watchKey;
 
     @Inject
-    public ClusterManager(RouterConfig config, RemoteInfoFactory remoteInfoFactory, LifeCycleManager lifeCycleManager,
-                          CustomSchedulerManager schedulerManager)
+    public ClusterManager(RouterConfig config, RemoteInfoFactory remoteInfoFactory, CustomSchedulerManager schedulerManager)
+            throws IOException
     {
-        this.routerConfig = requireNonNull(config, "config is null");
+        this.configFile = Paths.get(requireNonNull(config, "config is null").getConfigFile());
         this.remoteInfoFactory = requireNonNull(remoteInfoFactory, "remoteInfoFactory is null");
-        this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifecycleManager is null");
         this.schedulerManager = schedulerManager;
-        onConfigChangeDetection();
-        this.initializeServerWeights();
+        reloadConfig();
+        initializeServerWeights();
+        watchService = FileSystems.getDefault().newWatchService();
+        Path parentDir = configFile.getParent();
+        log.info("Router config watch service monitoring %s", parentDir);
+        watchKey = parentDir.register(watchService,
+                new WatchEvent.Kind[] {ENTRY_CREATE, ENTRY_MODIFY},
+                SensitivityWatchEventModifier.HIGH);
+        log.info("Successfully registered watch service for %s", parentDir);
+        scheduledExecutorService = newSingleThreadScheduledExecutor();
+        configDetection = scheduledExecutorService.scheduleAtFixedRate(this::monitorConfig, 0, 3, SECONDS);
     }
 
-    protected void onConfigChangeDetection()
+    protected void monitorConfig()
     {
-        RouterSpec newRouterSpec = parseRouterConfig(routerConfig)
+        boolean reload = false;
+        try {
+            WatchKey key = watchService.poll(1, SECONDS);
+            if (key == null) {
+                return;
+            }
+            List<WatchEvent<?>> events = key.pollEvents();
+            log.debug("Changes to router config directory detected");
+            for (WatchEvent<?> event : events) {
+                log.debug("Event detected: %s, path: %s", event.kind().name(), event.context());
+                Path changed = (Path) event.context();
+                if (changed.endsWith(configFile.getFileName())) {
+                    reload = true;
+                    break;
+                }
+                else {
+                    log.debug("Change to %s ignored by ClusterManager (config file is %s)", event.context(), configFile);
+                }
+            }
+            key.reset();
+        }
+        catch (ClosedWatchServiceException e) {
+            log.warn("Watch service closed. Future updates configuration changes will not be detected.");
+        }
+        catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            log.warn("Watch service interrupted while waiting for configuration updates");
+        }
+
+        if (reload) {
+            reloadConfig();
+        }
+    }
+
+    protected void reloadConfig()
+    {
+        RouterSpec newRouterSpec = parseRouterConfig(configFile)
                 .orElseThrow(() -> new PrestoException(CONFIGURATION_INVALID, "Failed to load router config"));
         Map<String, GroupSpec> newGroups = newRouterSpec.getGroups().stream().collect(toImmutableMap(GroupSpec::getName, group -> group));
         List<SelectorRuleSpec> newGroupSelectors = ImmutableList.copyOf(newRouterSpec.getSelectors());
@@ -132,54 +182,6 @@ protected void onConfigChangeDetection()
         currentConfig.set(new ClusterManagerConfig(newGroups, newGroupSelectors, newScheduler, newSchedulerType));
     }
 
-    @PostConstruct
-    public void startConfigReloadTaskFileWatcher()
-    {
-        CompletableFuture.supplyAsync(() -> {
-            try (WatchService watchService = FileSystems.getDefault().newWatchService()) {
-                File routerConfigFile = new File(routerConfig.getConfigFile());
-                log.info("Router config watch service monitoring %s", routerConfig.getConfigFile());
-                Path parentDir = routerConfigFile.toPath().getParent();
-                parentDir.register(
-                        watchService,
-                        new WatchEvent.Kind[] {
-                                StandardWatchEventKinds.ENTRY_MODIFY,
-                                StandardWatchEventKinds.ENTRY_CREATE,
-                                StandardWatchEventKinds.ENTRY_DELETE,
-                                StandardWatchEventKinds.OVERFLOW},
-                        SensitivityWatchEventModifier.HIGH);
-                isWatchServiceStarted.set(true);
-                log.info("Successfully registered watch service for %s", parentDir);
-
-                while (true) {
-                    WatchKey key = watchService.take();
-                    log.info("Changes to router config directory detected: %s", routerConfigFile);
-                    for (WatchEvent<?> event : key.pollEvents()) {
-                        log.debug("Event detected: %s, path: %s", event.kind().name(), event.context());
-                        Path changed = (Path) event.context();
-                        if (changed.endsWith(routerConfigFile.getName())) {
-                            try {
-                                onConfigChangeDetection();
-                            }
-                            catch (Exception e) {
-                                log.error("Exception in config reload");
-                            }
-                        }
-                        else {
-                            log.debug("Config change to %s ignored by ClusterManager (config file is %s)", event.context(), routerConfigFile.getName());
-                        }
-                    }
-                    key.reset();
-                }
-            }
-            catch (IOException | InterruptedException e) {
-                log.error("Exception in file watcher loop while monitoring %s, %s", routerConfig.getConfigFile(), e);
-                lifeCycleManager.stop();
-                throw new RuntimeException(e);
-            }
-        });
-    }
-
     public List<URI> getAllClusters()
     {
         return currentConfig.get().getGroups().values().stream()
@@ -199,9 +201,9 @@ public Optional<URI> getDestination(RequestInfo requestInfo)
         GroupSpec groupSpec = config.getGroups().get(target.get());
 
         List<URI> healthyClusterURIs = groupSpec.getMembers().stream().filter((entry) ->
-                Optional.ofNullable(remoteClusterInfos.get(entry))
-                        .map(RemoteClusterInfo::isHealthy)
-                        .orElse(false))
+                        Optional.ofNullable(remoteClusterInfos.get(entry))
+                                .map(RemoteClusterInfo::isHealthy)
+                                .orElse(false))
                 .collect(Collectors.toList());
 
         if (healthyClusterURIs.isEmpty()) {
@@ -285,6 +287,21 @@ public boolean getIsWatchServiceStarted()
         return isWatchServiceStarted.get();
     }
 
+    @PreDestroy
+    @Override
+    public void close()
+            throws Exception
+    {
+        try {
+            watchKey.cancel();
+            watchService.close();
+        }
+        finally {
+            configDetection.cancel(true);
+            scheduledExecutorService.shutdownNow();
+        }
+    }
+
     public static class ClusterStatusTracker
     {
         private final Logger log = Logger.get(ClusterStatusTracker.class);

presto-router/src/main/java/com/facebook/presto/router/cluster/RemoteState.java

@@ -31,6 +31,7 @@
 import javax.annotation.concurrent.ThreadSafe;
 
 import java.net.URI;
+import java.nio.file.Paths;
 import java.util.Optional;
 import java.util.concurrent.Future;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -71,7 +72,7 @@ public RemoteState(HttpClient httpClient, URI remoteUri, RemoteStateConfig remot
         this.isHealthy = new AtomicBoolean(true);
         this.httpClient = requireNonNull(httpClient, "httpClient is null");
         this.remoteUri = requireNonNull(remoteUri, "remoteUri is null");
-        RouterSpec routerSpec = parseRouterConfig(routerConfig)
+        RouterSpec routerSpec = parseRouterConfig(Paths.get(routerConfig.getConfigFile()))
                 .orElseThrow(() -> new PrestoException(CONFIGURATION_INVALID, "Failed to load router config"));
         this.routerUserCredentials = routerSpec.getUserCredentials();
         this.clusterUnhealthyTimeout = remoteStateConfig.getClusterUnhealthyTimeout();

presto-router/src/main/java/com/facebook/presto/router/predictor/PredictorManager.java

@@ -21,6 +21,7 @@
 import javax.inject.Inject;
 
 import java.net.URI;
+import java.nio.file.Paths;
 import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -45,7 +46,7 @@ public class PredictorManager
     @Inject
     public PredictorManager(RemoteQueryFactory remoteQueryFactory, RouterConfig config)
     {
-        RouterSpec routerSpec = parseRouterConfig(config)
+        RouterSpec routerSpec = parseRouterConfig(Paths.get(config.getConfigFile()))
                 .orElseThrow(() -> new PrestoException(CONFIGURATION_INVALID, "Failed to load router config"));
 
         this.remoteQueryFactory = requireNonNull(remoteQueryFactory, "");

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

@@ -13,11 +13,11 @@
  */
 package com.facebook.presto.router;
 
-import com.facebook.airlift.bootstrap.LifeCycleManager;
 import com.facebook.presto.router.cluster.ClusterManager;
 import com.facebook.presto.router.cluster.RemoteInfoFactory;
 import com.facebook.presto.router.scheduler.CustomSchedulerManager;
 
+import java.io.IOException;
 import java.util.concurrent.BrokenBarrierException;
 import java.util.concurrent.CyclicBarrier;
 import java.util.concurrent.TimeUnit;
@@ -28,18 +28,18 @@ public class BarrierClusterManager
 {
     private final CyclicBarrier barrier;
 
-    public BarrierClusterManager(RouterConfig config, RemoteInfoFactory remoteInfoFactory, CyclicBarrier barrier, LifeCycleManager lifeCycleManager, CustomSchedulerManager schedulerManager)
+    public BarrierClusterManager(RouterConfig config, RemoteInfoFactory remoteInfoFactory, CyclicBarrier barrier, CustomSchedulerManager schedulerManager)
+            throws IOException
     {
-        super(config, remoteInfoFactory, lifeCycleManager, schedulerManager);
+        super(config, remoteInfoFactory, schedulerManager);
         this.barrier = barrier;
-        startConfigReloadTaskFileWatcher();
     }
 
     @Override
-    protected void onConfigChangeDetection()
+    protected void reloadConfig()
     {
         try {
-            super.onConfigChangeDetection();
+            super.reloadConfig();
             if (barrier != null) {
                 barrier.await(5, TimeUnit.SECONDS);
             }