ATLAS-2926: ZipSink OOM

repository/src/main/java/org/apache/atlas/repository/impexp/ZipSink.java

@@ -22,6 +22,7 @@ import org.apache.atlas.model.impexp.AtlasExportResult;
 import org.apache.atlas.model.instance.AtlasEntity;
 import org.apache.atlas.model.typedef.AtlasTypesDef;
 import org.apache.atlas.type.AtlasType;
+import org.apache.commons.lang.StringUtils;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -103,11 +104,31 @@ public class ZipSink {
 
         ZipEntry e = new ZipEntry(entryName);
         zipOutputStream.putNextEntry(e);
-
-        zipOutputStream.write(payload.getBytes());
+        writeBytes(payload);
         zipOutputStream.closeEntry();
     }
 
+    private void writeBytes(String payload) throws IOException {
+        splitAndWriteBytes(payload, 10 * 1024 * 1024, zipOutputStream);
+    }
+
+    static void splitAndWriteBytes(String msg, int bufferSize, OutputStream os) throws IOException {
+        int numberOfSplits = (int) Math.ceil(((float) msg.length()) / bufferSize);
+        if (numberOfSplits == 0) {
+            numberOfSplits = 1;
+        } else {
+            if (LOG.isDebugEnabled()) {
+                LOG.info("ZipSink: number of splits: {}", numberOfSplits);
+            }
+        }
+
+        for (int i = 0, start = 0; i < numberOfSplits; i++, start += bufferSize) {
+            int end = bufferSize + start;
+            String s = StringUtils.substring(msg, start, end);
+            os.write(s.getBytes());
+        }
+    }
+
     public boolean hasEntity(String guid) {
         return guids.contains(guid);
     }

repository/src/test/java/org/apache/atlas/repository/impexp/ZipSinkTest.java

@@ -30,12 +30,14 @@ import org.testng.annotations.Test;
 import java.io.ByteArrayInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.io.OutputStream;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 
+import static org.mockito.Mockito.mock;
 import static org.testng.Assert.*;
 
 public class ZipSinkTest {
@@ -45,6 +47,22 @@ public class ZipSinkTest {
     private AtlasExportResult defaultExportResult;
     private String knownEntityGuidFormat = "111-222-333-%s";
 
+    private class MockOutputStream extends OutputStream {
+        List<byte[]> collected = new ArrayList<>();
+
+        @Override
+        public void write(int b) throws IOException {
+        }
+
+        @Override
+        public void write(byte[] bytes) {
+            collected.add(bytes);
+        }
+
+        public List<byte[]> getCollected()  {
+            return collected;
+        }
+    };
 
     private void initZipSinkWithExportOrder() throws AtlasBaseException {
         zipSink = new ZipSink(byteArrayOutputStream);
@@ -207,4 +225,26 @@ public class ZipSinkTest {
         String json = AtlasType.toJson(defaultExportResult);
         return json.equals(s);
     }
+
+    @Test
+    public void splitTest() throws IOException {
+        assertSplit("ABCDEFGHIJKLMNOPQRSTUVWXYZ01", 7, new String[] {"ABCDEFG", "HIJKLMN", "OPQRSTU", "VWXYZ01"});
+        assertSplit("ABCDEFGHIJKLMNOPQRSTUVWXYZ", 7, new String[] {"ABCDEFG", "HIJKLMN", "OPQRSTU", "VWXYZ"});
+    }
+
+    private void assertSplit(String msg, int bufferSize, String[] splits) throws IOException {
+        MockOutputStream os = getOutputStream();
+        ZipSink.splitAndWriteBytes(msg, bufferSize, os);
+        assertEquals(os.getCollected().size(), splits.length);
+
+        for (int i = 0; i < os.collected.size(); i++) {
+            byte[] bytes = os.getCollected().get(i);
+            String s = new String(bytes);
+            assertEquals(s, splits[i]);
+        }
+    }
+
+    private MockOutputStream getOutputStream() {
+        return new MockOutputStream();
+    }
 }