Merge pull request #22744 from vespa-engine/hmusum/support-file-distribution-compress-single-file

Support compressing files (not just directories) in file distribution [run-systemtest]

3 files changed, 43 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(); }