This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 71696f333af Add pipe receiver memory protection (#18090)
71696f333af is described below
commit 71696f333af3f71e2bea946754436b90247b0d4e
Author: Caideyipi <[email protected]>
AuthorDate: Tue Jul 7 10:31:11 2026 +0800
Add pipe receiver memory protection (#18090)
* Add pipe receiver memory protection
* Add pipe receiver memory protection tests
* Address receiver memory review comments
---
.../apache/iotdb/db/i18n/DataNodePipeMessages.java | 3 +
.../apache/iotdb/db/i18n/DataNodePipeMessages.java | 3 +
.../protocol/thrift/IoTDBDataNodeReceiver.java | 116 ++++++++++++++++++---
.../apache/iotdb/commons/i18n/PipeMessages.java | 2 +
.../apache/iotdb/commons/i18n/PipeMessages.java | 2 +
.../commons/pipe/receiver/IoTDBFileReceiver.java | 37 ++++++-
.../thrift/common/PipeTransferSliceReqHandler.java | 37 +++++++
.../thrift/request/PipeTransferCompressedReq.java | 21 ++++
.../pipe/receiver/IoTDBFileReceiverTest.java | 89 ++++++++++++++++
.../common/PipeTransferSliceReqBuilderTest.java | 52 +++++++++
.../request/PipeTransferCompressedReqTest.java | 37 +++++++
11 files changed, 383 insertions(+), 16 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
index dad8dcbc7f4..2f3fdd5ff64 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
@@ -1243,6 +1243,9 @@ public final class DataNodePipeMessages {
public static final String
RECEIVER_ID_UNSUPPORTED_STATEMENT_TYPE_FOR_REDIRECTION =
"Receiver id = {}: Unsupported statement type {} for redirection.";
public static final String RECEIVER_IS_READY = "Receiver-{} is ready";
+ public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_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";
public static final String REGISTER_WITH_INTERVAL_IN_SECONDS_SUCCESSFULLY =
"Register {} with interval in seconds {} successfully.";
public static final String SOCKET_CLOSED_WHEN_EXECUTING_READTILLFULL =
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
index b9b042aff52..b2b67f4d4b5 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
@@ -1172,6 +1172,9 @@ public final class DataNodePipeMessages {
public static final String
RECEIVER_ID_UNSUPPORTED_STATEMENT_TYPE_FOR_REDIRECTION =
"Receiver id = {}：不支持的 statement type {} 用于 redirection。";
public static final String RECEIVER_IS_READY = "Receiver-{} 已就绪";
+ public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
+ "执行 %s 时暂时内存不足。请求内存：%d bytes，已用内存：%d bytes，可用内存：%d bytes，"
+ + "总非浮动内存：%d bytes";
public static final String REGISTER_WITH_INTERVAL_IN_SECONDS_SUCCESSFULLY =
"成功注册 {}，间隔时间（秒）{}。";
public static final String SOCKET_CLOSED_WHEN_EXECUTING_READTILLFULL =
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 8944466847d..f4a4baed694 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
@@ -192,6 +192,7 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
private static final PipeConfig PIPE_CONFIG = PipeConfig.getInstance();
private PipeMemoryBlock allocatedMemoryBlock;
+ private final List<PipeMemoryBlock> allocatedSliceMemoryBlocks = new
ArrayList<>();
private final Set<String> autoCreatedTreeDatabases =
ConcurrentHashMap.newKeySet();
private final Set<String> conflictedTreeDatabases =
ConcurrentHashMap.newKeySet();
@@ -219,7 +220,7 @@ public class IoTDBDataNodeReceiver extends
IoTDBFileReceiver {
if (PipeRequestType.isValidatedRequestType(rawRequestType)) {
final PipeRequestType requestType =
PipeRequestType.valueOf(rawRequestType);
if (requestType != PipeRequestType.TRANSFER_SLICE) {
- sliceReqHandler.clear();
+ clearSliceReqHandler();
}
switch (requestType) {
case HANDSHAKE_DATANODE_V1:
@@ -442,8 +443,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);
@@ -823,22 +834,96 @@ 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(
+ DataNodePipeMessages.RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT,
+ action,
+ requestedMemorySizeInBytes,
+
PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(),
+
PipeDataNodeResourceManager.memory().getFreeMemorySizeInBytes(),
+
PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes()));
}
/**
@@ -1395,6 +1480,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/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java
b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java
index 08e12488b03..afb977f85a2 100644
---
a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java
+++
b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java
@@ -595,6 +595,8 @@ public final class PipeMessages {
"Receiver id = %s: Handshake failed, response status = %s.";
public static final String RECEIVER_HANDSHAKE_FAILED_LOGIN =
"Receiver id = %s: Handshake failed because login failed, response
status = %s.";
+ public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
+ "Temporarily out of memory when %s. Requested memory: %d bytes. Root
cause: %s";
public static final String RECEIVER_USER_LOGIN_SUCCESS =
"Receiver id = {}: User {} login successfully.";
public static final String RECEIVER_EXITED =
diff --git
a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java
b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java
index 0c855952e99..87f1601a754 100644
---
a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java
+++
b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java
@@ -577,6 +577,8 @@ public final class PipeMessages {
"接收器 id = %s：握手失败，响应状态 = %s。";
public static final String RECEIVER_HANDSHAKE_FAILED_LOGIN =
"接收器 id = %s：因登录失败导致握手失败，响应状态 = %s。";
+ public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
+ "执行 %s 时暂时内存不足。请求内存：%d bytes。根因：%s";
public static final String RECEIVER_USER_LOGIN_SUCCESS =
"接收器 id = {}：用户 {} 登录成功。";
public static final String RECEIVER_EXITED =
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 8304161c9d0..f05e5069eed 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
@@ -24,6 +24,7 @@ import org.apache.iotdb.commons.audit.IAuditEntity;
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.i18n.PipeMessages;
import org.apache.iotdb.commons.pipe.config.PipeConfig;
import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
@@ -386,7 +387,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
@@ -419,6 +420,18 @@ 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, PipeMessages.RECEIVER_FAILED_WRITE_FILE_PIECE,
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, e, PipeMessages.RECEIVER_FAILED_WRITE_FILE_PIECE,
receiverId.get(), req);
@@ -435,6 +448,28 @@ 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(
+ PipeMessages.RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT,
+ 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 c95d4ec4306..1e219dbfd27 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
@@ -43,8 +43,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
@@ -59,6 +65,7 @@ public class PipeTransferSliceReqHandler {
originReqType = req.getOriginReqType();
originBodySize = req.getOriginBodySize();
sliceCount = req.getSliceCount();
+ receivedBodySize = 0;
} else {
LOGGER.warn(
PipeMessages.INVALID_STATE_SLICE,
@@ -107,10 +114,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();
@@ -134,5 +170,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 8f6f36edf1b..5dc1f1f6661 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
@@ -113,6 +113,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 5a1ed3e653a..f774d725bae 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,10 +22,13 @@ 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.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.response.PipeTransferFilePieceResp;
import org.apache.iotdb.rpc.TSStatusCode;
import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp;
@@ -138,10 +141,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;
DummyFileReceiver(final File baseDir) {
receiverFileBaseDir = baseDir;
@@ -166,6 +214,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);
+ }
+
TPipeTransferResp sealFileV1(final String fileName, final long fileLength)
throws IOException {
return
handleTransferFileSealV1(DummyFileSealReqV1.toTPipeTransferReq(fileName,
fileLength));
}
@@ -228,6 +293,14 @@ public class IoTDBFileReceiverTest {
return new TSStatus(200);
}
+ @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) {
@@ -252,6 +325,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 cc836cae945..40ec6d2d62c 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
@@ -37,6 +37,7 @@ import org.junit.Test;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.Optional;
public class PipeTransferSliceReqBuilderTest {
@@ -123,6 +124,57 @@ public class PipeTransferSliceReqBuilderTest {
Assert.assertEquals(2, sliceReq.getSliceCount());
}
+ @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
index ccdbc5086a5..a00530b8c68 100644
---
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
@@ -51,6 +51,17 @@ public class PipeTransferCompressedReqTest {
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);
@@ -61,6 +72,32 @@ public class PipeTransferCompressedReqTest {
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();
