Repository: atlas Updated Branches: refs/heads/master 7217f6781 -> 78e41ac25
ATLAS-2926: ZipSink OOM Project: http://git-wip-us.apache.org/repos/asf/atlas/repo Commit: http://git-wip-us.apache.org/repos/asf/atlas/commit/78e41ac2 Tree: http://git-wip-us.apache.org/repos/asf/atlas/tree/78e41ac2 Diff: http://git-wip-us.apache.org/repos/asf/atlas/diff/78e41ac2 Branch: refs/heads/master Commit: 78e41ac25d6dba65a52d61c43f4f98c5956d1e7d Parents: 7217f67 Author: Ashutosh Mestry <[email protected]> Authored: Thu Oct 18 22:20:21 2018 -0700 Committer: Ashutosh Mestry <[email protected]> Committed: Fri Oct 19 11:05:12 2018 -0700 ---------------------------------------------------------------------- .../apache/atlas/repository/impexp/ZipSink.java | 25 +++++++++++- .../atlas/repository/impexp/ZipSinkTest.java | 40 ++++++++++++++++++++ 2 files changed, 63 insertions(+), 2 deletions(-) ---------------------------------------------------------------------- http://git-wip-us.apache.org/repos/asf/atlas/blob/78e41ac2/repository/src/main/java/org/apache/atlas/repository/impexp/ZipSink.java ---------------------------------------------------------------------- diff --git a/repository/src/main/java/org/apache/atlas/repository/impexp/ZipSink.java b/repository/src/main/java/org/apache/atlas/repository/impexp/ZipSink.java index 9928c54..6375454 100644 --- a/repository/src/main/java/org/apache/atlas/repository/impexp/ZipSink.java +++ b/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); } http://git-wip-us.apache.org/repos/asf/atlas/blob/78e41ac2/repository/src/test/java/org/apache/atlas/repository/impexp/ZipSinkTest.java ---------------------------------------------------------------------- diff --git a/repository/src/test/java/org/apache/atlas/repository/impexp/ZipSinkTest.java b/repository/src/test/java/org/apache/atlas/repository/impexp/ZipSinkTest.java index 0988a30..cf6d16b 100644 --- a/repository/src/test/java/org/apache/atlas/repository/impexp/ZipSinkTest.java +++ b/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(); + } }
