handle zstd compressed files

2 files changed, 52 insertions, 19 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/handler/LogReader.java b/container-core/src/main/java/com/yahoo/container/handler/LogReader.java
index 2d60dc3f37b..93881b52eb6 100644
--- a/container-core/src/main/java/com/yahoo/container/handler/LogReader.java
+++ b/container-core/src/main/java/com/yahoo/container/handler/LogReader.java
@@ -43,8 +43,8 @@ import static java.nio.charset.StandardCharsets.UTF_8;
  * @author jonmv
  */
 class LogReader {
-    static final Pattern logArchivePathPattern = Pattern.compile("(\\d{4})/(\\d{2})/(\\d{2})/(\\d{2})-\\d+(.gz)?");
-    static final Pattern vespaLogPathPattern = Pattern.compile("vespa\\.log(?:-(\\d{4})-(\\d{2})-(\\d{2})\\.(\\d{2})-(\\d{2})-(\\d{2})(?:.gz)?)?");
+    static final Pattern logArchivePathPattern = Pattern.compile("(\\d{4})/(\\d{2})/(\\d{2})/(\\d{2})-\\d+(\\.gz|\\.zst)?");
+    static final Pattern vespaLogPathPattern = Pattern.compile("vespa\\.log(?:-(\\d{4})-(\\d{2})-(\\d{2})\\.(\\d{2})-(\\d{2})-(\\d{2})(?:\\.gz|\\.zst)?)?");
 
     private final Path logDirectory;
     private final Pattern logFilePattern;
@@ -97,23 +97,44 @@ class LogReader {
         private final double to;
         private final Optional<String> hostname;
         private LineWithTimestamp next;
+        private Process zcat = null;
 
-        private LogLineIterator(Path log, double from, double to, Optional<String> hostname) throws IOException {
-            boolean zipped = log.toString().endsWith(".gz");
-            InputStream in = InputStream.nullInputStream();
+        private InputStream openFile(Path log) {
+            boolean gzipped = log.toString().endsWith(".gz");
+            boolean is_zstd = log.toString().endsWith(".zst");
             try {
-                in = Files.newInputStream(log);
-            }
-            catch (NoSuchFileException e) { // File may have been compressed since we found it.
-                if ( ! zipped)
+                if (gzipped) {
+                    var in_gz = Files.newInputStream(log);
+                    return new GZIPInputStream(in_gz);
+                } else if (is_zstd) {
+                    var pb = new ProcessBuilder("zstdcat", log.toString());
+                    pb.redirectError(ProcessBuilder.Redirect.DISCARD);
+                    zcat = pb.start();
+                    zcat.getOutputStream().close();
+                    return zcat.getInputStream();
+                } else {
                     try {
-                        in = Files.newInputStream(Paths.get(log + ".gz"));
-                        zipped = true;
+                        return Files.newInputStream(log);
+                    } catch (NoSuchFileException e) { // File may have been compressed since we found it.
+                        Path p = Paths.get(log + ".gz");
+                        if (Files.exists(p)) {
+                            return openFile(p);
+                        }
+                        p = Paths.get(log + ".zst");
+                        if (Files.exists(p)) {
+                            return openFile(p);
+                        }
                     }
-                    catch (NoSuchFileException ignored) { }
+                }
+            } catch (IOException ignored) {
             }
+            // failure fallback:
+            return InputStream.nullInputStream();
+        }
 
-            this.reader = new BufferedReader(new InputStreamReader(zipped ? new GZIPInputStream(in) : in, UTF_8));
+        private LogLineIterator(Path log, double from, double to, Optional<String> hostname) throws IOException {
+            InputStream in = openFile(log);
+            this.reader = new BufferedReader(new InputStreamReader(in, UTF_8));
             this.from = from;
             this.to = to;
             this.hostname = hostname;
@@ -135,6 +156,9 @@ class LogReader {
         @Override
         public void close() throws IOException {
             reader.close();
+            if (zcat != null) {
+                zcat.destroy();
+            }
         }
 
         private LineWithTimestamp readNext() {
diff --git a/container-core/src/test/java/com/yahoo/container/handler/LogReaderTest.java b/container-core/src/test/java/com/yahoo/container/handler/LogReaderTest.java
index 120ee49031d..e1bd0da457e 100644
--- a/container-core/src/test/java/com/yahoo/container/handler/LogReaderTest.java
+++ b/container-core/src/test/java/com/yahoo/container/handler/LogReaderTest.java
@@ -1,9 +1,11 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.container.handler;
 
-import com.yahoo.vespa.test.file.TestFileSystem;
+import com.yahoo.compress.ZstdCompressor;
 import org.junit.Before;
 import org.junit.Test;
+import org.junit.Rule;
+import org.junit.rules.TemporaryFolder;
 
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
@@ -22,8 +24,9 @@ import static org.junit.Assert.assertEquals;
 
 public class LogReaderTest {
 
-    private final FileSystem fileSystem = TestFileSystem.create();
-    private final Path logDirectory = fileSystem.getPath("/opt/vespa/logs");
+    @Rule
+    public TemporaryFolder folder = new TemporaryFolder();
+    private Path logDirectory;
 
     private static final String logv11 = "3600.2\tnode1.com\t5480\tcontainer\tstdout\tinfo\tfourth\n";
     private static final String logv   = "90000.1\tnode1.com\t5480\tcontainer\tstdout\tinfo\tlast\n";
@@ -34,12 +37,13 @@ public class LogReaderTest {
 
     @Before
     public void setup() throws IOException {
+        logDirectory = folder.newFolder("opt/vespa/logs").toPath();
         // Log archive paths and file names indicate what hour they contain logs for, with the start of that hour.
         // Multiple entries may exist for each hour.
         Files.createDirectories(logDirectory.resolve("1970/01/01"));
-        Files.write(logDirectory.resolve("1970/01/01/00-0.gz"), compress(log100));
+        Files.write(logDirectory.resolve("1970/01/01/00-0.gz"), compress1(log100));
         Files.write(logDirectory.resolve("1970/01/01/00-1"), log101.getBytes(UTF_8));
-        Files.write(logDirectory.resolve("1970/01/01/01-0.gz"), compress(log110));
+        Files.write(logDirectory.resolve("1970/01/01/01-0.zst"), compress2(log110));
 
         Files.createDirectories(logDirectory.resolve("1970/01/02"));
         Files.write(logDirectory.resolve("1970/01/02/00-0"), log200.getBytes(UTF_8));
@@ -86,7 +90,7 @@ public class LogReaderTest {
         assertEquals(log101 + log100 + log200, baos.toString(UTF_8));
     }
 
-    private byte[] compress(String input) throws IOException {
+    private byte[] compress1(String input) throws IOException {
         ByteArrayOutputStream baos = new ByteArrayOutputStream();
         OutputStream zip = new GZIPOutputStream(baos);
         zip.write(input.getBytes());
@@ -94,4 +98,9 @@ public class LogReaderTest {
         return baos.toByteArray();
     }
 
+    private byte[] compress2(String input) throws IOException {
+        byte[] data = input.getBytes();
+        return new ZstdCompressor().compress(data, 0, data.length);
+    }
+
 }