Support compressing files (not just directories) in file distribution

Use file distribution compression feature flags and support compressing single
files

4 files changed, 60 insertions, 23 deletions

diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReceiver.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReceiver.java
index 81a8944149c..da52caefb7d 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReceiver.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReceiver.java
@@ -22,8 +22,6 @@ import java.util.concurrent.atomic.AtomicInteger;
 import java.util.logging.Level;
 import java.util.logging.Logger;
 
-import static com.yahoo.vespa.filedistribution.FileReferenceData.Type;
-
 /**
  * When asking for a file reference, this handles RPC callbacks from config server with file data and metadata.
  * Uses the same Supervisor as the original caller that requests files, so communication uses the same
@@ -109,32 +107,40 @@ public class FileReceiver {
             verifyHash(hash);
 
             File file = new File(fileReferenceDir, fileName);
-            File decompressedDir = null;
             try {
-                if (fileType == Type.file) {
-                    try {
+                switch (fileType) {
+                    case file:
+                        log.log(Level.FINE, () -> "Uncompressed file reference " + fileName + ", storing as " + file.getAbsolutePath());
                         Files.createDirectories(fileReferenceDir.toPath());
-                    } catch (IOException e) {
-                        log.log(Level.SEVERE, "Failed creating directory (" + fileReferenceDir.toPath() + "): " + e.getMessage(), e);
-                        throw new RuntimeException("Failed creating directory (" + fileReferenceDir.toPath() + "): ", e);
-                    }
-                    log.log(Level.FINE, () -> "Uncompressed file, moving to " + file.getAbsolutePath());
-                    moveFileToDestination(inprogressFile, file);
-                } else {
-                    decompressedDir = Files.createTempDirectory(tmpDir.toPath(), "archive").toFile();
-                    new FileReferenceCompressor(fileType).decompress(inprogressFile, decompressedDir);
-                    moveFileToDestination(decompressedDir, fileReferenceDir);
+                        moveFileToDestination(inprogressFile, file);
+                        break;
+                    case compressed:
+                        log.log(Level.FINE, () -> "Compressed file reference (directory)" + fileName + ", storing in " + fileReferenceDir.getAbsolutePath());
+                        decompress(fileType, fileReferenceDir, inprogressFile);
+                        break;
+                    default:
+                        throw new RuntimeException("Unknown file type " + fileType);
                 }
             } catch (IOException e) {
                 log.log(Level.SEVERE, "Failed writing file: " + e.getMessage(), e);
                 throw new RuntimeException("Failed writing file: ", e);
             } finally {
                 deletePath(inprogressFile);
-                deletePath(decompressedDir);
             }
             return file;
         }
 
+        void decompress(FileReferenceData.Type fileType, File fileReferenceDir, File file) throws IOException {
+            File decompressedDir = null;
+            try {
+                decompressedDir = Files.createTempDirectory(tmpDir.toPath(), "archive").toFile();
+                new FileReferenceCompressor(fileType).decompress(file, decompressedDir);
+                moveFileToDestination(decompressedDir, fileReferenceDir);
+            } finally {
+                deletePath(decompressedDir);
+            }
+        }
+
         double percentageReceived() {
             return (double)currentFileSize/(double)fileSize;
         }
diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java
index c36bcd22606..d0957297f79 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java
@@ -47,13 +47,16 @@ public class FileReferenceCompressor {
         return outputFile;
     }
 
-    public File compress(File directory, File outputFile) throws IOException {
-        return compress(directory,
-                        Files.find(Paths.get(directory.getAbsolutePath()),
-                                   recurseDepth,
-                                   (p, basicFileAttributes) -> basicFileAttributes.isRegularFile())
-                             .map(Path::toFile).collect(Collectors.toList()),
-                        outputFile);
+    public File compress(File file, File outputFile) throws IOException {
+        if (file.isDirectory())
+            return compress(file,
+                            Files.find(Paths.get(file.getAbsolutePath()),
+                                       recurseDepth,
+                                       (p, basicFileAttributes) -> basicFileAttributes.isRegularFile())
+                                 .map(Path::toFile).collect(Collectors.toList()),
+                            outputFile);
+        else
+            return compress(file.getParentFile(), List.of(file), outputFile);
     }
 
     public void decompress(File inputFile, File outputDir) throws IOException {
diff --git a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceData.java b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceData.java
index 03f8d184f94..f1ac8edaabb 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceData.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceData.java
@@ -64,6 +64,17 @@ public abstract class FileReferenceData {
      */
     public abstract void close();
 
+    public static Type from(String type) {
+        switch (type) {
+            case "none":
+                return Type.file;
+            case "gzip":
+                return Type.compressed;
+            default:
+                throw new RuntimeException("Unknown compression type " + type);
+        }
+    }
+
     @Override
     public String toString() { return fileReference.value() + " (" + filename + "), " + type.name(); }
 
diff --git a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java
index 5c15f945ae3..4996a439cfe 100644
--- a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java
+++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java
@@ -15,6 +15,8 @@ import java.io.FileWriter;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.file.Files;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import static com.yahoo.vespa.filedistribution.FileReferenceData.Type.compressed;
 import static org.junit.Assert.assertEquals;
@@ -65,6 +67,21 @@ public class FileReceiverTest {
         assertEquals("2", IOUtils.readFile(new File(downloadDir, "b")));
     }
 
+    @Test
+    public void receiveCompressedSingleFile() throws IOException{
+        File inputFile = temporaryFolder.newFile("a-file");
+        FileWriter writer = new FileWriter(inputFile);
+        String content = IntStream.range(1, 1000).mapToObj(a -> "surely this can be compressed").collect(Collectors.joining(","));
+        writer.write(content);
+        writer.close();
+
+        File tempFile = temporaryFolder.newFile();
+        File file = new FileReferenceCompressor(compressed).compress(inputFile, tempFile);
+        transferCompressedData(new FileReference("ref"), "a-file", IOUtils.readFileBytes(file));
+        File downloadDir = new File(root, "ref");
+        assertEquals(content, IOUtils.readFile(new File(downloadDir, "a-file")));
+    }
+
     private void transferPartsAndAssert(FileReference ref, String fileName, String all, int numParts) throws IOException {
         byte [] allContent = Utf8.toBytes(all);