Support specifying compression type for file distribution requests

Add accepted compression types in file distribution requests.
Older clients (e.g. Vespa 7.x) will not have that in requests,
so fallback to gzip in that case.

8 files changed, 114 insertions, 40 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 3fe727c6f9d..f0c34e83713 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
@@ -29,6 +29,7 @@ import java.nio.file.Path;
 import java.time.Duration;
 import java.time.Instant;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ThreadPoolExecutor;
@@ -36,6 +37,9 @@ import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static com.yahoo.vespa.config.server.filedistribution.FileDistributionUtil.getOtherConfigServersInCluster;
+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.Type.compressed;
 
 public class FileServer {
@@ -114,17 +118,14 @@ public class FileServer {
 
     FileDirectory getRootDir() { return root; }
 
-    void startFileServing(FileReference fileReference, Receiver target) {
-        if (root.getFile(fileReference).exists())
-            serveFile(fileReference, target);
-    }
+    void startFileServing(FileReference reference, Receiver target, Set<CompressionType> acceptedCompressionTypes) {
+        if ( ! root.getFile(reference).exists()) return;
 
-    private void serveFile(FileReference reference, Receiver target) {
         File file = root.getFile(reference);
         log.log(Level.FINE, () -> "Start serving " + reference + " with file '" + file.getAbsolutePath() + "'");
         FileReferenceData fileData = EmptyFileReferenceData.empty(reference, file.getName());
         try {
-            fileData = readFileReferenceData(reference);
+            fileData = readFileReferenceData(reference, acceptedCompressionTypes);
             target.receive(fileData, new ReplayStatus(0, "OK"));
             log.log(Level.FINE, () -> "Done serving " + reference.value() + " with file '" + file.getAbsolutePath() + "'");
         } catch (IOException e) {
@@ -138,19 +139,23 @@ public class FileServer {
         }
     }
 
-    private FileReferenceData readFileReferenceData(FileReference reference) throws IOException {
+    private FileReferenceData readFileReferenceData(FileReference reference, Set<CompressionType> acceptedCompressionTypes) throws IOException {
         File file = root.getFile(reference);
 
         if (file.isDirectory()) {
             Path tempFile = Files.createTempFile("filereferencedata", reference.value());
-            File compressedFile = new FileReferenceCompressor(compressed).compress(file.getParentFile(), tempFile.toFile());
+            CompressionType compressionType = chooseCompressionType(acceptedCompressionTypes);
+            File compressedFile = new FileReferenceCompressor(compressed, compressionType).compress(file.getParentFile(), tempFile.toFile());
             return new LazyTemporaryStorageFileReferenceData(reference, file.getName(), compressed, compressedFile);
         } else {
             return new LazyFileReferenceData(reference, file.getName(), FileReferenceData.Type.file, file);
         }
     }
 
-    public void serveFile(FileReference fileReference, boolean downloadFromOtherSourceIfNotFound, Request request, Receiver receiver) {
+    public void serveFile(FileReference fileReference,
+                          boolean downloadFromOtherSourceIfNotFound,
+                          Set<CompressionType> acceptedCompressionTypes,
+                          Request request, Receiver receiver) {
         if (executor instanceof ThreadPoolExecutor)
             log.log(Level.FINE, () -> "Active threads: " + ((ThreadPoolExecutor) executor).getActiveCount());
 
@@ -158,7 +163,7 @@ public class FileServer {
         Instant deadline = Instant.now().plus(timeout);
         String client = request.target().toString();
         executor.execute(() -> {
-            var result = serveFileInternal(fileReference, downloadFromOtherSourceIfNotFound, client, receiver, deadline);
+            var result = serveFileInternal(fileReference, downloadFromOtherSourceIfNotFound, client, receiver, deadline, acceptedCompressionTypes);
             request.returnValues()
                    .add(new Int32Value(result.getCode()))
                    .add(new StringValue(result.getDescription()));
@@ -170,7 +175,8 @@ public class FileServer {
                                                 boolean downloadFromOtherSourceIfNotFound,
                                                 String client,
                                                 Receiver receiver,
-                                                Instant deadline) {
+                                                Instant deadline,
+                                                Set<CompressionType> acceptedCompressionTypes) {
         if (Instant.now().isAfter(deadline)) {
             log.log(Level.INFO, () -> "Deadline exceeded for request for file reference '" + fileReference + "' from " + client);
             return FileApiErrorCodes.TIMEOUT;
@@ -180,7 +186,7 @@ public class FileServer {
         try {
             var fileReferenceDownload = new FileReferenceDownload(fileReference, client, downloadFromOtherSourceIfNotFound);
             fileExists = hasFileDownloadIfNeeded(fileReferenceDownload);
-            if (fileExists) startFileServing(fileReference, receiver);
+            if (fileExists) startFileServing(fileReference, receiver, acceptedCompressionTypes);
         } catch (IllegalArgumentException e) {
             fileExists = false;
             log.warning("Failed serving file reference '" + fileReference + "', request from " + client + " failed with: " + e.getMessage());
@@ -189,6 +195,11 @@ public class FileServer {
         return (fileExists ? FileApiErrorCodes.OK : FileApiErrorCodes.NOT_FOUND);
     }
 
+    // TODO: Use lz4 for testing only, add zstd when we have support for (de)compressing zstd input and output streams
+    private CompressionType chooseCompressionType(Set<CompressionType> acceptedCompressionTypes) {
+        return acceptedCompressionTypes.contains(lz4) ? lz4 : gzip;
+    }
+
     boolean hasFileDownloadIfNeeded(FileReferenceDownload fileReferenceDownload) {
         FileReference fileReference = fileReferenceDownload.fileReference();
         if (hasFile(fileReference)) return true;
diff --git a/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java b/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java
index 6e919ae7f2e..687cb1d3cca 100644
--- a/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java
+++ b/configserver/src/main/java/com/yahoo/vespa/config/server/rpc/RpcServer.java
@@ -1,9 +1,9 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.config.server.rpc;
 
-import com.yahoo.component.annotation.Inject;
 import com.yahoo.cloud.config.ConfigserverConfig;
 import com.yahoo.component.Version;
+import com.yahoo.component.annotation.Inject;
 import com.yahoo.concurrent.ThreadFactoryFactory;
 import com.yahoo.config.FileReference;
 import com.yahoo.config.provision.ApplicationId;
@@ -44,13 +44,14 @@ import com.yahoo.vespa.filedistribution.FileDownloader;
 import com.yahoo.vespa.filedistribution.FileReceiver;
 import com.yahoo.vespa.filedistribution.FileReferenceData;
 import com.yahoo.vespa.filedistribution.FileReferenceDownload;
-
 import java.nio.ByteBuffer;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletionService;
 import java.util.concurrent.ConcurrentHashMap;
@@ -62,8 +63,11 @@ import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.logging.Level;
 import java.util.logging.Logger;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType;
+
 /**
  * An RPC server class that handles the config protocol RPC method "getConfigV3".
  * Mandatory hooks need to be implemented by subclasses.
@@ -221,7 +225,7 @@ public class RpcServer implements Runnable, ReloadListener, TenantListener {
         getSupervisor().addMethod(new Method("printStatistics", "", "s", this::printStatistics)
                                   .methodDesc("printStatistics")
                                   .returnDesc(0, "statistics", "Statistics for server"));
-        getSupervisor().addMethod(new Method("filedistribution.serveFile", "si", "is", this::serveFile));
+        getSupervisor().addMethod(new Method("filedistribution.serveFile", "si*", "is", this::serveFile));
         getSupervisor().addMethod(new Method("filedistribution.setFileReferencesToDownload", "S", "i", this::setFileReferencesToDownload)
                                      .methodDesc("set which file references to download")
                                      .paramDesc(0, "file references", "file reference to download")
@@ -566,10 +570,18 @@ public class RpcServer implements Runnable, ReloadListener, TenantListener {
         rpcAuthorizer.authorizeFileRequest(request)
                 .thenRun(() -> { // okay to do in authorizer thread as serveFile is async
                     FileServer.Receiver receiver = new ChunkedFileReceiver(request.target());
-                    fileServer.serveFile(new FileReference(request.parameters().get(0).asString()),
-                                         request.parameters().get(1).asInt32() == 0,
-                                         request,
-                                         receiver);
+
+                    FileReference reference = new FileReference(request.parameters().get(0).asString());
+                    boolean downloadFromOtherSourceIfNotFound = request.parameters().get(1).asInt32() == 0;
+                    Set<FileReferenceData.CompressionType> acceptedCompressionTypes = Set.of(CompressionType.gzip);
+                    // Newer clients specify accepted compression types in request
+                    if (request.parameters().size() > 2)
+                        acceptedCompressionTypes = Arrays.stream(request.parameters().get(2).asStringArray())
+                                                         .map(CompressionType::valueOf)
+                                                         .collect(Collectors.toSet());
+                    log.log(Level.FINE, "acceptedCompressionTypes=" + acceptedCompressionTypes);
+
+                    fileServer.serveFile(reference, downloadFromOtherSourceIfNotFound, acceptedCompressionTypes, request, receiver);
                 });
     }
 
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 861e6967620..9498db1d1e0 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
@@ -14,15 +14,17 @@ import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
-
 import java.io.File;
 import java.io.IOException;
 import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType.gzip;
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType.lz4;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
@@ -75,8 +77,13 @@ public class FileServerTest {
         File dir = getFileServerRootDir();
         IOUtils.writeFile(dir + "/12y/f1", "dummy-data", true);
         CompletableFuture<byte []> content = new CompletableFuture<>();
-        fileServer.startFileServing(new FileReference("12y"), new FileReceiver(content));
+        fileServer.startFileServing(new FileReference("12y"), new FileReceiver(content), Set.of(gzip));
         assertEquals(new String(content.get()), "dummy-data");
+
+        IOUtils.writeFile(dir + "/12z/f1", "dummy-data-2", true);
+        content = new CompletableFuture<>();
+        fileServer.startFileServing(new FileReference("12z"), new FileReceiver(content), Set.of(gzip, lz4));
+        assertEquals(new String(content.get()), "dummy-data-2");
     }
 
     @Test
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..65c6dd5931d 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReceiver.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReceiver.java
@@ -23,6 +23,7 @@ import java.util.logging.Level;
 import java.util.logging.Logger;
 
 import static com.yahoo.vespa.filedistribution.FileReferenceData.Type;
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType;
 
 /**
  * When asking for a file reference, this handles RPC callbacks from config server with file data and metadata.
@@ -48,7 +49,8 @@ public class FileReceiver {
         private final StreamingXXHash64 hasher;
         private final int sessionId;
         private final FileReference reference;
-        private final FileReferenceData.Type fileType;
+        private final Type fileType;
+        private final CompressionType compressionType;
         private final String fileName;
         private final long fileSize;
         private long currentFileSize;
@@ -58,13 +60,18 @@ public class FileReceiver {
         private final File tmpDir;
         private final File inprogressFile;
 
-        Session(File downloadDirectory, int sessionId, FileReference reference,
-                FileReferenceData.Type fileType, String fileName, long fileSize)
-        {
+        Session(File downloadDirectory,
+                int sessionId,
+                FileReference reference,
+                Type fileType,
+                FileReferenceData.CompressionType compressionType,
+                String fileName,
+                long fileSize) {
             this.hasher = XXHashFactory.fastestInstance().newStreamingHash64(0);
             this.sessionId = sessionId;
             this.reference = reference;
             this.fileType = fileType;
+            this.compressionType = compressionType;
             this.fileName = fileName;
             this.fileSize = fileSize;
             currentFileSize = 0;
@@ -122,7 +129,7 @@ public class FileReceiver {
                     moveFileToDestination(inprogressFile, file);
                 } else {
                     decompressedDir = Files.createTempDirectory(tmpDir.toPath(), "archive").toFile();
-                    new FileReferenceCompressor(fileType).decompress(inprogressFile, decompressedDir);
+                    new FileReferenceCompressor(fileType, compressionType).decompress(inprogressFile, decompressedDir);
                     moveFileToDestination(decompressedDir, fileReferenceDir);
                 }
             } catch (IOException e) {
@@ -161,11 +168,12 @@ public class FileReceiver {
     // receiveFile after getting a serveFile method call). handler needs to implement receiveFile* methods
     private List<Method> receiveFileMethod() {
         List<Method> methods = new ArrayList<>();
-        methods.add(new Method(RECEIVE_META_METHOD, "sssl", "ii", this::receiveFileMeta)
+        methods.add(new Method(RECEIVE_META_METHOD, "sssl*", "ii", this::receiveFileMeta)
                 .paramDesc(0, "filereference", "file reference to download")
                 .paramDesc(1, "filename", "filename")
                 .paramDesc(2, "type", "'file' or 'compressed'")
                 .paramDesc(3, "filelength", "length in bytes of file")
+                .paramDesc(3, "compressionType", "compression type: gzip, lz4, zstd")
                 .returnDesc(0, "ret", "0 if success, 1 otherwise")
                 .returnDesc(1, "session-id", "Session id to be used for this transfer"));
         methods.add(new Method(RECEIVE_PART_METHOD, "siix", "i", this::receiveFilePart)
@@ -220,8 +228,11 @@ public class FileReceiver {
         log.log(Level.FINE, () -> "Received method call '" + req.methodName() + "' with parameters : " + req.parameters());
         FileReference reference = new FileReference(req.parameters().get(0).asString());
         String fileName = req.parameters().get(1).asString();
-        String type = req.parameters().get(2).asString();
+        Type type = FileReferenceData.Type.valueOf(req.parameters().get(2).asString());
         long fileSize = req.parameters().get(3).asInt64();
+        CompressionType compressionType = (req.parameters().size() > 4)
+                ? CompressionType.valueOf(req.parameters().get(4).asString())
+                : CompressionType.gzip; // fallback/legacy compression type
         int sessionId = nextSessionId.getAndIncrement();
         int retval = 0;
         synchronized (sessions) {
@@ -231,7 +242,7 @@ public class FileReceiver {
             } else {
                 try {
                     sessions.put(sessionId, new Session(downloadDirectory, sessionId, reference,
-                                                        FileReferenceData.Type.valueOf(type),fileName, fileSize));
+                                                        type, compressionType, fileName, fileSize));
                 } catch (Exception e) {
                     retval = 1;
                 }
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..b485e6ded86 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceCompressor.java
@@ -2,6 +2,8 @@
 package com.yahoo.vespa.filedistribution;
 
 import com.google.common.io.ByteStreams;
+import net.jpountz.lz4.LZ4BlockInputStream;
+import net.jpountz.lz4.LZ4BlockOutputStream;
 import org.apache.commons.compress.archivers.ArchiveEntry;
 import org.apache.commons.compress.archivers.ArchiveInputStream;
 import org.apache.commons.compress.archivers.ArchiveOutputStream;
@@ -35,9 +37,11 @@ public class FileReferenceCompressor {
     private static final int recurseDepth = 100;
 
     private final FileReferenceData.Type type;
+    private final FileReferenceData.CompressionType compressionType;
 
-    public FileReferenceCompressor(FileReferenceData.Type type) {
+    public FileReferenceCompressor(FileReferenceData.Type type, FileReferenceData.CompressionType compressionType) {
         this.type = Objects.requireNonNull(type, "Type cannot be null");
+        this.compressionType = Objects.requireNonNull(compressionType, "Compression type cannot be null");
     }
 
     public File compress(File baseDir, List<File> inputFiles, File outputFile) throws IOException {
@@ -114,9 +118,17 @@ public class FileReferenceCompressor {
     }
 
     private OutputStream compressedOutputStream(File outputFile) throws IOException {
+        log.log(Level.FINE, () -> "Compressing with type " + type + " and compression type " + compressionType);
         switch (type) {
             case compressed:
-                return new GZIPOutputStream(new FileOutputStream(outputFile));
+                switch (compressionType) {
+                    case gzip:
+                        return new GZIPOutputStream(new FileOutputStream(outputFile));
+                    case lz4:
+                        return new LZ4BlockOutputStream(new FileOutputStream(outputFile));
+                    default:
+                        throw new RuntimeException("Unknown compression type " + compressionType);
+                }
             case file:
                 return new FileOutputStream(outputFile);
             default:
@@ -125,9 +137,17 @@ public class FileReferenceCompressor {
     }
 
     private InputStream decompressedInputStream(File inputFile) throws IOException {
+        log.log(Level.FINE, () -> "Decompressing with type " + type + " and compression type " + compressionType);
         switch (type) {
             case compressed:
-                return new GZIPInputStream(new FileInputStream(inputFile));
+                switch (compressionType) {
+                    case gzip:
+                        return new GZIPInputStream(new FileInputStream(inputFile));
+                    case lz4:
+                        return new LZ4BlockInputStream(new FileInputStream(inputFile));
+                    default:
+                        throw new RuntimeException("Unknown compression type " + compressionType);
+                }
             case file:
                 return new FileInputStream(inputFile);
             default:
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..d14f690b2d3 100644
--- a/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceData.java
+++ b/filedistribution/src/main/java/com/yahoo/vespa/filedistribution/FileReferenceData.java
@@ -13,6 +13,7 @@ import java.nio.ByteBuffer;
 public abstract class FileReferenceData {
 
     public enum Type { file, compressed }
+    public enum CompressionType { gzip, lz4, zstd }
 
     private final FileReference fileReference;
     private final String filename;
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 637fbbbd0a7..f5cd1760e89 100644
--- a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
+++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileDownloaderTest.java
@@ -32,6 +32,7 @@ import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 
 import static com.yahoo.jrt.ErrorCode.CONNECTION;
+import static com.yahoo.vespa.filedistribution.FileReferenceData.CompressionType.gzip;
 import static com.yahoo.vespa.filedistribution.FileReferenceData.Type.compressed;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -148,7 +149,7 @@ public class FileDownloaderTest {
             File barFile = new File(subdir, "really-long-filename-over-100-bytes-aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa");
             IOUtils.writeFile(barFile, "bar", false);
 
-            File tarFile = new FileReferenceCompressor(compressed).compress(tempPath.toFile(), Arrays.asList(fooFile, barFile), new File(tempPath.toFile(), filename));
+            File tarFile = new FileReferenceCompressor(compressed, gzip).compress(tempPath.toFile(), Arrays.asList(fooFile, barFile), new File(tempPath.toFile(), filename));
             byte[] tarredContent = IOUtils.readFileBytes(tarFile);
             receiveFile(fileReference, filename, compressed, tarredContent);
             Optional<File> downloadedFile = getFile(fileReference);
@@ -291,7 +292,7 @@ public class FileDownloaderTest {
                              FileReferenceData.Type type, byte[] content) {
         XXHash64 hasher = XXHashFactory.fastestInstance().hash64();
         FileReceiver.Session session =
-                new FileReceiver.Session(downloadDir, 1, fileReference, type, filename, content.length);
+                new FileReceiver.Session(downloadDir, 1, fileReference, type, gzip, filename, content.length);
         session.addPart(0, content);
         File file = session.close(hasher.hash(ByteBuffer.wrap(content), 0));
         fileDownloader.downloads().completedDownloading(fileReference, file);
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..84e7a07340e 100644
--- a/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java
+++ b/filedistribution/src/test/java/com/yahoo/vespa/filedistribution/FileReceiverTest.java
@@ -16,7 +16,11 @@ import java.io.IOException;
 import java.nio.ByteBuffer;
 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.Type.compressed;
+import static com.yahoo.vespa.filedistribution.FileReferenceData.Type.file;
 import static org.junit.Assert.assertEquals;
 
 public class FileReceiverTest {
@@ -58,18 +62,25 @@ public class FileReceiverTest {
         writerB.close();
 
         File tempFile = temporaryFolder.newFile();
-        File file = new FileReferenceCompressor(compressed).compress(dirWithFiles, tempFile);
-        transferCompressedData(new FileReference("ref"), "a", IOUtils.readFileBytes(file));
+        File file = new FileReferenceCompressor(compressed, gzip).compress(dirWithFiles, tempFile);
+        transferCompressedData(gzip, 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 {
         byte [] allContent = Utf8.toBytes(all);
 
-        FileReceiver.Session session = new FileReceiver.Session(root, 1, ref,
-                FileReferenceData.Type.file, fileName, allContent.length);
+        FileReceiver.Session session = new FileReceiver.Session(root, 1, ref, file, gzip, fileName, allContent.length);
         int partSize = (allContent.length+(numParts-1))/numParts;
         ByteBuffer bb = ByteBuffer.wrap(allContent);
         for (int i = 0, pos = 0; i < numParts; i++) {
@@ -87,8 +98,8 @@ public class FileReceiverTest {
         assertEquals(all, Utf8.toString(allReadBytes));
     }
 
-    private void transferCompressedData(FileReference ref, String fileName, byte[] data) {
-        FileReceiver.Session session = new FileReceiver.Session(root, 1, ref, compressed, fileName, data.length);
+    private void transferCompressedData(CompressionType compressionType, FileReference ref, String fileName, byte[] data) {
+        FileReceiver.Session session = new FileReceiver.Session(root, 1, ref, compressed, compressionType, fileName, data.length);
         session.addPart(0, data);
         session.close(hasher.hash(ByteBuffer.wrap(data), 0));
     }