From 96bdc03ac0af705b43ae8f6756846d752f8b1533 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=A9=AC=E5=AD=90=E5=9D=A4?= <55695098+DanielWang2035@users.noreply.github.com> Date: Tue, 11 Jun 2024 14:32:46 +0800 Subject: [PATCH] init --- .../connector/compressor/PipeCompressor.java | 19 ++++++++++ .../compressor/PipeGZIPCompressor.java | 7 ++++ .../compressor/PipeLZ4Compressor.java | 7 ++++ .../compressor/PipeLZMA2Compressor.java | 7 ++++ .../compressor/PipeSnappyCompressor.java | 7 ++++ .../compressor/PipeZSTDCompressor.java | 5 +++ .../request/PipeTransferCompressedReq.java | 36 ++++++++++--------- 7 files changed, 72 insertions(+), 16 deletions(-) diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressor.java index 4d5113e4c20e0..a7f8405f7015a 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeCompressor.java @@ -49,8 +49,27 @@ protected PipeCompressor(PipeCompressionType compressionType) { public abstract byte[] compress(byte[] data) throws IOException; + /** + * Decompress the byte array to a byte array. NOTE: the length of the decompressed byte array is + * not provided in this method, and some decompressors (LZ4) may construct large byte arrays, + * leading to potential OOM. + * + * @param byteArray the byte array to be decompressed + * @return the decompressed byte array + * @throws IOException + */ public abstract byte[] decompress(byte[] byteArray) throws IOException; + /** + * Decompress the byte array to a byte array with a known length. + * + * @param byteArray the byte array to be decompressed + * @param decompressedLength the length of the decompressed byte array + * @return the decompressed byte array + * @throws IOException + */ + public abstract byte[] decompress(byte[] byteArray, int decompressedLength) throws IOException; + public byte serialize() { return compressionType.getIndex(); } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeGZIPCompressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeGZIPCompressor.java index 2cc1887555cd7..568122d7035b5 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeGZIPCompressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeGZIPCompressor.java @@ -44,4 +44,11 @@ public byte[] compress(byte[] data) throws IOException { public byte[] decompress(byte[] byteArray) throws IOException { return DECOMPRESSOR.uncompress(byteArray); } + + @Override + public byte[] decompress(byte[] byteArray, int decompressedLength) throws IOException { + byte[] uncompressed = new byte[decompressedLength]; + DECOMPRESSOR.uncompress(byteArray, 0, byteArray.length, uncompressed, 0); + return uncompressed; + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZ4Compressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZ4Compressor.java index 5d6db90e66100..52a2412a30f39 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZ4Compressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZ4Compressor.java @@ -44,4 +44,11 @@ public byte[] compress(byte[] data) throws IOException { public byte[] decompress(byte[] byteArray) throws IOException { return DECOMPRESSOR.uncompress(byteArray); } + + @Override + public byte[] decompress(byte[] byteArray, int decompressedLength) throws IOException { + byte[] uncompressed = new byte[decompressedLength]; + DECOMPRESSOR.uncompress(byteArray, 0, byteArray.length, uncompressed, 0); + return uncompressed; + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZMA2Compressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZMA2Compressor.java index 9e7c0a8b8d372..eebb703f9bed4 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZMA2Compressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeLZMA2Compressor.java @@ -44,4 +44,11 @@ public byte[] compress(byte[] data) throws IOException { public byte[] decompress(byte[] byteArray) throws IOException { return DECOMPRESSOR.uncompress(byteArray); } + + @Override + public byte[] decompress(byte[] byteArray, int decompressedLength) throws IOException { + byte[] uncompressed = new byte[decompressedLength]; + DECOMPRESSOR.uncompress(byteArray, 0, byteArray.length, uncompressed, 0); + return uncompressed; + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeSnappyCompressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeSnappyCompressor.java index 72333c1e929bd..284b29610a44a 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeSnappyCompressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeSnappyCompressor.java @@ -44,4 +44,11 @@ public byte[] compress(byte[] data) throws IOException { public byte[] decompress(byte[] byteArray) throws IOException { return DECOMPRESSOR.uncompress(byteArray); } + + @Override + public byte[] decompress(byte[] byteArray, int decompressedLength) throws IOException { + byte[] uncompressed = new byte[decompressedLength]; + DECOMPRESSOR.uncompress(byteArray, 0, byteArray.length, uncompressed, 0); + return uncompressed; + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java index 50e2e1f845cf7..f5e4d6476e844 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/compressor/PipeZSTDCompressor.java @@ -41,4 +41,9 @@ public byte[] compress(byte[] data) throws IOException { public byte[] decompress(byte[] byteArray) { return Zstd.decompress(byteArray, (int) Zstd.decompressedSize(byteArray, 0, byteArray.length)); } + + @Override + public byte[] decompress(byte[] byteArray, int decompressedLength) { + return Zstd.decompress(byteArray, decompressedLength); + } } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/payload/thrift/request/PipeTransferCompressedReq.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/payload/thrift/request/PipeTransferCompressedReq.java index f5e114949bd75..b92b2d4256a93 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/payload/thrift/request/PipeTransferCompressedReq.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/connector/payload/thrift/request/PipeTransferCompressedReq.java @@ -44,8 +44,10 @@ public static TPipeTransferReq toTPipeTransferReq( // version // type: TRANSFER_COMPRESSED // body: - // (byte) count of compressors - // (bytes) 1 byte for each compressor + // (byte) count of compressors (n) + // (n*3 bytes) for each compressor: + // (byte) compressor type + // (int) length of uncompressed bytes // compressed req: // (byte) version // (2 bytes) type @@ -56,18 +58,17 @@ public static TPipeTransferReq toTPipeTransferReq( try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { - ReadWriteIOUtils.write((byte) compressors.size(), outputStream); - for (final PipeCompressor compressor : compressors) { - ReadWriteIOUtils.write(compressor.serialize(), outputStream); - } - byte[] body = BytesUtils.concatByteArrayList( Arrays.asList( new byte[] {originalReq.version}, BytesUtils.shortToBytes(originalReq.type), originalReq.getBody())); + + ReadWriteIOUtils.write((byte) compressors.size(), outputStream); for (final PipeCompressor compressor : compressors) { + ReadWriteIOUtils.write(compressor.serialize(), outputStream); + ReadWriteIOUtils.write(body.length, outputStream); body = compressor.compress(body); } outputStream.write(body); @@ -84,17 +85,19 @@ public static TPipeTransferReq fromTPipeTransferReq(final TPipeTransferReq trans final ByteBuffer compressedBuffer = transferReq.body; final List compressors = new ArrayList<>(); + final List uncompressedLengths = new ArrayList<>(); final int compressorsSize = ReadWriteIOUtils.readByte(compressedBuffer); for (int i = 0; i < compressorsSize; ++i) { compressors.add( PipeCompressorFactory.getCompressor(ReadWriteIOUtils.readByte(compressedBuffer))); + uncompressedLengths.add(ReadWriteIOUtils.readInt(compressedBuffer)); } byte[] body = new byte[compressedBuffer.remaining()]; compressedBuffer.get(body); for (int i = compressors.size() - 1; i >= 0; --i) { - body = compressors.get(i).decompress(body); + body = compressors.get(i).decompress(body, uncompressedLengths.get(i)); } final ByteBuffer decompressedBuffer = ByteBuffer.wrap(body); @@ -115,26 +118,27 @@ public static byte[] toTPipeTransferReqBytes( // The generated bytes consists of: // (byte) version // (2 bytes) type: TRANSFER_COMPRESSED - // (byte) count of compressors - // (bytes) 1 byte for each compressor + // (byte) count of compressors (n) + // (n*3 bytes) for each compressor: + // (byte) compressor type + // (int) length of uncompressed bytes // compressed req: // (byte) version // (2 bytes) type // (bytes) body try (final PublicBAOS byteArrayOutputStream = new PublicBAOS(); final DataOutputStream outputStream = new DataOutputStream(byteArrayOutputStream)) { + byte[] body = rawReqInBytes; + ReadWriteIOUtils.write(IoTDBConnectorRequestVersion.VERSION_1.getVersion(), outputStream); ReadWriteIOUtils.write(PipeRequestType.TRANSFER_COMPRESSED.getType(), outputStream); ReadWriteIOUtils.write((byte) compressors.size(), outputStream); for (final PipeCompressor compressor : compressors) { ReadWriteIOUtils.write(compressor.serialize(), outputStream); + ReadWriteIOUtils.write(body.length, outputStream); + body = compressor.compress(body); } - - byte[] compressedReq = rawReqInBytes; - for (final PipeCompressor compressor : compressors) { - compressedReq = compressor.compress(compressedReq); - } - outputStream.write(compressedReq); + outputStream.write(body); return byteArrayOutputStream.toByteArray(); }