From 8262835be0dee586fbe1862f62f02fd69e7b2f27 Mon Sep 17 00:00:00 2001 From: Harald Musum Date: Thu, 12 Jan 2023 13:09:28 +0100 Subject: Support zstd as compression type for file references --- .../filedistribution/FileReferenceCompressor.java | 21 ++++++++++++--------- .../vespa/filedistribution/FileReceiverTest.java | 19 +++++++++---------- 2 files changed, 21 insertions(+), 19 deletions(-) (limited to 'filedistribution') 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 2ea41a9fd8c..7b45c12bf61 100644 --- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java +++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java @@ -1,6 +1,8 @@ // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. package com.yahoo.vespa.filedistribution; +import ai.vespa.airlift.zstd.ZstdInputStream; +import com.yahoo.compress.ZstdOutputStream; import net.jpountz.lz4.LZ4BlockInputStream; import net.jpountz.lz4.LZ4BlockOutputStream; import org.apache.commons.compress.archivers.ArchiveEntry; @@ -21,7 +23,6 @@ import java.util.List; import java.util.Objects; import java.util.logging.Level; import java.util.logging.Logger; -import java.util.stream.Collectors; import java.util.zip.GZIPInputStream; import java.util.zip.GZIPOutputStream; @@ -118,32 +119,34 @@ public class FileReferenceCompressor { } private OutputStream compressedOutputStream(File outputFile) throws IOException { + FileOutputStream out = new FileOutputStream(outputFile); switch (type) { case compressed: log.log(Level.FINE, () -> "Compressing with compression type " + compressionType); return switch (compressionType) { - case gzip -> new GZIPOutputStream(new FileOutputStream(outputFile)); - case lz4 -> new LZ4BlockOutputStream(new FileOutputStream(outputFile)); - default -> throw new RuntimeException("Unknown compression type " + compressionType); + case gzip -> new GZIPOutputStream(out); + case lz4 -> new LZ4BlockOutputStream(out); + case zstd -> new ZstdOutputStream(out); }; case file: - return new FileOutputStream(outputFile); + return out; default: throw new RuntimeException("Unknown file reference type " + type); } } private InputStream decompressedInputStream(File inputFile) throws IOException { + FileInputStream in = new FileInputStream(inputFile); switch (type) { case compressed: log.log(Level.FINE, () -> "Decompressing with compression type " + compressionType); return switch (compressionType) { - case gzip -> new GZIPInputStream(new FileInputStream(inputFile)); - case lz4 -> new LZ4BlockInputStream(new FileInputStream(inputFile)); - default -> throw new RuntimeException("Unknown compression type " + compressionType); + case gzip -> new GZIPInputStream(in); + case lz4 -> new LZ4BlockInputStream(in); + case zstd -> new ZstdInputStream(in); }; case file: - return new FileInputStream(inputFile); + return in; default: throw new RuntimeException("Unknown file reference type " + type); } 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 84e7a07340e..59ebd571431 100644 --- a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java +++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java @@ -19,6 +19,7 @@ import java.nio.file.Files; import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType; import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType.gzip; import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType.lz4; +import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType.zstd; import static com.yahoo.vespa.filedistribution.FileReferenceData.Type.compressed; import static com.yahoo.vespa.filedistribution.FileReferenceData.Type.file; import static org.junit.Assert.assertEquals; @@ -61,20 +62,18 @@ public class FileReceiverTest { writerB.write("2"); writerB.close(); + testWithCompression(dirWithFiles, gzip); + testWithCompression(dirWithFiles, lz4); + testWithCompression(dirWithFiles, zstd); + } + + private void testWithCompression(File dirWithFiles, CompressionType compressionType) throws IOException { File tempFile = temporaryFolder.newFile(); - File file = new FileReferenceCompressor(compressed, gzip).compress(dirWithFiles, tempFile); - transferCompressedData(gzip, new FileReference("ref"), "a", IOUtils.readFileBytes(file)); + File file = new FileReferenceCompressor(compressed, compressionType).compress(dirWithFiles, tempFile); + transferCompressedData(compressionType, new FileReference("ref"), "a", IOUtils.readFileBytes(file)); File downloadDir = new File(root, "ref"); assertEquals("1", IOUtils.readFile(new File(downloadDir, "a"))); assertEquals("2", IOUtils.readFile(new File(downloadDir, "b"))); - - tempFile = temporaryFolder.newFile(); - FileReferenceCompressor compressor = new FileReferenceCompressor(compressed, lz4); - file = compressor.compress(dirWithFiles, tempFile); - transferCompressedData(lz4, new FileReference("ref"), "a", IOUtils.readFileBytes(file)); - downloadDir = new File(root, "ref"); - assertEquals("1", IOUtils.readFile(new File(downloadDir, "a"))); - assertEquals("2", IOUtils.readFile(new File(downloadDir, "b"))); } private void transferPartsAndAssert(FileReference ref, String fileName, String all, int numParts) throws IOException { -- cgit v1.2.3