Use a flag to choose which connection pool to use for file distribution

3 files changed, 29 insertions, 12 deletions

diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java
index cfc7a88568c..f4801c5a7ea 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/filedistribution/FileServer.java
@@ -16,11 +16,14 @@ import com.yahoo.vespa.config.JRTConnectionPool;
 import com.yahoo.vespa.defaults.Defaults;
 import com.yahoo.vespa.filedistribution.CompressedFileReference;
 import com.yahoo.vespa.filedistribution.EmptyFileReferenceData;
+import com.yahoo.vespa.filedistribution.FileDistributionConnectionPool;
 import com.yahoo.vespa.filedistribution.FileDownloader;
 import com.yahoo.vespa.filedistribution.FileReferenceData;
 import com.yahoo.vespa.filedistribution.FileReferenceDownload;
 import com.yahoo.vespa.filedistribution.LazyFileReferenceData;
 import com.yahoo.vespa.filedistribution.LazyTemporaryStorageFileReferenceData;
+import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.Flags;
 import com.yahoo.yolean.Exceptions;
 
 import java.io.File;
@@ -74,14 +77,15 @@ public class FileServer {
 
     @SuppressWarnings("WeakerAccess") // Created by dependency injection
     @Inject
-    public FileServer(ConfigserverConfig configserverConfig) {
+    public FileServer(ConfigserverConfig configserverConfig, FlagSource flagSource) {
         this(new File(Defaults.getDefaults().underVespaHome(configserverConfig.fileReferencesDir())),
-             createFileDownloader(getOtherConfigServersInCluster(configserverConfig)));
+             createFileDownloader(getOtherConfigServersInCluster(configserverConfig),
+             Flags.USE_FILE_DISTRIBUTION_CONNECTION_POOL.bindTo(flagSource).value()));
     }
 
     // For testing only
     public FileServer(File rootDir) {
-        this(rootDir, createFileDownloader(List.of()));
+        this(rootDir, createFileDownloader(List.of(), true));
     }
 
     public FileServer(File rootDir, FileDownloader fileDownloader) {
@@ -204,18 +208,22 @@ public class FileServer {
         executor.shutdown();
     }
 
-    private static FileDownloader createFileDownloader(List<String> configServers) {
+    private static FileDownloader createFileDownloader(List<String> configServers, boolean useFileDistributionConnectionPool) {
         Supervisor supervisor = new Supervisor(new Transport("filedistribution-pool")).setDropEmptyBuffers(true);
         return new FileDownloader(configServers.isEmpty()
                                           ? FileDownloader.emptyConnectionPool()
-                                          : createConnectionPool(configServers, supervisor),
+                                          : createConnectionPool(configServers, supervisor, useFileDistributionConnectionPool),
                                   supervisor);
     }
 
-    private static ConnectionPool createConnectionPool(List<String> configServers, Supervisor supervisor) {
-        return configServers.size() > 0
-                ? new JRTConnectionPool(new ConfigSourceSet(configServers), supervisor)
-                : FileDownloader.emptyConnectionPool();
+    private static ConnectionPool createConnectionPool(List<String> configServers, Supervisor supervisor, boolean useFileDistributionConnectionPool) {
+        ConfigSourceSet configSourceSet = new ConfigSourceSet(configServers);
+
+        if (configServers.size() == 0) return FileDownloader.emptyConnectionPool();
+
+        return useFileDistributionConnectionPool
+                ? new FileDistributionConnectionPool(configSourceSet, supervisor)
+                : new JRTConnectionPool(configSourceSet, supervisor);
     }
 
 }
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ApplicationPackageMaintainer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ApplicationPackageMaintainer.java
index 1f30781ff03..08c300220df 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ApplicationPackageMaintainer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/maintenance/ApplicationPackageMaintainer.java
@@ -16,7 +16,9 @@ import com.yahoo.vespa.curator.Curator;
 import com.yahoo.vespa.defaults.Defaults;
 import com.yahoo.vespa.filedistribution.FileDownloader;
 import com.yahoo.vespa.filedistribution.FileReferenceDownload;
+import com.yahoo.vespa.filedistribution.FileDistributionConnectionPool;
 import com.yahoo.vespa.flags.FlagSource;
+import com.yahoo.vespa.flags.Flags;
 
 import java.io.File;
 import java.time.Duration;
@@ -40,6 +42,8 @@ public class ApplicationPackageMaintainer extends ConfigServerMaintainer {
     private final File downloadDirectory;
     private final ConfigserverConfig configserverConfig;
     private final Supervisor supervisor;
+    private final boolean useFileDistributionConnectionPool;
+
 
     ApplicationPackageMaintainer(ApplicationRepository applicationRepository,
                                  Curator curator,
@@ -49,7 +53,8 @@ public class ApplicationPackageMaintainer extends ConfigServerMaintainer {
         this.applicationRepository = applicationRepository;
         this.configserverConfig = applicationRepository.configserverConfig();
         this.supervisor = new Supervisor(new Transport("filedistribution-pool")).setDropEmptyBuffers(true);
-        downloadDirectory = new File(Defaults.getDefaults().underVespaHome(configserverConfig.fileReferencesDir()));
+        this.downloadDirectory = new File(Defaults.getDefaults().underVespaHome(configserverConfig.fileReferencesDir()));
+        this.useFileDistributionConnectionPool = Flags.USE_FILE_DISTRIBUTION_CONNECTION_POOL.bindTo(flagSource).value();
     }
 
     @Override
@@ -90,7 +95,10 @@ public class ApplicationPackageMaintainer extends ConfigServerMaintainer {
     }
 
     private FileDownloader createFileDownloader() {
-        return new FileDownloader(new JRTConnectionPool(new ConfigSourceSet(getOtherConfigServersInCluster(configserverConfig)), supervisor),
+        ConfigSourceSet configSourceSet = new ConfigSourceSet(getOtherConfigServersInCluster(configserverConfig));
+        return new FileDownloader(useFileDistributionConnectionPool
+                                          ? new FileDistributionConnectionPool(configSourceSet, supervisor)
+                                          : new JRTConnectionPool(configSourceSet, supervisor),
                                   supervisor,
                                   downloadDirectory);
     }
diff --git a/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/FileServerTest.java b/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/FileServerTest.java
index 29ec11bad26..f85ca37a351 100644
--- a/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/FileServerTest.java
+++ b/configserver/src/test/java/com/yahoo/vespa/config/server/filedistribution/FileServerTest.java
@@ -10,6 +10,7 @@ import com.yahoo.net.HostName;
 import com.yahoo.vespa.filedistribution.FileDownloader;
 import com.yahoo.vespa.filedistribution.FileReferenceData;
 import com.yahoo.vespa.filedistribution.FileReferenceDownload;
+import com.yahoo.vespa.flags.InMemoryFlagSource;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -117,7 +118,7 @@ public class FileServerTest {
     private FileServer createFileServer(ConfigserverConfig.Builder configBuilder) throws IOException {
         File fileReferencesDir = temporaryFolder.newFolder();
         configBuilder.fileReferencesDir(fileReferencesDir.getAbsolutePath());
-        return new FileServer(new ConfigserverConfig(configBuilder));
+        return new FileServer(new ConfigserverConfig(configBuilder), new InMemoryFlagSource());
     }
 
     private static class FileReceiver implements FileServer.Receiver {