Skip to content

Commit 71696f3

Browse files
CaideyipiCaideyipi
authored
Add pipe receiver memory protection (#18090)
* Add pipe receiver memory protection * Add pipe receiver memory protection tests * Address receiver memory review comments
1 parent 35cf94a commit 71696f3

11 files changed

Lines changed: 383 additions & 16 deletions

File tree

‎iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1243,6 +1243,9 @@ public final class DataNodePipeMessages {
12431243
public static final String RECEIVER_ID_UNSUPPORTED_STATEMENT_TYPE_FOR_REDIRECTION =
12441244
"Receiver id = {}: Unsupported statement type {} for redirection.";
12451245
public static final String RECEIVER_IS_READY = "Receiver-{} is ready";
1246+
public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
1247+
"Temporarily out of memory when %s. Requested memory: %d bytes, used memory: %d bytes, "
1248+
+ "free memory: %d bytes, total non-floating memory: %d bytes";
12461249
public static final String REGISTER_WITH_INTERVAL_IN_SECONDS_SUCCESSFULLY =
12471250
"Register {} with interval in seconds {} successfully.";
12481251
public static final String SOCKET_CLOSED_WHEN_EXECUTING_READTILLFULL =

‎iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1172,6 +1172,9 @@ public final class DataNodePipeMessages {
11721172
public static final String RECEIVER_ID_UNSUPPORTED_STATEMENT_TYPE_FOR_REDIRECTION =
11731173
"Receiver id = {}：不支持的 statement type {} 用于 redirection。";
11741174
public static final String RECEIVER_IS_READY = "Receiver-{} 已就绪";
1175+
public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
1176+
"执行 %s 时暂时内存不足。请求内存：%d bytes，已用内存：%d bytes，可用内存：%d bytes，"
1177+
+ "总非浮动内存：%d bytes";
11751178
public static final String REGISTER_WITH_INTERVAL_IN_SECONDS_SUCCESSFULLY =
11761179
"成功注册 {}，间隔时间（秒）{}。";
11771180
public static final String SOCKET_CLOSED_WHEN_EXECUTING_READTILLFULL =

‎iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java‎

Lines changed: 101 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -192,6 +192,7 @@ public class IoTDBDataNodeReceiver extends IoTDBFileReceiver {
192192
private static final PipeConfig PIPE_CONFIG = PipeConfig.getInstance();
193193

194194
private PipeMemoryBlock allocatedMemoryBlock;
195+
private final List<PipeMemoryBlock> allocatedSliceMemoryBlocks = new ArrayList<>();
195196
private final Set<String> autoCreatedTreeDatabases = ConcurrentHashMap.newKeySet();
196197
private final Set<String> conflictedTreeDatabases = ConcurrentHashMap.newKeySet();
197198

@@ -219,7 +220,7 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) {
219220
if (PipeRequestType.isValidatedRequestType(rawRequestType)) {
220221
final PipeRequestType requestType = PipeRequestType.valueOf(rawRequestType);
221222
if (requestType != PipeRequestType.TRANSFER_SLICE) {
222-
sliceReqHandler.clear();
223+
clearSliceReqHandler();
223224
}
224225
switch (requestType) {
225226
case HANDSHAKE_DATANODE_V1:
@@ -442,8 +443,18 @@ public synchronized TPipeTransferResp receive(final TPipeTransferReq req) {
442443
}
443444
case TRANSFER_COMPRESSED:
444445
{
446+
long requestedMemorySizeInBytes = 0;
445447
try {
446-
return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
448+
requestedMemorySizeInBytes =
449+
PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(req);
450+
try (final PipeMemoryBlock ignored =
451+
tryAllocateReceiverMemory(requestedMemorySizeInBytes)) {
452+
return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
453+
}
454+
} catch (final PipeRuntimeOutOfMemoryCriticalException e) {
455+
return new TPipeTransferResp(
456+
getReceiverTemporaryUnavailableStatus(
457+
"decompressing pipe transfer request", requestedMemorySizeInBytes, e));
447458
} finally {
448459
PipeDataNodeReceiverMetrics.getInstance()
449460
.recordTransferCompressedTimer(System.nanoTime() - startTime);
@@ -823,22 +834,96 @@ private String getConfigReceiverId() {
823834
}
824835

825836
private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq pipeTransferSliceReq) {
826-
final boolean isInorder = sliceReqHandler.receiveSlice(pipeTransferSliceReq);
827-
if (!isInorder) {
837+
final long sliceBodySizeInBytes = getSliceBodySizeInBytes(pipeTransferSliceReq);
838+
long requestedMemorySizeInBytes = sliceBodySizeInBytes;
839+
String memoryAction = "buffering sliced pipe transfer request";
840+
PipeMemoryBlock sliceMemoryBlock = null;
841+
try {
842+
sliceMemoryBlock = tryAllocateReceiverMemory(sliceBodySizeInBytes);
843+
844+
final boolean isInorder = sliceReqHandler.receiveSlice(pipeTransferSliceReq);
845+
if (!isInorder) {
846+
closeMemoryBlock(sliceMemoryBlock);
847+
clearSliceReqHandler();
848+
return new TPipeTransferResp(
849+
RpcUtils.getStatus(
850+
TSStatusCode.PIPE_TRANSFER_SLICE_OUT_OF_ORDER,
851+
"Slice request is out of order, please check the request sequence."));
852+
}
853+
854+
allocatedSliceMemoryBlocks.add(sliceMemoryBlock);
855+
sliceMemoryBlock = null;
856+
857+
if (pipeTransferSliceReq.getSliceIndex() + 1 < pipeTransferSliceReq.getSliceCount()) {
858+
return new TPipeTransferResp(
859+
RpcUtils.getStatus(
860+
TSStatusCode.SUCCESS_STATUS,
861+
"Slice received, waiting for more slices to complete the request."));
862+
}
863+
864+
memoryAction = "assembling sliced pipe transfer request";
865+
requestedMemorySizeInBytes = pipeTransferSliceReq.getOriginBodySize();
866+
try (final PipeMemoryBlock ignored = tryAllocateReceiverMemory(requestedMemorySizeInBytes)) {
867+
final Optional<TPipeTransferReq> req = sliceReqHandler.makeReqIfComplete();
868+
if (!req.isPresent()) {
869+
return new TPipeTransferResp(
870+
RpcUtils.getStatus(
871+
TSStatusCode.SUCCESS_STATUS,
872+
"Slice received, waiting for more slices to complete the request."));
873+
}
874+
clearSliceReqHandler();
875+
return receive(req.get());
876+
}
877+
} catch (final PipeRuntimeOutOfMemoryCriticalException e) {
878+
closeMemoryBlock(sliceMemoryBlock);
879+
clearSliceReqHandler();
828880
return new TPipeTransferResp(
829-
RpcUtils.getStatus(
830-
TSStatusCode.PIPE_TRANSFER_SLICE_OUT_OF_ORDER,
831-
"Slice request is out of order, please check the request sequence."));
881+
getReceiverTemporaryUnavailableStatus(memoryAction, requestedMemorySizeInBytes, e));
882+
} catch (final RuntimeException e) {
883+
closeMemoryBlock(sliceMemoryBlock);
884+
clearSliceReqHandler();
885+
throw e;
832886
}
833-
final Optional<TPipeTransferReq> req = sliceReqHandler.makeReqIfComplete();
834-
if (!req.isPresent()) {
835-
return new TPipeTransferResp(
836-
RpcUtils.getStatus(
837-
TSStatusCode.SUCCESS_STATUS,
838-
"Slice received, waiting for more slices to complete the request."));
887+
}
888+
889+
private long getSliceBodySizeInBytes(final PipeTransferSliceReq pipeTransferSliceReq) {
890+
return pipeTransferSliceReq.getSliceBody() == null
891+
? 0
892+
: pipeTransferSliceReq.getSliceBody().length;
893+
}
894+
895+
private void clearSliceReqHandler() {
896+
sliceReqHandler.clear();
897+
allocatedSliceMemoryBlocks.forEach(this::closeMemoryBlock);
898+
allocatedSliceMemoryBlocks.clear();
899+
}
900+
901+
private void closeMemoryBlock(final PipeMemoryBlock memoryBlock) {
902+
if (Objects.nonNull(memoryBlock)) {
903+
memoryBlock.close();
839904
}
840-
// sliceReqHandler will be cleared in the receive(req) method
841-
return receive(req.get());
905+
}
906+
907+
private PipeMemoryBlock tryAllocateReceiverMemory(final long requestedMemorySizeInBytes)
908+
throws PipeRuntimeOutOfMemoryCriticalException {
909+
return PipeDataNodeResourceManager.memory()
910+
.forceAllocate(Math.max(requestedMemorySizeInBytes, 0));
911+
}
912+
913+
@Override
914+
protected TSStatus getReceiverTemporaryUnavailableStatus(
915+
final String action,
916+
final long requestedMemorySizeInBytes,
917+
final PipeRuntimeOutOfMemoryCriticalException e) {
918+
return new TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode())
919+
.setMessage(
920+
String.format(
921+
DataNodePipeMessages.RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT,
922+
action,
923+
requestedMemorySizeInBytes,
924+
PipeDataNodeResourceManager.memory().getUsedMemorySizeInBytes(),
925+
PipeDataNodeResourceManager.memory().getFreeMemorySizeInBytes(),
926+
PipeDataNodeResourceManager.memory().getTotalNonFloatingMemorySizeInBytes()));
842927
}
843928

844929
/**
@@ -1395,6 +1480,7 @@ private TSStatus executeStatementForTableModelWithPermissionCheck(
13951480

13961481
@Override
13971482
public synchronized void handleExit() {
1483+
clearSliceReqHandler();
13981484
if (Objects.nonNull(configReceiverId.get())) {
13991485
try {
14001486
ClusterConfigTaskExecutor.getInstance().handlePipeConfigClientExit(configReceiverId.get());

‎iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -595,6 +595,8 @@ public final class PipeMessages {
595595
"Receiver id = %s: Handshake failed, response status = %s.";
596596
public static final String RECEIVER_HANDSHAKE_FAILED_LOGIN =
597597
"Receiver id = %s: Handshake failed because login failed, response status = %s.";
598+
public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
599+
"Temporarily out of memory when %s. Requested memory: %d bytes. Root cause: %s";
598600
public static final String RECEIVER_USER_LOGIN_SUCCESS =
599601
"Receiver id = {}: User {} login successfully.";
600602
public static final String RECEIVER_EXITED =

‎iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -577,6 +577,8 @@ public final class PipeMessages {
577577
"接收器 id = %s：握手失败，响应状态 = %s。";
578578
public static final String RECEIVER_HANDSHAKE_FAILED_LOGIN =
579579
"接收器 id = %s：因登录失败导致握手失败，响应状态 = %s。";
580+
public static final String RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT =
581+
"执行 %s 时暂时内存不足。请求内存：%d bytes。根因：%s";
580582
public static final String RECEIVER_USER_LOGIN_SUCCESS =
581583
"接收器 id = {}：用户 {} 登录成功。";
582584
public static final String RECEIVER_EXITED =

‎iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/receiver/IoTDBFileReceiver.java‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
import org.apache.iotdb.commons.audit.UserEntity;
2525
import org.apache.iotdb.commons.conf.CommonDescriptor;
2626
import org.apache.iotdb.commons.exception.IllegalPathException;
27+
import org.apache.iotdb.commons.exception.pipe.PipeRuntimeOutOfMemoryCriticalException;
2728
import org.apache.iotdb.commons.i18n.PipeMessages;
2829
import org.apache.iotdb.commons.pipe.config.PipeConfig;
2930
import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
@@ -386,7 +387,7 @@ protected final TPipeTransferResp handleTransferFilePiece(
386387
final PipeTransferFilePieceReq req,
387388
final boolean isRequestThroughAirGap,
388389
final boolean isSingleFile) {
389-
try {
390+
try (final AutoCloseable ignored = tryAllocateMemoryForFilePiece(req)) {
390391
updateWritingFileIfNeeded(req.getFileName(), isSingleFile);
391392

392393
// If the request is through air gap, the sender will resend the file piece from the beginning
@@ -419,6 +420,18 @@ protected final TPipeTransferResp handleTransferFilePiece(
419420
writingFileWriter.write(req.getFilePiece());
420421
return PipeTransferFilePieceResp.toTPipeTransferResp(
421422
RpcUtils.SUCCESS_STATUS, writingFileWriter.length());
423+
} catch (final PipeRuntimeOutOfMemoryCriticalException e) {
424+
final TSStatus status =
425+
getReceiverTemporaryUnavailableStatus(
426+
"receiving pipe file piece", getFilePieceSizeInBytes(req), e);
427+
PipeLogger.log(
428+
LOGGER::warn, e, PipeMessages.RECEIVER_FAILED_WRITE_FILE_PIECE, receiverId.get(), req);
429+
try {
430+
return PipeTransferFilePieceResp.toTPipeTransferResp(
431+
status, PipeTransferFilePieceResp.ERROR_END_OFFSET);
432+
} catch (Exception ex) {
433+
return PipeTransferFilePieceResp.toTPipeTransferResp(status);
434+
}
422435
} catch (final Exception e) {
423436
PipeLogger.log(
424437
LOGGER::warn, e, PipeMessages.RECEIVER_FAILED_WRITE_FILE_PIECE, receiverId.get(), req);
@@ -435,6 +448,28 @@ protected final TPipeTransferResp handleTransferFilePiece(
435448
}
436449
}
437450

451+
protected AutoCloseable tryAllocateMemoryForFilePiece(final PipeTransferFilePieceReq req)
452+
throws PipeRuntimeOutOfMemoryCriticalException {
453+
return () -> {};
454+
}
455+
456+
protected TSStatus getReceiverTemporaryUnavailableStatus(
457+
final String action,
458+
final long requestedMemorySizeInBytes,
459+
final PipeRuntimeOutOfMemoryCriticalException e) {
460+
return new TSStatus(TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode())
461+
.setMessage(
462+
String.format(
463+
PipeMessages.RECEIVER_TEMPORARILY_OUT_OF_MEMORY_FORMAT,
464+
action,
465+
requestedMemorySizeInBytes,
466+
e.getMessage()));
467+
}
468+
469+
private static long getFilePieceSizeInBytes(final PipeTransferFilePieceReq req) {
470+
return req.getFilePiece() == null ? 0 : req.getFilePiece().length;
471+
}
472+
438473
protected final void updateWritingFileIfNeeded(final String fileName, final boolean isSingleFile)
439474
throws IOException {
440475
if (isFileExistedAndNameCorrect(fileName)) {

‎iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/common/PipeTransferSliceReqHandler.java‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,8 +43,14 @@ public class PipeTransferSliceReqHandler {
4343

4444
private int sliceCount = -1;
4545
private final List<byte[]> sliceBodies = new ArrayList<>();
46+
private long receivedBodySize = 0;
4647

4748
public boolean receiveSlice(final PipeTransferSliceReq req) {
49+
if (!isValidSliceReq(req)) {
50+
clear();
51+
return false;
52+
}
53+
4854
if (orderId == -1
4955
|| originReqType == -1
5056
|| originBodySize == -1
@@ -59,6 +65,7 @@ public boolean receiveSlice(final PipeTransferSliceReq req) {
5965
originReqType = req.getOriginReqType();
6066
originBodySize = req.getOriginBodySize();
6167
sliceCount = req.getSliceCount();
68+
receivedBodySize = 0;
6269
} else {
6370
LOGGER.warn(
6471
PipeMessages.INVALID_STATE_SLICE,
@@ -107,10 +114,39 @@ public boolean receiveSlice(final PipeTransferSliceReq req) {
107114
return false;
108115
}
109116

117+
if (receivedBodySize + req.getSliceBody().length > originBodySize) {
118+
LOGGER.warn(
119+
"Received slice body size {} exceeds origin body size {}",
120+
receivedBodySize + req.getSliceBody().length,
121+
originBodySize);
122+
clear();
123+
return false;
124+
}
125+
110126
sliceBodies.add(req.getSliceBody());
127+
receivedBodySize += req.getSliceBody().length;
128+
129+
if (sliceBodies.size() == sliceCount && receivedBodySize != originBodySize) {
130+
LOGGER.warn(
131+
"Received slice body size {} is not equal to origin body size {}",
132+
receivedBodySize,
133+
originBodySize);
134+
clear();
135+
return false;
136+
}
137+
111138
return true;
112139
}
113140

141+
private boolean isValidSliceReq(final PipeTransferSliceReq req) {
142+
return req.getOriginBodySize() >= 0
143+
&& req.getSliceCount() > 0
144+
&& req.getSliceIndex() >= 0
145+
&& req.getSliceIndex() < req.getSliceCount()
146+
&& req.getSliceBody() != null
147+
&& req.getSliceBody().length <= req.getOriginBodySize();
148+
}
149+
114150
public Optional<TPipeTransferReq> makeReqIfComplete() {
115151
if (sliceBodies.size() != sliceCount) {
116152
return Optional.empty();
@@ -134,5 +170,6 @@ public void clear() {
134170
originBodySize = -1;
135171
sliceCount = -1;
136172
sliceBodies.clear();
173+
receivedBodySize = 0;
137174
}
138175
}

‎iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/thrift/request/PipeTransferCompressedReq.java‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,27 @@ public static TPipeTransferReq fromTPipeTransferReq(final TPipeTransferReq trans
113113
return decompressedReq;
114114
}
115115

116+
/** Get the largest intermediate decompressed body size without consuming the request body. */
117+
public static int getMaxDecompressedLengthInBytes(final TPipeTransferReq transferReq) {
118+
final ByteBuffer compressedBuffer = transferReq.body.duplicate();
119+
120+
int maxDecompressedLength = compressedBuffer.remaining();
121+
final int compressorsSize = ReadWriteIOUtils.readByte(compressedBuffer);
122+
for (int i = 0; i < compressorsSize; ++i) {
123+
ReadWriteIOUtils.readByte(compressedBuffer);
124+
final int decompressedLength = ReadWriteIOUtils.readInt(compressedBuffer);
125+
checkDecompressedLength(decompressedLength);
126+
maxDecompressedLength = Math.max(maxDecompressedLength, decompressedLength);
127+
}
128+
return maxDecompressedLength;
129+
}
130+
131+
/** Get the largest additional decompressed body size beyond the current transfer frame. */
132+
public static int getMaxAdditionalDecompressedLengthInBytes(final TPipeTransferReq transferReq) {
133+
final int transferFrameBodySize = transferReq.body.duplicate().remaining();
134+
return Math.max(0, getMaxDecompressedLengthInBytes(transferReq) - transferFrameBodySize);
135+
}
136+
116137
/** This method is used to prevent decompression bomb attacks. */
117138
private static void checkDecompressedLength(final int decompressedLength)
118139
throws IllegalArgumentException {

0 commit comments

Comments
 (0)