Wire in use of compression types and flags for file distribution

VESPA_FILE_DISTRIBUTION_ACCEPTED_COMPRESSION_TYPES is not set anywhere
yet, will be set through host-admin

4 files changed, 54 insertions, 16 deletions

diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java
index 5941ed536a8..dc87ae2a0b4 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileDownloader.java
@@ -6,11 +6,11 @@ import com.yahoo.jrt.Supervisor;
 import com.yahoo.vespa.config.Connection;
 import com.yahoo.vespa.config.ConnectionPool;
 import com.yahoo.vespa.defaults.Defaults;
-
 import java.io.File;
 import java.time.Duration;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Future;
@@ -19,6 +19,8 @@ import java.util.concurrent.TimeoutException;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType;
+
 /**
  * Handles downloads of files (file references only for now)
  *
@@ -37,19 +39,20 @@ public class FileDownloader implements AutoCloseable {
     private final FileReferenceDownloader fileReferenceDownloader;
     private final Downloads downloads = new Downloads();
 
-    public FileDownloader(ConnectionPool connectionPool, Supervisor supervisor, Duration timeout) {
-        this(connectionPool, supervisor, defaultDownloadDirectory, timeout, defaultSleepBetweenRetries);
+    public FileDownloader(ConnectionPool connectionPool, Supervisor supervisor, Duration timeout, Set<CompressionType> acceptedCompressionTypes) {
+        this(connectionPool, supervisor, defaultDownloadDirectory, timeout, defaultSleepBetweenRetries, acceptedCompressionTypes);
     }
 
-    public FileDownloader(ConnectionPool connectionPool, Supervisor supervisor, File downloadDirectory, Duration timeout) {
-        this(connectionPool, supervisor, downloadDirectory, timeout, defaultSleepBetweenRetries);
+    public FileDownloader(ConnectionPool connectionPool, Supervisor supervisor, File downloadDirectory, Duration timeout, Set<CompressionType> acceptedCompressionTypes) {
+        this(connectionPool, supervisor, downloadDirectory, timeout, defaultSleepBetweenRetries, acceptedCompressionTypes);
     }
 
     public FileDownloader(ConnectionPool connectionPool,
                           Supervisor supervisor,
                           File downloadDirectory,
                           Duration timeout,
-                          Duration sleepBetweenRetries) {
+                          Duration sleepBetweenRetries,
+                          Set<CompressionType> acceptedCompressionTypes) {
         this.connectionPool = connectionPool;
         this.supervisor = supervisor;
         this.downloadDirectory = downloadDirectory;
@@ -60,7 +63,8 @@ public class FileDownloader implements AutoCloseable {
                                                                    downloads,
                                                                    timeout,
                                                                    sleepBetweenRetries,
-                                                                   downloadDirectory);
+                                                                   downloadDirectory,
+                                                                   acceptedCompressionTypes);
     }
 
     public Optional<File> getFile(FileReferenceDownload fileReferenceDownload) {
diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownload.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownload.java
index 8d6f428eaef..f3a8cf9299d 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownload.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownload.java
@@ -18,7 +18,7 @@ public class FileReferenceDownload {
     private final boolean downloadFromOtherSourceIfNotFound;
     private final String client;
 
-    public FileReferenceDownload(FileReference fileReference,  String client) {
+    public FileReferenceDownload(FileReference fileReference, String client) {
         this(fileReference, client, true);
     }
 
diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
index 0267feb9ffc..7078c5aae6c 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceDownloader.java
@@ -5,14 +5,16 @@ import com.yahoo.concurrent.DaemonThreadFactory;
 import com.yahoo.config.FileReference;
 import com.yahoo.jrt.Int32Value;
 import com.yahoo.jrt.Request;
+import com.yahoo.jrt.StringArray;
 import com.yahoo.jrt.StringValue;
 import com.yahoo.vespa.config.Connection;
 import com.yahoo.vespa.config.ConnectionPool;
-
 import java.io.File;
 import java.time.Duration;
 import java.time.Instant;
+import java.util.Objects;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
@@ -20,8 +22,10 @@ import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType;
+
 /**
- * Downloads file reference using rpc requests to config server and keeps track of files being downloaded
+ * Downloads file reference from config server and keeps track of files being downloaded
  *
  * @author hmusum
  */
@@ -38,12 +42,14 @@ public class FileReferenceDownloader {
     private final Duration sleepBetweenRetries;
     private final Duration rpcTimeout;
     private final File downloadDirectory;
+    private final Set<CompressionType> acceptedCompressionTypes;
 
     FileReferenceDownloader(ConnectionPool connectionPool,
                             Downloads downloads,
                             Duration timeout,
                             Duration sleepBetweenRetries,
-                            File downloadDirectory) {
+                            File downloadDirectory,
+                            Set<CompressionType> acceptedCompressionTypes) {
         this.connectionPool = connectionPool;
         this.downloads = downloads;
         this.downloadTimeout = timeout;
@@ -51,6 +57,7 @@ public class FileReferenceDownloader {
         this.downloadDirectory = downloadDirectory;
         String timeoutString = System.getenv("VESPA_CONFIGPROXY_FILEDOWNLOAD_RPC_TIMEOUT");
         this.rpcTimeout = Duration.ofSeconds(timeoutString == null ? 30 : Integer.parseInt(timeoutString));
+        this.acceptedCompressionTypes = requireNonEmpty(acceptedCompressionTypes);
     }
 
     private void waitUntilDownloadStarted(FileReferenceDownload fileReferenceDownload) {
@@ -132,6 +139,9 @@ public class FileReferenceDownloader {
         Request request = new Request("filedistribution.serveFile");
         request.parameters().add(new StringValue(fileReferenceDownload.fileReference().value()));
         request.parameters().add(new Int32Value(fileReferenceDownload.downloadFromOtherSourceIfNotFound() ? 0 : 1));
+        String[] temp = new String[acceptedCompressionTypes.size()];
+        acceptedCompressionTypes.stream().map(Enum::name).toList().toArray(temp);
+        request.parameters().add(new StringArray(temp));
         return request;
     }
 
@@ -160,4 +170,9 @@ public class FileReferenceDownloader {
         }
     }
 
+    private static Set<CompressionType> requireNonEmpty(Set<CompressionType> s) {
+        if (Objects.requireNonNull(s).isEmpty()) throw new IllegalArgumentException("set must be non-empty");
+        return s;
+    }
+
 }
diff --git a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
index f5cd1760e89..629ea5915df 100644
--- a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
+++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
@@ -17,7 +17,6 @@ import net.jpountz.xxhash.XXHashFactory;
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
-
 import java.io.File;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -26,6 +25,7 @@ import java.nio.file.Path;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
@@ -41,6 +41,7 @@ import static org.junit.Assert.fail;
 
 public class FileDownloaderTest {
     private static final Duration sleepBetweenRetries = Duration.ofMillis(10);
+    private static final Set<FileReferenceData.CompressionType> acceptedCompressionTypes = Set.of(gzip);
 
     private MockConnection connection;
     private FileDownloader fileDownloader;
@@ -53,7 +54,7 @@ public class FileDownloaderTest {
             downloadDir = Files.createTempDirectory("filedistribution").toFile();
             connection = new MockConnection();
             supervisor = new Supervisor(new Transport()).setDropEmptyBuffers(true);
-            fileDownloader = new FileDownloader(connection, supervisor, downloadDir, Duration.ofSeconds(1), sleepBetweenRetries);
+            fileDownloader = createDownloader(connection, Duration.ofSeconds(1));
         } catch (IOException e) {
             e.printStackTrace();
             fail(e.getMessage());
@@ -167,7 +168,7 @@ public class FileDownloaderTest {
 
     @Test
     public void getFileWhenConnectionError() throws IOException {
-        fileDownloader = new FileDownloader(connection, supervisor, downloadDir, Duration.ofSeconds(2), sleepBetweenRetries);
+        fileDownloader = createDownloader(connection, Duration.ofSeconds(2));
         File downloadDir = fileDownloader.downloadDirectory();
 
         int timesToFail = 2;
@@ -201,7 +202,7 @@ public class FileDownloaderTest {
     public void getFileWhenDownloadInProgress() throws IOException, ExecutionException, InterruptedException {
         ExecutorService executor = Executors.newFixedThreadPool(2);
         String filename = "abc.jar";
-        fileDownloader = new FileDownloader(connection, supervisor, downloadDir, Duration.ofSeconds(3), sleepBetweenRetries);
+        fileDownloader = createDownloader(connection, Duration.ofSeconds(3));
         File downloadDir = fileDownloader.downloadDirectory();
 
         // Delay response so that we can make a second request while downloading the file from the first request
@@ -241,7 +242,7 @@ public class FileDownloaderTest {
         Duration timeout = Duration.ofMillis(200);
         MockConnection connectionPool = new MockConnection();
         connectionPool.setResponseHandler(new MockConnection.WaitResponseHandler(timeout.plus(Duration.ofMillis(1000))));
-        FileDownloader fileDownloader = new FileDownloader(connectionPool, supervisor, downloadDir, timeout, sleepBetweenRetries);
+        FileDownloader fileDownloader = createDownloader(connectionPool, timeout);
         FileReference xyzzy = new FileReference("xyzzy");
         // Should download since we do not have the file on disk
         fileDownloader.downloadIfNeeded(new FileReferenceDownload(xyzzy, "test"));
@@ -264,6 +265,16 @@ public class FileDownloaderTest {
         assertEquals("content", IOUtils.readFile(downloadedFile));
     }
 
+    @Test
+    public void testCompressionTypes() {
+        try {
+            createDownloader(connection, Duration.ofSeconds(1), Set.of());
+            fail("expected to fail when set is empty");
+        } catch (IllegalArgumentException e) {
+            // ignore
+        }
+    }
+
     private void writeFileReference(File dir, String fileReferenceString, String fileName) throws IOException {
         File fileReferenceDir = new File(dir, fileReferenceString);
         fileReferenceDir.mkdir();
@@ -302,6 +313,14 @@ public class FileDownloaderTest {
         return fileDownloader.getFile(new FileReferenceDownload(fileReference, "test"));
     }
 
+    private FileDownloader createDownloader(MockConnection connection, Duration timeout) {
+        return  createDownloader(connection, timeout, acceptedCompressionTypes);
+    }
+
+    private FileDownloader createDownloader(MockConnection connection, Duration timeout, Set<FileReferenceData.CompressionType> acceptedCompressionTypes) {
+        return new FileDownloader(connection, supervisor, downloadDir, timeout, sleepBetweenRetries, acceptedCompressionTypes);
+    }
+
     private static class MockConnection implements ConnectionPool, com.yahoo.vespa.config.Connection {
 
         private ResponseHandler responseHandler;