Minor refactoring, move into separate package

3 files changed, 50 insertions, 9 deletions

diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/ProxyServer.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/ProxyServer.java
index d93eb819463..55e546072fc 100644
--- a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/ProxyServer.java
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/ProxyServer.java
@@ -9,12 +9,10 @@ import com.yahoo.jrt.Transport;
 import com.yahoo.log.LogLevel;
 import com.yahoo.log.LogSetup;
 import com.yahoo.log.event.Event;
-import com.yahoo.vespa.config.JRTConnectionPool;
 import com.yahoo.vespa.config.RawConfig;
 import com.yahoo.vespa.config.TimingValues;
 import com.yahoo.vespa.config.protocol.JRTServerConfigRequest;
-import com.yahoo.vespa.filedistribution.FileDistributionRpcServer;
-import com.yahoo.vespa.filedistribution.FileDownloader;
+import com.yahoo.vespa.config.proxy.filedistribution.FileDistributionAndUrlDownload;
 import com.yahoo.yolean.system.CatchSignals;
 
 import java.util.List;
@@ -61,6 +59,7 @@ public class ProxyServer implements Runnable {
     private static final double timingValuesRatio = 0.8;
     private final static TimingValues defaultTimingValues;
     private final boolean delayedResponseHandling;
+    private final FileDistributionAndUrlDownload fileDistributionAndUrlDownload;
 
     private volatile Mode mode = new Mode(DEFAULT);
 
@@ -88,8 +87,7 @@ public class ProxyServer implements Runnable {
         this.rpcServer = createRpcServer(spec);
         clientUpdater = new ClientUpdater(rpcServer, statistics, delayedResponses);
         this.configClient = createClient(clientUpdater, delayedResponses, source, timingValues, memoryCache, configClient);
-        new FileDistributionRpcServer(supervisor, new FileDownloader(new JRTConnectionPool(source)));
-        new UrlDownloadRpcServer(supervisor);
+        this.fileDistributionAndUrlDownload = new FileDistributionAndUrlDownload(supervisor, source);
     }
 
     static ProxyServer createTestServer(ConfigSourceSet source) {
@@ -267,6 +265,7 @@ public class ProxyServer implements Runnable {
         if (statistics != null) {
             statistics.stop();
         }
+        fileDistributionAndUrlDownload.close();
     }
 
     MemoryCache getMemoryCache() {
diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/FileDistributionAndUrlDownload.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/FileDistributionAndUrlDownload.java
new file mode 100644
index 00000000000..4eef3c40df4
--- /dev/null
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/FileDistributionAndUrlDownload.java
@@ -0,0 +1,32 @@
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.config.proxy.filedistribution;
+
+import com.yahoo.config.subscription.ConfigSourceSet;
+import com.yahoo.jrt.Supervisor;
+import com.yahoo.vespa.config.JRTConnectionPool;
+import com.yahoo.vespa.filedistribution.FileDistributionRpcServer;
+import com.yahoo.vespa.filedistribution.FileDownloader;
+
+import java.util.stream.Stream;
+
+/**
+ * Keeps track of file distribution and url download rpc servers.
+ *
+ * @author hmusum
+ */
+public class FileDistributionAndUrlDownload {
+
+    private final FileDistributionRpcServer fileDistributionRpcServer;
+    private final UrlDownloadRpcServer urlDownloadRpcServer;
+
+    public FileDistributionAndUrlDownload(Supervisor supervisor, ConfigSourceSet source) {
+        fileDistributionRpcServer = new FileDistributionRpcServer(supervisor, new FileDownloader(new JRTConnectionPool(source)));
+        urlDownloadRpcServer = new UrlDownloadRpcServer(supervisor);
+    }
+
+    public void close() {
+        fileDistributionRpcServer.close();
+        urlDownloadRpcServer.close();
+    }
+
+}
diff --git a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/UrlDownloadRpcServer.java b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloadRpcServer.java
index 711c43340cb..9d89f1d10b2 100644
--- a/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/UrlDownloadRpcServer.java
+++ b/config-proxy/src/main/java/com/yahoo/vespa/config/proxy/filedistribution/UrlDownloadRpcServer.java
@@ -1,5 +1,5 @@
-// Copyright 2019 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-package com.yahoo.vespa.config.proxy;
+// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.config.proxy.filedistribution;
 
 import com.yahoo.concurrent.DaemonThreadFactory;
 import com.yahoo.jrt.Method;
@@ -26,6 +26,7 @@ import java.nio.channels.ReadableByteChannel;
 import java.nio.file.Files;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 import java.util.logging.Logger;
 
 import static com.yahoo.vespa.config.UrlDownloader.DOES_NOT_EXIST;
@@ -37,7 +38,7 @@ import static com.yahoo.vespa.config.UrlDownloader.INTERNAL_ERROR;
  *
  * @author lesters
  */
-public class UrlDownloadRpcServer {
+class UrlDownloadRpcServer {
     private final static Logger log = Logger.getLogger(UrlDownloadRpcServer.class.getName());
 
     private static final String CONTENTS_FILE_NAME = "contents";
@@ -45,7 +46,7 @@ public class UrlDownloadRpcServer {
 
     private final File downloadBaseDir;
     private final ExecutorService rpcDownloadExecutor = Executors.newFixedThreadPool(Math.max(8, Runtime.getRuntime().availableProcessors()),
-                                                                                     new DaemonThreadFactory("Rpc download executor"));
+                                                                                     new DaemonThreadFactory("Rpc URL download executor"));
 
     UrlDownloadRpcServer(Supervisor supervisor) {
         supervisor.addMethod(new Method("url.waitFor", "s", "s", this::download)
@@ -55,6 +56,15 @@ public class UrlDownloadRpcServer {
         downloadBaseDir = new File(Defaults.getDefaults().underVespaHome("var/db/vespa/download"));
     }
 
+    void close() {
+        rpcDownloadExecutor.shutdownNow();
+        try {
+            rpcDownloadExecutor.awaitTermination(10, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private void download(Request req) {
         req.detach();
         rpcDownloadExecutor.execute(() -> downloadFile(req));