Add ZipBuilder

2 files changed, 131 insertions, 0 deletions

diff --git a/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/ZipBuilder.java b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/ZipBuilder.java
new file mode 100644
index 00000000000..e3ede999c11
--- /dev/null
+++ b/controller-server/src/main/java/com/yahoo/vespa/hosted/controller/deployment/ZipBuilder.java
@@ -0,0 +1,65 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.deployment;
+
+import org.apache.commons.io.IOUtils;
+
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipInputStream;
+import java.util.zip.ZipOutputStream;
+
+/**
+ * Utility class to build zipped content by adding already zipped byte content or
+ * adding new unzipped entries.
+ *
+ * @author freva
+ */
+public class ZipBuilder implements AutoCloseable {
+
+    private final ByteArrayOutputStream byteArrayOutputStream;
+    private final ZipOutputStream zipOutputStream;
+
+    public ZipBuilder(int initialSize) {
+        byteArrayOutputStream = new ByteArrayOutputStream(initialSize);
+        zipOutputStream = new ZipOutputStream(byteArrayOutputStream);
+    }
+
+    public void add(byte[] zippedContent) {
+        try (ZipInputStream zin = new ZipInputStream(new ByteArrayInputStream(zippedContent))) {
+            for (ZipEntry entry = zin.getNextEntry(); entry != null; entry = zin.getNextEntry()) {
+                zipOutputStream.putNextEntry(entry);
+                IOUtils.copy(zin, zipOutputStream);
+                zipOutputStream.closeEntry();
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to add zipped content", e);
+        }
+    }
+
+    public void add(String entryName, byte[] content) {
+        try {
+            zipOutputStream.putNextEntry(new ZipEntry(entryName));
+            zipOutputStream.write(content);
+            zipOutputStream.closeEntry();
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to add entry " + entryName, e);
+        }
+    }
+
+    /** @return zipped byte array */
+    public byte[] toByteArray() {
+        return byteArrayOutputStream.toByteArray();
+    }
+
+    @Override
+    public void close() {
+        try {
+            zipOutputStream.close();
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to close zip output stream", e);
+        }
+    }
+}
diff --git a/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ZipBuilderTest.java b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ZipBuilderTest.java
new file mode 100644
index 00000000000..d54fd2124f0
--- /dev/null
+++ b/controller-server/src/test/java/com/yahoo/vespa/hosted/controller/deployment/ZipBuilderTest.java
@@ -0,0 +1,66 @@
+// Copyright 2018 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.controller.deployment;
+
+import org.apache.commons.io.IOUtils;
+import org.junit.Test;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.UncheckedIOException;
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.zip.ZipEntry;
+import java.util.zip.ZipInputStream;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author freva
+ */
+public class ZipBuilderTest {
+
+    @Test
+    public void test() {
+        Map<String, String> expected = new HashMap<>();
+        expected.put("dir/myfile", "my content");
+        expected.put("rootfile", "this is root");
+        expected.put("dir/newfile", "new file");
+        expected.put("dir/dir2/file", "nested file");
+
+        try (ZipBuilder zipBuilder1 = new ZipBuilder(100);
+             ZipBuilder zipBuilder2 = new ZipBuilder(1000)) {
+
+            // Add the entries to both zip builders one by one
+            Iterator<Map.Entry<String, String>> entries = expected.entrySet().iterator();
+            for (int i = 0; entries.hasNext(); i++) {
+                Map.Entry<String, String> entry = entries.next();
+                (i % 2 == 0 ? zipBuilder1 : zipBuilder2)
+                        .add(entry.getKey(), entry.getValue().getBytes(StandardCharsets.UTF_8));
+            }
+
+            // Add the zipped data from zip1 to zip2
+            zipBuilder2.add(zipBuilder1.toByteArray());
+
+            System.out.println(zipBuilder1.toByteArray().length);
+            System.out.println(zipBuilder2.toByteArray().length);
+            Map<String, String> actual = unzipToMap(zipBuilder2.toByteArray());
+
+            assertEquals(expected, actual);
+        }
+    }
+
+    Map<String, String> unzipToMap(byte[] zippedContent) {
+        Map<String, String> contents = new HashMap<>();
+        try (ZipInputStream zin = new ZipInputStream(new ByteArrayInputStream(zippedContent))) {
+            for (ZipEntry entry = zin.getNextEntry(); entry != null; entry = zin.getNextEntry()) {
+                if (entry.isDirectory()) continue;
+                contents.put(entry.getName(), IOUtils.toString(zin, StandardCharsets.UTF_8));
+            }
+        } catch (IOException e) {
+            throw new UncheckedIOException("Failed to read zipped content", e);
+        }
+        return contents;
+    }
+}
\ No newline at end of file