This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch dev/1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/dev/1.3 by this push:
new 4df9f9961a9 Add pipe receiver memory protection (#18090) (#18133)
4df9f9961a9 is described below
commit 4df9f9961a9600a2e146d6866ec232f18625f1a9
Author: Caideyipi <[email protected]>
AuthorDate: Fri Jul 10 10:20:43 2026 +0800
Add pipe receiver memory protection (#18090) (#18133)
* Add pipe receiver memory protection
* Add pipe receiver memory protection tests
* Address receiver memory review comments
(cherry picked from commit 71696f333af3f71e2bea946754436b90247b0d4e)
---
.../protocol/thrift/IoTDBDataNodeReceiver.java | 117 +++++++++++++---
.../commons/pipe/receiver/IoTDBFileReceiver.java | 39 +++++-
.../thrift/common/PipeTransferSliceReqHandler.java | 37 +++++
.../thrift/request/PipeTransferCompressedReq.java | 21 +++
.../pipe/receiver/IoTDBFileReceiverTest.java | 89 ++++++++++++
.../common/PipeTransferSliceReqBuilderTest.java | 53 ++++++++
.../request/PipeTransferCompressedReqTest.java | 149 +++++++++++++++++++++
7 files changed, 489 insertions(+), 16 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java
index 761c0183c9b..945b4a55209 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java
@@ -151,6 +151,7 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
private static final SessionManager SESSION_MANAGER =
SessionManager.getInstance();
private PipeMemoryBlock allocatedMemoryBlock;
+ private final List<PipeMemoryBlock> allocatedSliceMemoryBlocks = new
ArrayList<>();
static {
try {
@@ -172,7 +173,7 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
if (PipeRequestType.isValidatedRequestType(rawRequestType)) {
final PipeRequestType requestType =
PipeRequestType.valueOf(rawRequestType);
if (requestType != PipeRequestType.TRANSFER_SLICE) {
- sliceReqHandler.clear();
+ clearSliceReqHandler();
}
final TPipeTransferResp authResp =
checkPipeTransferAuthenticated(requestType);
if (Objects.nonNull(authResp)) {
@@ -350,8 +351,18 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
}
case TRANSFER_COMPRESSED:
{
+ long requestedMemorySizeInBytes = 0;
try {
- return
receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
+ requestedMemorySizeInBytes =
+
PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(req);
+ try (final PipeMemoryBlock ignored =
+ tryAllocateReceiverMemory(requestedMemorySizeInBytes)) {
+ return
receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
+ }
+ } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+ return new TPipeTransferResp(
+ getReceiverTemporaryUnavailableStatus(
+ "decompressing pipe transfer request",
requestedMemorySizeInBytes, e));
} finally {
PipeDataNodeReceiverMetrics.getInstance()
.recordTransferCompressedTimer(System.nanoTime() -
startTime);
@@ -698,22 +709,97 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
}
private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq
pipeTransferSliceReq) {
- final boolean isInorder =
sliceReqHandler.receiveSlice(pipeTransferSliceReq);
- if (!isInorder) {
+ final long sliceBodySizeInBytes =
getSliceBodySizeInBytes(pipeTransferSliceReq);
+ long requestedMemorySizeInBytes = sliceBodySizeInBytes;
+ String memoryAction = "buffering sliced pipe transfer request";
+ PipeMemoryBlock sliceMemoryBlock = null;
+ try {
+ sliceMemoryBlock = tryAllocateReceiverMemory(sliceBodySizeInBytes);
+
+ final boolean isInorder =
sliceReqHandler.receiveSlice(pipeTransferSliceReq);
+ if (!isInorder) {
+ closeMemoryBlock(sliceMemoryBlock);
+ clearSliceReqHandler();
+ return new TPipeTransferResp(
+ RpcUtils.getStatus(
+ TSStatusCode.PIPE_TRANSFER_SLICE_OUT_OF_ORDER,
+ "Slice request is out of order, please check the request
sequence."));
+ }
+
+ allocatedSliceMemoryBlocks.add(sliceMemoryBlock);
+ sliceMemoryBlock = null;
+
+ if (pipeTransferSliceReq.getSliceIndex() + 1 <
pipeTransferSliceReq.getSliceCount()) {
+ return new TPipeTransferResp(
+ RpcUtils.getStatus(
+ TSStatusCode.SUCCESS_STATUS,
+ "Slice received, waiting for more slices to complete the
request."));
+ }
+
+ memoryAction = "assembling sliced pipe transfer request";
+ requestedMemorySizeInBytes = pipeTransferSliceReq.getOriginBodySize();
+ try (final PipeMemoryBlock ignored =
tryAllocateReceiverMemory(requestedMemorySizeInBytes)) {
+ final Optional<TPipeTransferReq> req =
sliceReqHandler.makeReqIfComplete();
+ if (!req.isPresent()) {
+ return new TPipeTransferResp(
+ RpcUtils.getStatus(
+ TSStatusCode.SUCCESS_STATUS,
+ "Slice received, waiting for more slices to complete the
request."));
+ }
+ clearSliceReqHandler();
+ return receive(req.get());
+ }
+ } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+ closeMemoryBlock(sliceMemoryBlock);
+ clearSliceReqHandler();
return new TPipeTransferResp(
- RpcUtils.getStatus(
- TSStatusCode.PIPE_TRANSFER_SLICE_OUT_OF_ORDER,
- "Slice request is out of order, please check the request
sequence."));
+ getReceiverTemporaryUnavailableStatus(memoryAction,
requestedMemorySizeInBytes, e));
+ } catch (final RuntimeException e) {
+ closeMemoryBlock(sliceMemoryBlock);
+ clearSliceReqHandler();
+ throw e;
}
- final Optional<TPipeTransferReq> req = sliceReqHandler.makeReqIfComplete();
- if (!req.isPresent()) {
- return new TPipeTransferResp(
- RpcUtils.getStatus(
- TSStatusCode.SUCCESS_STATUS,
- "Slice received, waiting for more slices to complete the
request."));
+ }
+
+ private long getSliceBodySizeInBytes(final PipeTransferSliceReq
pipeTransferSliceReq) {
+ return pipeTransferSliceReq.getSliceBody() == null
+ ? 0
+ : pipeTransferSliceReq.getSliceBody().length;
+ }
+
+ private void clearSliceReqHandler() {
+ sliceReqHandler.clear();
+ allocatedSliceMemoryBlocks.forEach(this::closeMemoryBlock);
+ allocatedSliceMemoryBlocks.clear();
+ }
+
+ private void closeMemoryBlock(final PipeMemoryBlock memoryBlock) {
+ if (Objects.nonNull(memoryBlock)) {
+ memoryBlock.close();
}
- // sliceReqHandler will be cleared in the receive(req) method
- return receive(req.get());
+ }
+
+ private PipeMemoryBlock tryAllocateReceiverMemory(final long
requestedMemorySizeInBytes)
+ throws PipeRuntimeOutOfMemoryCriticalException {
+ return PipeDataNodeResourceManager.memory()
+ .forceAllocate(Math.max(requestedMemorySizeInBytes, 0));
+ }
+
+ @Override
+ protected TSStatus getReceiverTemporaryUnavailableStatus(
+ final String action,
+ final long requestedMemorySizeInBytes,
+ final PipeRuntimeOutOfMemoryCriticalException e) {
+ return new
TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode())
+ .setMessage(
+ String.format(
+ "Temporarily out of memory when %s. Requested memory: %d
bytes, used memory: %d "
+ + "bytes, free memory: %d bytes, total non-floating
memory: %d bytes",
+ action,
+ requestedMemorySizeInBytes,
+
PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(),
+
PipeDataNodeResourceManager.memory().getFreeMemorySizeInBytes(),
+
PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes()));
}
/**
@@ -945,6 +1031,7 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
@Override
public synchronized void handleExit() {
+ clearSliceReqHandler();
if (Objects.nonNull(configReceiverId.get())) {
try {
ClusterConfigTaskExecutor.getInstance().handlePipeConfigClientExit(configReceiverId.get());
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java
index 2375f0d4072..3223305dc80 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java
@@ -23,6 +23,7 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.audit.UserEntity;
import org.apache.iotdb.commons.conf.CommonDescriptor;
import org.apache.iotdb.commons.exception.IllegalPathException;
+import
org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
import org.apache.iotdb.commons.pipe.config.PipeConfig;
import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
import org.apache.iotdb.commons.pipe.resource.log.PipeLogger;
@@ -425,7 +426,7 @@ public abstract class IoTDBFileReceiver implements
IoTDBReceiver {
final PipeTransferFilePieceReq req,
final boolean isRequestThroughAirGap,
final boolean isSingleFile) {
- try {
+ try (final AutoCloseable ignored = tryAllocateMemoryForFilePiece(req)) {
updateWritingFileIfNeeded(req.getFileName(), isSingleFile);
// If the request is through air gap, the sender will resend the file
piece from the beginning
@@ -460,6 +461,22 @@ public abstract class IoTDBFileReceiver implements
IoTDBReceiver {
writingFileWriter.write(req.getFilePiece());
return PipeTransferFilePieceResp.toTPipeTransferResp(
RpcUtils.SUCCESS_STATUS, writingFileWriter.length());
+ } catch (final PipeRuntimeOutOfMemoryCriticalException e) {
+ final TSStatus status =
+ getReceiverTemporaryUnavailableStatus(
+ "receiving pipe file piece", getFilePieceSizeInBytes(req), e);
+ PipeLogger.log(
+ LOGGER::warn,
+ e,
+ "Receiver id = %s: Failed to write file piece from req %s.",
+ receiverId.get(),
+ req);
+ try {
+ return PipeTransferFilePieceResp.toTPipeTransferResp(
+ status, PipeTransferFilePieceResp.ERROR_END_OFFSET);
+ } catch (Exception ex) {
+ return PipeTransferFilePieceResp.toTPipeTransferResp(status);
+ }
} catch (final Exception e) {
PipeLogger.log(
LOGGER::warn,
@@ -480,6 +497,26 @@ public abstract class IoTDBFileReceiver implements
IoTDBReceiver {
}
}
+ protected AutoCloseable tryAllocateMemoryForFilePiece(final
PipeTransferFilePieceReq req)
+ throws PipeRuntimeOutOfMemoryCriticalException {
+ return () -> {};
+ }
+
+ protected TSStatus getReceiverTemporaryUnavailableStatus(
+ final String action,
+ final long requestedMemorySizeInBytes,
+ final PipeRuntimeOutOfMemoryCriticalException e) {
+ return new
TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode())
+ .setMessage(
+ String.format(
+ "Temporarily out of memory when %s. Requested memory: %d
bytes. Root cause: %s",
+ action, requestedMemorySizeInBytes, e.getMessage()));
+ }
+
+ private static long getFilePieceSizeInBytes(final PipeTransferFilePieceReq
req) {
+ return req.getFilePiece() == null ? 0 : req.getFilePiece().length;
+ }
+
protected final void updateWritingFileIfNeeded(final String fileName, final
boolean isSingleFile)
throws IOException {
if (isFileExistedAndNameCorrect(fileName)) {
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java
index 09edb4bbe4a..ead9ea277d6 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java
@@ -42,8 +42,14 @@ public class PipeTransferSliceReqHandler {
private int sliceCount = -1;
private final List<byte[]> sliceBodies = new ArrayList<>();
+ private long receivedBodySize = 0;
public boolean receiveSlice(final PipeTransferSliceReq req) {
+ if (!isValidSliceReq(req)) {
+ clear();
+ return false;
+ }
+
if (orderId == -1
|| originReqType == -1
|| originBodySize == -1
@@ -58,6 +64,7 @@ public class PipeTransferSliceReqHandler {
originReqType = req.getOriginReqType();
originBodySize = req.getOriginBodySize();
sliceCount = req.getSliceCount();
+ receivedBodySize = 0;
} else {
LOGGER.warn(
"Invalid state: orderId={}, originReqType={}, originBodySize={},
sliceCount={}, sliceBodies.size={}",
@@ -104,10 +111,39 @@ public class PipeTransferSliceReqHandler {
return false;
}
+ if (receivedBodySize + req.getSliceBody().length > originBodySize) {
+ LOGGER.warn(
+ "Received slice body size {} exceeds origin body size {}",
+ receivedBodySize + req.getSliceBody().length,
+ originBodySize);
+ clear();
+ return false;
+ }
+
sliceBodies.add(req.getSliceBody());
+ receivedBodySize += req.getSliceBody().length;
+
+ if (sliceBodies.size() == sliceCount && receivedBodySize !=
originBodySize) {
+ LOGGER.warn(
+ "Received slice body size {} is not equal to origin body size {}",
+ receivedBodySize,
+ originBodySize);
+ clear();
+ return false;
+ }
+
return true;
}
+ private boolean isValidSliceReq(final PipeTransferSliceReq req) {
+ return req.getOriginBodySize() >= 0
+ && req.getSliceCount() > 0
+ && req.getSliceIndex() >= 0
+ && req.getSliceIndex() < req.getSliceCount()
+ && req.getSliceBody() != null
+ && req.getSliceBody().length <= req.getOriginBodySize();
+ }
+
public Optional<TPipeTransferReq> makeReqIfComplete() {
if (sliceBodies.size() != sliceCount) {
return Optional.empty();
@@ -131,5 +167,6 @@ public class PipeTransferSliceReqHandler {
originBodySize = -1;
sliceCount = -1;
sliceBodies.clear();
+ receivedBodySize = 0;
}
}
diff --git
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java
index 73bac4c05b4..7aa639140be 100644
---
a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java
+++
b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java
@@ -112,6 +112,27 @@ public class PipeTransferCompressedReq extends
TPipeTransferReq {
return decompressedReq;
}
+ /** Get the largest intermediate decompressed body size without consuming
the request body. */
+ public static int getMaxDecompressedLengthInBytes(final TPipeTransferReq
transferReq) {
+ final ByteBuffer compressedBuffer = transferReq.body.duplicate();
+
+ int maxDecompressedLength = compressedBuffer.remaining();
+ final int compressorsSize = ReadWriteIOUtils.readByte(compressedBuffer);
+ for (int i = 0; i < compressorsSize; ++i) {
+ ReadWriteIOUtils.readByte(compressedBuffer);
+ final int decompressedLength =
ReadWriteIOUtils.readInt(compressedBuffer);
+ checkDecompressedLength(decompressedLength);
+ maxDecompressedLength = Math.max(maxDecompressedLength,
decompressedLength);
+ }
+ return maxDecompressedLength;
+ }
+
+ /** Get the largest additional decompressed body size beyond the current
transfer frame. */
+ public static int getMaxAdditionalDecompressedLengthInBytes(final
TPipeTransferReq transferReq) {
+ final int transferFrameBodySize = transferReq.body.duplicate().remaining();
+ return Math.max(0, getMaxDecompressedLengthInBytes(transferReq) -
transferFrameBodySize);
+ }
+
/** This method is used to prevent decompression bomb attacks. */
private static void checkDecompressedLength(final int decompressedLength)
throws IllegalArgumentException {
diff --git
a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java
index b0180e5b124..a03f8b8e61a 100644
---
a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java
+++
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiverTest.java
@@ -22,12 +22,15 @@ package org.apache.iotdb.commons.pipe.receiver;
import org.apache.iotdb.common.rpc.thrift.TSStatus;
import org.apache.iotdb.commons.conf.CommonDescriptor;
import org.apache.iotdb.commons.exception.IllegalPathException;
+import
org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType;
+import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFilePieceReq;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV1;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferFileSealReqV2;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV1Req;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV2Req;
+import
org.apache.iotdb.commons.pipe.sink.payload.thrift.response.PipeTransferFilePieceResp;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp;
@@ -200,10 +203,55 @@ public class IoTDBFileReceiverTest {
}
}
+ @Test
+ public void
testFilePieceMemoryAllocationFailureReturnsTemporaryUnavailable() throws
Exception {
+ final Path baseDir = Files.createTempDirectory("iotdb-file-receiver-test");
+ final DummyFileReceiver receiver = new DummyFileReceiver(baseDir.toFile());
+ try {
+ receiver.setFailFilePieceMemoryAllocation(true);
+
+ final TPipeTransferResp response =
+ receiver.writeFilePiece("normal.tsfile", 0, new byte[] {1, 2, 3});
+ final PipeTransferFilePieceResp filePieceResp =
+ PipeTransferFilePieceResp.fromTPipeTransferResp(response);
+
+ Assert.assertEquals(
+
TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode(),
+ response.getStatus().getCode());
+ Assert.assertTrue(response.getStatus().getMessage().contains("no memory
for file piece"));
+ Assert.assertEquals(
+ PipeTransferFilePieceResp.ERROR_END_OFFSET,
filePieceResp.getEndWritingOffset());
+
Assert.assertFalse(receiver.getWritingFileInBaseDir("normal.tsfile").exists());
+ } finally {
+ receiver.handleExit();
+ }
+ }
+
+ @Test
+ public void testFilePieceMemoryAllocationIsClosedAfterWrite() throws
Exception {
+ final Path baseDir = Files.createTempDirectory("iotdb-file-receiver-test");
+ final DummyFileReceiver receiver = new DummyFileReceiver(baseDir.toFile());
+ try {
+ final TPipeTransferResp response =
+ receiver.writeFilePiece("normal.tsfile", 0, new byte[] {1, 2, 3});
+ final PipeTransferFilePieceResp filePieceResp =
+ PipeTransferFilePieceResp.fromTPipeTransferResp(response);
+
+ Assert.assertEquals(
+ TSStatusCode.SUCCESS_STATUS.getStatusCode(),
response.getStatus().getCode());
+ Assert.assertEquals(3, filePieceResp.getEndWritingOffset());
+ Assert.assertEquals(1, receiver.getFilePieceMemoryCloseCount());
+ } finally {
+ receiver.handleExit();
+ }
+ }
+
private static class DummyFileReceiver extends IoTDBFileReceiver {
private final File receiverFileBaseDir;
private TSStatus loadFileV1Status = new
TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
+ private boolean failFilePieceMemoryAllocation = false;
+ private int filePieceMemoryCloseCount = 0;
private int loginCallCount = 0;
DummyFileReceiver(final File baseDir) {
@@ -233,6 +281,23 @@ public class IoTDBFileReceiverTest {
loadFileV1Status = status;
}
+ void setFailFilePieceMemoryAllocation(final boolean
failFilePieceMemoryAllocation) {
+ this.failFilePieceMemoryAllocation = failFilePieceMemoryAllocation;
+ }
+
+ int getFilePieceMemoryCloseCount() {
+ return filePieceMemoryCloseCount;
+ }
+
+ TPipeTransferResp writeFilePiece(
+ final String fileName, final long startWritingOffset, final byte[]
filePiece)
+ throws IOException {
+ return handleTransferFilePiece(
+ DummyFilePieceReq.toTPipeTransferReq(fileName, startWritingOffset,
filePiece),
+ false,
+ true);
+ }
+
void setHasPipeHandshakeCredential(final boolean
hasPipeHandshakeCredential) {
this.hasPipeHandshakeCredential = hasPipeHandshakeCredential;
}
@@ -303,6 +368,14 @@ public class IoTDBFileReceiverTest {
return new TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode());
}
+ @Override
+ protected AutoCloseable tryAllocateMemoryForFilePiece(final
PipeTransferFilePieceReq req) {
+ if (failFilePieceMemoryAllocation) {
+ throw new PipeRuntimeOutOfMemoryCriticalException("no memory for file
piece");
+ }
+ return () -> filePieceMemoryCloseCount++;
+ }
+
@Override
protected TSStatus loadFileV1(
final PipeTransferFileSealReqV1 req, final String fileAbsolutePath) {
@@ -327,6 +400,22 @@ public class IoTDBFileReceiverTest {
}
}
+ private static class DummyFilePieceReq extends PipeTransferFilePieceReq {
+
+ static DummyFilePieceReq toTPipeTransferReq(
+ final String fileName, final long startWritingOffset, final byte[]
filePiece)
+ throws IOException {
+ return (DummyFilePieceReq)
+ new DummyFilePieceReq()
+ .convertToTPipeTransferReq(fileName, startWritingOffset,
filePiece);
+ }
+
+ @Override
+ protected PipeRequestType getPlanType() {
+ return PipeRequestType.TRANSFER_TS_FILE_PIECE;
+ }
+ }
+
private static class DummyHandshakeReq extends PipeTransferHandshakeV1Req {
static DummyHandshakeReq toTPipeTransferReq(final String
timestampPrecision)
diff --git
a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java
index 290ce397980..857a62ac75c 100644
---
a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java
+++
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqBuilderTest.java
@@ -30,7 +30,9 @@ import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
+import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Optional;
public class PipeTransferSliceReqBuilderTest {
@@ -91,6 +93,57 @@ public class PipeTransferSliceReqBuilderTest {
createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 4),
bodySizeLimit));
}
+ @Test
+ public void testSliceReqHandlerRejectsOversizedSlices() throws IOException {
+ final TPipeTransferReq req =
createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 6);
+ final PipeTransferSliceReq firstSlice =
+ PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 0, 2,
req.body.duplicate(), 0, 4);
+ final PipeTransferSliceReq oversizedSecondSlice =
+ PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 1, 2,
req.body.duplicate(), 2, 6);
+
+ final PipeTransferSliceReqHandler handler = new
PipeTransferSliceReqHandler();
+ Assert.assertTrue(handler.receiveSlice(firstSlice));
+ Assert.assertFalse(handler.receiveSlice(oversizedSecondSlice));
+ Assert.assertFalse(handler.makeReqIfComplete().isPresent());
+ }
+
+ @Test
+ public void testSliceReqHandlerAssemblesCompleteRequest() throws IOException
{
+ final TPipeTransferReq req =
createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 6);
+ final PipeTransferSliceReq firstSlice =
+ PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 0, 2,
req.body.duplicate(), 0, 4);
+ final PipeTransferSliceReq secondSlice =
+ PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 1, 2,
req.body.duplicate(), 4, 6);
+
+ final PipeTransferSliceReqHandler handler = new
PipeTransferSliceReqHandler();
+ Assert.assertTrue(handler.receiveSlice(firstSlice));
+ Assert.assertFalse(handler.makeReqIfComplete().isPresent());
+ Assert.assertTrue(handler.receiveSlice(secondSlice));
+
+ final Optional<TPipeTransferReq> completedReq =
handler.makeReqIfComplete();
+ Assert.assertTrue(completedReq.isPresent());
+
+ final TPipeTransferReq assembledReq = completedReq.get();
+ Assert.assertEquals(IoTDBSinkRequestVersion.VERSION_1.getVersion(),
assembledReq.getVersion());
+ Assert.assertEquals(req.getType(), assembledReq.getType());
+
+ final byte[] assembledBody = new byte[assembledReq.body.remaining()];
+ assembledReq.body.get(assembledBody);
+ Assert.assertArrayEquals(new byte[] {0, 1, 2, 3, 4, 5}, assembledBody);
+ }
+
+ @Test
+ public void testSliceReqHandlerRejectsInvalidMetadata() throws IOException {
+ final TPipeTransferReq req =
createReq(IoTDBSinkRequestVersion.VERSION_1.getVersion(), 2);
+ final PipeTransferSliceReq invalidSlice =
+ PipeTransferSliceReq.toTPipeTransferReq(7, req.getType(), 0, 0,
req.body.duplicate(), 0, 1);
+
+ final PipeTransferSliceReqHandler handler = new
PipeTransferSliceReqHandler();
+
+ Assert.assertFalse(handler.receiveSlice(invalidSlice));
+ Assert.assertFalse(handler.makeReqIfComplete().isPresent());
+ }
+
private static TPipeTransferReq createReq(final byte version, final int
bodySize) {
final byte[] body = new byte[bodySize];
for (int i = 0; i < body.length; ++i) {
diff --git
a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java
new file mode 100644
index 00000000000..a00530b8c68
--- /dev/null
+++
b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReqTest.java
@@ -0,0 +1,149 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.commons.pipe.sink.payload.thrift.request;
+
+import org.apache.iotdb.commons.pipe.sink.compressor.PipeCompressor;
+import org.apache.iotdb.commons.pipe.sink.compressor.PipeCompressorFactory;
+import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
+
+import org.apache.tsfile.utils.BytesUtils;
+import org.apache.tsfile.utils.PublicBAOS;
+import org.apache.tsfile.utils.ReadWriteIOUtils;
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.DataOutputStream;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+public class PipeTransferCompressedReqTest {
+
+ @Test
+ public void testPipeTransferCompressedReq() throws IOException {
+ final TPipeTransferReq originalReq = new TPipeTransferReq();
+ originalReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion();
+ originalReq.type = PipeRequestType.TRANSFER_TABLET_BINARY.getType();
+ originalReq.body = ByteBuffer.wrap(new byte[] {1, 2, 3, 4});
+
+ final TPipeTransferReq compressedReq =
+ PipeTransferCompressedReq.toTPipeTransferReq(
+ originalReq,
+ Collections.singletonList(
+ PipeCompressorFactory.getCompressor(
+ PipeCompressor.PipeCompressionType.GZIP.getIndex())));
+ final int compressedBodySize = compressedReq.body.remaining();
+ final int maxDecompressedLength =
+ Math.max(compressedBodySize, originalReq.body.remaining() + 3);
+ Assert.assertEquals(
+ maxDecompressedLength,
+
PipeTransferCompressedReq.getMaxDecompressedLengthInBytes(compressedReq));
+ Assert.assertEquals(
+ maxDecompressedLength - compressedBodySize,
+
PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(compressedReq));
+ Assert.assertEquals(0, compressedReq.body.position());
+
+ final TPipeTransferReq decompressedReq =
+ PipeTransferCompressedReq.fromTPipeTransferReq(compressedReq);
+
+ Assert.assertEquals(IoTDBSinkRequestVersion.VERSION_1.getVersion(),
compressedReq.version);
+ Assert.assertEquals(PipeRequestType.TRANSFER_COMPRESSED.getType(),
compressedReq.type);
+ Assert.assertEquals(originalReq.version, decompressedReq.version);
+ Assert.assertEquals(originalReq.type, decompressedReq.type);
+ Assert.assertArrayEquals(originalReq.getBody(), decompressedReq.getBody());
+ }
+
+ @Test
+ public void testAdditionalDecompressedLengthExcludesTransferFrameBody()
throws IOException {
+ final TPipeTransferReq originalReq = new TPipeTransferReq();
+ originalReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion();
+ originalReq.type = PipeRequestType.TRANSFER_TABLET_BINARY.getType();
+ final byte[] highlyCompressibleBody = new byte[16 * 1024];
+ Arrays.fill(highlyCompressibleBody, (byte) 1);
+ originalReq.body = ByteBuffer.wrap(highlyCompressibleBody);
+
+ final TPipeTransferReq compressedReq =
+ PipeTransferCompressedReq.toTPipeTransferReq(
+ originalReq,
+ Collections.singletonList(
+ PipeCompressorFactory.getCompressor(
+ PipeCompressor.PipeCompressionType.GZIP.getIndex())));
+ final int compressedBodySize = compressedReq.body.remaining();
+ final int maxDecompressedLength =
+
PipeTransferCompressedReq.getMaxDecompressedLengthInBytes(compressedReq);
+
+ Assert.assertTrue(maxDecompressedLength > compressedBodySize);
+ Assert.assertEquals(
+ maxDecompressedLength - compressedBodySize,
+
PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(compressedReq));
+ Assert.assertEquals(0, compressedReq.body.position());
+ }
+
+ @Test
+ public void testPipeTransferCompressedReqFromLegacyV13Body() throws
IOException {
+ final TPipeTransferReq originalReq = new TPipeTransferReq();
+ originalReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion();
+ originalReq.type = PipeRequestType.TRANSFER_TABLET_BINARY.getType();
+ originalReq.body = ByteBuffer.wrap(new byte[] {1, 2, 3, 4});
+
+ final TPipeTransferReq compressedReq = new TPipeTransferReq();
+ compressedReq.version = IoTDBSinkRequestVersion.VERSION_1.getVersion();
+ compressedReq.type = PipeRequestType.TRANSFER_COMPRESSED.getType();
+ compressedReq.body =
+ serializeLegacyCompressedBody(
+ originalReq,
+ Collections.singletonList(
+ PipeCompressorFactory.getCompressor(
+ PipeCompressor.PipeCompressionType.GZIP.getIndex())));
+
+ final TPipeTransferReq decompressedReq =
+ PipeTransferCompressedReq.fromTPipeTransferReq(compressedReq);
+
+ Assert.assertEquals(originalReq.version, decompressedReq.version);
+ Assert.assertEquals(originalReq.type, decompressedReq.type);
+ Assert.assertArrayEquals(originalReq.getBody(), decompressedReq.getBody());
+ }
+
+ private static ByteBuffer serializeLegacyCompressedBody(
+ final TPipeTransferReq originalReq, final List<PipeCompressor>
compressors)
+ throws IOException {
+ try (final PublicBAOS byteArrayOutputStream = new PublicBAOS();
+ final DataOutputStream outputStream = new
DataOutputStream(byteArrayOutputStream)) {
+ 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);
+
+ return ByteBuffer.wrap(byteArrayOutputStream.getBuf(), 0,
byteArrayOutputStream.size());
+ }
+ }
+}
