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 7310b2bc0eb Fix oversized Load TsFile piece dispatch (#18627)
7310b2bc0eb is described below
commit 7310b2bc0eb73952d88c84307fdb5a13b3660e6a
Author: Caideyipi <[email protected]>
AuthorDate: Fri Sep 18 15:47:36 2026 +0800
Fix oversized Load TsFile piece dispatch (#18627)
* Fix oversized Load TsFile piece dispatch
* Address Load TsFile slicing review feedback
---
.../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 3 +
.../iotdb/db/i18n/DataNodeQueryMessages.java | 5 +
.../iotdb/db/i18n/StorageEngineMessages.java | 17 ++
.../apache/iotdb/db/i18n/DataNodeMiscMessages.java | 3 +
.../iotdb/db/i18n/DataNodeQueryMessages.java | 5 +
.../iotdb/db/i18n/StorageEngineMessages.java | 17 ++
.../impl/DataNodeInternalRPCServiceImpl.java | 41 +++-
.../scheduler/load/LoadTsFileDispatcherImpl.java | 169 +++++++++++++---
.../iotdb/db/storageengine/StorageEngine.java | 32 +++
.../db/storageengine/load/LoadTsFileManager.java | 67 +++++-
.../load/LoadTsFilePieceNodeAssembler.java | 170 ++++++++++++++++
.../DataNodeInternalRPCServiceImplLoadTest.java | 94 +++++++++
.../load/LoadTsFileDispatcherImplTest.java | 225 +++++++++++++++++++++
.../load/LoadTsFilePieceNodeAssemblerTest.java | 215 ++++++++++++++++++++
.../src/main/thrift/datanode.thrift | 5 +
15 files changed, 1035 insertions(+), 33 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
index 12f4e374ca7..b571df825ab 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
@@ -22,6 +22,9 @@ package org.apache.iotdb.db.i18n;
/** Compile-time i18n constants for DataNode misc subsystems (English). */
public final class DataNodeMiscMessages {
+ public static final String
MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA =
+ "Missing Load TsFile slice metadata: %s";
+
public static final String INVALID_PIPE_NAME =
"Invalid pipeName";
public static final String READ_OBJECT_CONTENT_FROM_REMOTE_FILE =
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
index 7b60fe71c2b..771e3a4517d 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
@@ -21,6 +21,11 @@ package org.apache.iotdb.db.i18n;
public final class DataNodeQueryMessages {
+ public static final String
EXCEPTION_INVALID_THRIFT_MAXIMUM_FRAME_SIZE_ARG_FROM_ARG_A639588B =
+ "Invalid Thrift maximum frame size %d from %s";
+ public static final String
MESSAGE_FAILED_TO_DISPATCH_LOAD_COMMAND_ARG_TO_NODE_ARG_BECAUSE_OF_EXCEPTION_ARG_2D8A483D
=
+ "failed to dispatch load command %s to node %s because of exception: %s";
+
// --- Common ---
public static final String NO_MATCHED_DATABASE_PLEASE_CHECK_THE_PATH =
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java
index ccde46c4881..3d66a888696 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/StorageEngineMessages.java
@@ -21,6 +21,23 @@ package org.apache.iotdb.db.i18n;
public final class StorageEngineMessages {
+ public static final String
MESSAGE_LOAD_TSFILE_SLICE_BODY_IS_NULL_OR_EMPTY_2A65366C =
+ "Load TsFile slice body is null or empty";
+ public static final String
MESSAGE_INVALID_LOAD_TSFILE_SLICE_METADATA_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_379BF1B8
=
+ "Invalid Load TsFile slice metadata: sliceCount=%d, originBodySize=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_SLICE_METADATA_CHANGED_SLICECOUNT_ARG_EXPECTED_ARG_ORIGINBODYSIZE_ARG_EXPECTED_ARG_B16B3122
=
+ "Load TsFile slice metadata changed: sliceCount=%d (expected %d),
originBodySize=%d (expected %d)";
+ public static final String
MESSAGE_UNEXPECTED_LOAD_TSFILE_SLICE_INDEX_ARG_EXPECTED_ARG_SLICECOUNT_ARG_76260F62
=
+ "Unexpected Load TsFile slice index %d: expected %d, sliceCount=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_SLICE_EXCEEDS_ORIGINBODYSIZE_ASSEMBLEDSIZE_ARG_SLICESIZE_ARG_ORIGINBODYSIZE_ARG_0198E0D0
=
+ "Load TsFile slice exceeds originBodySize: assembledSize=%d,
sliceSize=%d, originBodySize=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_BODY_COMPLETED_BEFORE_THE_LAST_SLICE_RECEIVED_ARG_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_0425EA2C
=
+ "Load TsFile body completed before the last slice: received=%d,
sliceCount=%d, originBodySize=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_BODY_SIZE_MISMATCH_ASSEMBLEDSIZE_ARG_ORIGINBODYSIZE_ARG_5733FC4B
=
+ "Load TsFile body size mismatch: assembledSize=%d, originBodySize=%d";
+ public static final String
MESSAGE_MISSING_LOAD_TSFILE_ASSEMBLER_FOR_UUID_ARG_DATAREGION_ARG_SLICEINDEX_ARG_FF6EA463
=
+ "Missing Load TsFile assembler for uuid %s, DataRegion %s,
sliceIndex=%d";
+
public static final String IS_NOT_SUPPORTED =
" is not supported";
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
index d809fea1992..bf15d6cde07 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java
@@ -22,6 +22,9 @@ package org.apache.iotdb.db.i18n;
/** 编译时国际化常量 - DataNode 杂项子系统（中文）。 */
public final class DataNodeMiscMessages {
+ public static final String
MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA =
+ "缺少 Load TsFile 分片元数据：%s";
+
public static final String INVALID_PIPE_NAME =
"无效的 pipeName";
public static final String READ_OBJECT_CONTENT_FROM_REMOTE_FILE =
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
index eeefd24141d..24170c487e5 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodeQueryMessages.java
@@ -21,6 +21,11 @@ package org.apache.iotdb.db.i18n;
public final class DataNodeQueryMessages {
+ public static final String
EXCEPTION_INVALID_THRIFT_MAXIMUM_FRAME_SIZE_ARG_FROM_ARG_A639588B =
+ "Thrift 最大帧大小 %d 无效，来源：%s";
+ public static final String
MESSAGE_FAILED_TO_DISPATCH_LOAD_COMMAND_ARG_TO_NODE_ARG_BECAUSE_OF_EXCEPTION_ARG_2D8A483D
=
+ "向节点派发 Load 命令失败，命令：%s，节点：%s，异常：%s";
+
// --- Common ---
public static final String NO_MATCHED_DATABASE_PLEASE_CHECK_THE_PATH =
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java
index 1244879ed42..dfc7cea9871 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/StorageEngineMessages.java
@@ -21,6 +21,23 @@ package org.apache.iotdb.db.i18n;
public final class StorageEngineMessages {
+ public static final String
MESSAGE_LOAD_TSFILE_SLICE_BODY_IS_NULL_OR_EMPTY_2A65366C =
+ "Load TsFile 分片内容为 null 或为空";
+ public static final String
MESSAGE_INVALID_LOAD_TSFILE_SLICE_METADATA_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_379BF1B8
=
+ "Load TsFile 分片元数据无效：sliceCount=%d，originBodySize=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_SLICE_METADATA_CHANGED_SLICECOUNT_ARG_EXPECTED_ARG_ORIGINBODYSIZE_ARG_EXPECTED_ARG_B16B3122
=
+ "Load TsFile 分片元数据发生变化：sliceCount=%d（预期 %d），originBodySize=%d（预期 %d）";
+ public static final String
MESSAGE_UNEXPECTED_LOAD_TSFILE_SLICE_INDEX_ARG_EXPECTED_ARG_SLICECOUNT_ARG_76260F62
=
+ "Load TsFile 分片索引 %d 不符合预期：预期 %d，sliceCount=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_SLICE_EXCEEDS_ORIGINBODYSIZE_ASSEMBLEDSIZE_ARG_SLICESIZE_ARG_ORIGINBODYSIZE_ARG_0198E0D0
=
+ "Load TsFile 分片超出
originBodySize：assembledSize=%d，sliceSize=%d，originBodySize=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_BODY_COMPLETED_BEFORE_THE_LAST_SLICE_RECEIVED_ARG_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_0425EA2C
=
+ "Load TsFile
内容在最后一个分片之前已达到原始大小：received=%d，sliceCount=%d，originBodySize=%d";
+ public static final String
MESSAGE_LOAD_TSFILE_BODY_SIZE_MISMATCH_ASSEMBLEDSIZE_ARG_ORIGINBODYSIZE_ARG_5733FC4B
=
+ "Load TsFile 内容大小不匹配：assembledSize=%d，originBodySize=%d";
+ public static final String
MESSAGE_MISSING_LOAD_TSFILE_ASSEMBLER_FOR_UUID_ARG_DATAREGION_ARG_SLICEINDEX_ARG_FF6EA463
=
+ "缺少 Load TsFile 分片组装器：uuid %s，DataRegion %s，sliceIndex=%d";
+
public static final String IS_NOT_SUPPORTED =
" 不受支持";
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
index 34ef951296c..e7213b7822c 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImpl.java
@@ -627,12 +627,51 @@ public class DataNodeInternalRPCServiceImpl implements
IDataNodeRPCService.Iface
throw new UnsupportedOperationException();
}
+ @Override
+ public int getThriftMaxFrameSize() {
+ return IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize();
+ }
+
@Override
public TLoadResp sendTsFilePieceNode(final TTsFilePieceReq req) {
- LOGGER.info(DataNodeMiscMessages.RECEIVE_LOAD_NODE, req.uuid);
+ if (!req.isSetSliceIndex() || req.sliceIndex == 0) {
+ LOGGER.info(DataNodeMiscMessages.RECEIVE_LOAD_NODE, req.uuid);
+ }
final ConsensusGroupId groupId =
ConsensusGroupId.Factory.createFromTConsensusGroupId(req.consensusGroupId);
+ final boolean isSliced =
+ req.isSetSliceIndex() || req.isSetSliceCount() ||
req.isSetOriginBodySize();
+ if (isSliced) {
+ if (!req.isSetSliceIndex() || !req.isSetSliceCount() ||
!req.isSetOriginBodySize()) {
+ final List<String> missingFields = new ArrayList<>(3);
+ if (!req.isSetSliceIndex()) {
+ missingFields.add("sliceIndex");
+ }
+ if (!req.isSetSliceCount()) {
+ missingFields.add("sliceCount");
+ }
+ if (!req.isSetOriginBodySize()) {
+ missingFields.add("originBodySize");
+ }
+ return createTLoadResp(
+ RpcUtils.getStatus(
+ TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR,
+ String.format(
+
DataNodeMiscMessages.MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA,
+ String.join(", ", missingFields))));
+ }
+ return createTLoadResp(
+ StorageEngine.getInstance()
+ .writeLoadTsFileNodeSlice(
+ (DataRegionId) groupId,
+ req.body,
+ req.uuid,
+ req.sliceIndex,
+ req.sliceCount,
+ req.originBodySize));
+ }
+
final LoadTsFilePieceNode pieceNode = (LoadTsFilePieceNode)
PlanNodeType.deserialize(req.body);
if (pieceNode == null) {
return createTLoadResp(
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java
index 1ba39ffd0bd..80f09f17e28 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImpl.java
@@ -19,6 +19,7 @@
package org.apache.iotdb.db.queryengine.plan.scheduler.load;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
@@ -57,18 +58,21 @@ import org.apache.iotdb.mpp.rpc.thrift.TTsFilePieceReq;
import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.thrift.TApplicationException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeoutException;
@@ -82,6 +86,7 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
private static final int MAX_CONNECTION_TIMEOUT_MS = 24 * 60 * 60 * 1000; //
1 day
private static final int FIRST_ADJUSTMENT_TIMEOUT_MS = 6 * 60 * 60 * 1000;
// 6 hours
+ private static final int LOAD_TSFILE_PIECE_RPC_FRAME_RESERVED_BYTES = 1024;
private static final AtomicInteger CONNECTION_TIMEOUT_MS =
new
AtomicInteger(IoTDBDescriptor.getInstance().getConfig().getConnectionTimeoutInMS());
@@ -92,6 +97,7 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
internalServiceClientManager;
private ExecutorService executor;
private final boolean isGeneratedByPipe;
+ private final Map<TEndPoint, Integer> endPoint2ThriftMaxFrameSize = new
ConcurrentHashMap<>();
public LoadTsFileDispatcherImpl(
IClientManager<TEndPoint, SyncDataNodeInternalServiceClient>
internalServiceClientManager,
@@ -143,7 +149,7 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
private void dispatchOneInstance(FragmentInstance instance)
throws FragmentInstanceDispatchException {
- TTsFilePieceReq loadTsFileReq = null;
+ ByteBuffer body = null;
for (TDataNodeLocation dataNodeLocation :
instance.getRegionReplicaSet().getDataNodeLocations()) {
@@ -151,18 +157,107 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
if (isDispatchedToLocal(endPoint)) {
dispatchLocally(instance);
} else {
- if (loadTsFileReq == null) {
- loadTsFileReq =
- new TTsFilePieceReq(
-
instance.getFragment().getPlanNodeTree().serializeToByteBuffer(),
- uuid,
- instance.getRegionReplicaSet().getRegionId());
+ if (body == null) {
+ body =
instance.getFragment().getPlanNodeTree().serializeToByteBuffer();
}
- dispatchRemote(loadTsFileReq, endPoint);
+ dispatchRemote(body, instance.getRegionReplicaSet().getRegionId(),
endPoint);
}
}
}
+ private int getLoadTsFilePieceBodySizeLimit(final TEndPoint endPoint) throws
Exception {
+ final int localMaxFrameSize =
IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize();
+ Integer remoteMaxFrameSize = endPoint2ThriftMaxFrameSize.get(endPoint);
+ if (remoteMaxFrameSize == null) {
+ // An oversized frame is rejected before the RPC handler runs and closes
the connection, so
+ // the receiver's frame limit cannot be recovered from the sender's
transport exception.
+ try (SyncDataNodeInternalServiceClient client =
+ internalServiceClientManager.borrowClient(endPoint)) {
+ remoteMaxFrameSize = client.getThriftMaxFrameSize();
+ } catch (Exception e) {
+ if (!isUnknownMethod(e)) {
+ throw e;
+ }
+ // Older receivers still accept unsliced pieces. The failed RPC
invalidates its client,
+ // so dispatchRemote borrows another client before sending the piece.
+ remoteMaxFrameSize = localMaxFrameSize;
+ }
+ if (remoteMaxFrameSize <= 0) {
+ throw new IllegalArgumentException(
+ String.format(
+ DataNodeQueryMessages
+
.EXCEPTION_INVALID_THRIFT_MAXIMUM_FRAME_SIZE_ARG_FROM_ARG_A639588B,
+ remoteMaxFrameSize,
+ endPoint));
+ }
+ endPoint2ThriftMaxFrameSize.put(endPoint, remoteMaxFrameSize);
+ }
+ return Math.max(
+ 1,
+ Math.min(localMaxFrameSize, remoteMaxFrameSize)
+ - LOAD_TSFILE_PIECE_RPC_FRAME_RESERVED_BYTES);
+ }
+
+ private static boolean isUnknownMethod(Throwable e) {
+ do {
+ if (e instanceof TApplicationException
+ && ((TApplicationException) e).getType() ==
TApplicationException.UNKNOWN_METHOD) {
+ return true;
+ }
+ } while ((e = e.getCause()) != null);
+ return false;
+ }
+
+ static List<TTsFilePieceReq> splitTsFilePieceReq(
+ final ByteBuffer body,
+ final String uuid,
+ final TConsensusGroupId consensusGroupId,
+ final int bodySizeLimit) {
+ if (bodySizeLimit <= 0) {
+ throw new IllegalArgumentException();
+ }
+
+ final int originBodySize = body.remaining();
+ final int sliceCount = getSliceCount(originBodySize, bodySizeLimit);
+ final List<TTsFilePieceReq> requests = new ArrayList<>(sliceCount);
+ if (sliceCount == 1) {
+ requests.add(createTsFilePieceReq(body.duplicate(), uuid,
consensusGroupId));
+ return requests;
+ }
+
+ final int originPosition = body.position();
+ for (int sliceIndex = 0; sliceIndex < sliceCount; sliceIndex++) {
+ final int startOffset = sliceIndex * bodySizeLimit;
+ final int endOffset = startOffset + Math.min(bodySizeLimit,
originBodySize - startOffset);
+ final ByteBuffer slicedBody = body.duplicate();
+ slicedBody.position(originPosition + startOffset);
+ slicedBody.limit(originPosition + endOffset);
+ requests.add(
+ createTsFilePieceReq(slicedBody.slice(), uuid, consensusGroupId)
+ .setSliceIndex(sliceIndex)
+ .setSliceCount(sliceCount)
+ .setOriginBodySize(originBodySize));
+ }
+ return requests;
+ }
+
+ static int getSliceCount(final int bodySize, final int bodySizeLimit) {
+ if (bodySize < 0 || bodySizeLimit <= 0) {
+ throw new IllegalArgumentException();
+ }
+ return bodySize == 0 ? 1 : (bodySize - 1) / bodySizeLimit + 1;
+ }
+
+ private static TTsFilePieceReq createTsFilePieceReq(
+ final ByteBuffer body, final String uuid, final TConsensusGroupId
consensusGroupId) {
+ final TTsFilePieceReq request =
+ new
TTsFilePieceReq().setUuid(uuid).setConsensusGroupId(consensusGroupId);
+ // The generated setter copies the whole buffer, while these immutable
slices remain valid until
+ // all replicas have been dispatched.
+ request.body = body;
+ return request;
+ }
+
public void dispatchLocally(FragmentInstance instance) throws
FragmentInstanceDispatchException {
if (isGeneratedByPipe) {
LOGGER.debug(DataNodeQueryMessages.RECEIVE_LOAD_NODE_FROM_UUID, uuid);
@@ -222,28 +317,36 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
}
}
- private void dispatchRemote(TTsFilePieceReq loadTsFileReq, TEndPoint
endPoint)
+ private void dispatchRemote(
+ ByteBuffer body, TConsensusGroupId consensusGroupId, TEndPoint endPoint)
throws FragmentInstanceDispatchException {
boolean transferAttemptRecorded = false;
- try (SyncDataNodeInternalServiceClient client =
- internalServiceClientManager.borrowClient(endPoint)) {
- client.setTimeout(CONNECTION_TIMEOUT_MS.get());
-
- final TLoadResp loadResp = client.sendTsFilePieceNode(loadTsFileReq);
- if (!loadResp.isAccepted()) {
- recordTransferAttempt(
- endPoint,
- false,
- loadResp.isSetStatus()
- ? String.valueOf(loadResp.getStatus().getCode())
- : UserDataTransferErrorCode.REMOTE_REJECTED.name(),
- null);
+ try {
+ final List<TTsFilePieceReq> loadTsFileReqs =
+ splitTsFilePieceReq(
+ body, uuid, consensusGroupId,
getLoadTsFilePieceBodySizeLimit(endPoint));
+ try (SyncDataNodeInternalServiceClient client =
+ internalServiceClientManager.borrowClient(endPoint)) {
+ client.setTimeout(CONNECTION_TIMEOUT_MS.get());
+
+ for (final TTsFilePieceReq loadTsFileReq : loadTsFileReqs) {
+ final TLoadResp loadResp = client.sendTsFilePieceNode(loadTsFileReq);
+ if (!loadResp.isAccepted()) {
+ recordTransferAttempt(
+ endPoint,
+ false,
+ loadResp.isSetStatus()
+ ? String.valueOf(loadResp.getStatus().getCode())
+ : UserDataTransferErrorCode.REMOTE_REJECTED.name(),
+ null);
+ transferAttemptRecorded = true;
+ LOGGER.warn(loadResp.message);
+ throw new FragmentInstanceDispatchException(loadResp.status);
+ }
+ }
+ recordTransferAttempt(endPoint, true, null, null);
transferAttemptRecorded = true;
- LOGGER.warn(loadResp.message);
- throw new FragmentInstanceDispatchException(loadResp.status);
}
- recordTransferAttempt(endPoint, true, null, null);
- transferAttemptRecorded = true;
} catch (Exception e) {
if (!transferAttemptRecorded) {
recordTransferAttempt(endPoint, false, null, e);
@@ -252,8 +355,11 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
final String exceptionMessage =
String.format(
- "failed to dispatch load command %s to node %s because of
exception: %s",
- loadTsFileReq, endPoint, e);
+ DataNodeQueryMessages
+
.MESSAGE_FAILED_TO_DISPATCH_LOAD_COMMAND_ARG_TO_NODE_ARG_BECAUSE_OF_EXCEPTION_ARG_2D8A483D,
+ uuid,
+ endPoint,
+ e);
LOGGER.warn(exceptionMessage, e);
throw new FragmentInstanceDispatchException(
new TSStatus()
@@ -361,8 +467,11 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
final String exceptionMessage =
String.format(
- "failed to dispatch load command %s to node %s because of
exception: %s",
- loadCommandReq, endPoint, e);
+ DataNodeQueryMessages
+
.MESSAGE_FAILED_TO_DISPATCH_LOAD_COMMAND_ARG_TO_NODE_ARG_BECAUSE_OF_EXCEPTION_ARG_2D8A483D,
+ loadCommandReq,
+ endPoint,
+ e);
LOGGER.warn(exceptionMessage, e);
throw new FragmentInstanceDispatchException(
new TSStatus()
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java
index 1ac6c15488b..a06ec8c00b5 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/StorageEngine.java
@@ -38,6 +38,7 @@ import
org.apache.iotdb.commons.exception.IllegalPathException;
import org.apache.iotdb.commons.exception.ShutdownException;
import org.apache.iotdb.commons.exception.StartupException;
import org.apache.iotdb.commons.file.SystemFileFactory;
+import
org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType;
import org.apache.iotdb.commons.schema.ttl.TTLCache;
import org.apache.iotdb.commons.service.IService;
import org.apache.iotdb.commons.service.ServiceType;
@@ -79,6 +80,7 @@ import
org.apache.iotdb.db.storageengine.dataregion.wal.WALManager;
import org.apache.iotdb.db.storageengine.dataregion.wal.exception.WALException;
import
org.apache.iotdb.db.storageengine.dataregion.wal.recover.WALRecoverManager;
import org.apache.iotdb.db.storageengine.load.LoadTsFileManager;
+import org.apache.iotdb.db.storageengine.load.LoadTsFilePieceNodeAssembler;
import org.apache.iotdb.db.storageengine.load.limiter.LoadTsFileRateLimiter;
import org.apache.iotdb.db.storageengine.rescon.disk.TierManager;
import org.apache.iotdb.db.storageengine.rescon.memory.SystemInfo;
@@ -97,6 +99,7 @@ import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
import java.net.URL;
+import java.nio.ByteBuffer;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
@@ -1066,6 +1069,35 @@ public class StorageEngine implements IService {
return RpcUtils.SUCCESS_STATUS;
}
+ public TSStatus writeLoadTsFileNodeSlice(
+ final DataRegionId dataRegionId,
+ final ByteBuffer body,
+ final String uuid,
+ final int sliceIndex,
+ final int sliceCount,
+ final int originBodySize) {
+ final LoadTsFilePieceNodeAssembler.Result result =
+ loadTsFileManager.appendPieceNodeSlice(
+ dataRegionId, uuid, body, sliceIndex, sliceCount, originBodySize);
+ if (!result.isValid()) {
+ return RpcUtils.getStatus(
+ TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR,
result.getErrorMessage());
+ }
+ if (!result.isComplete()) {
+ return RpcUtils.SUCCESS_STATUS;
+ }
+
+ try {
+ final Object planNode = PlanNodeType.deserialize(result.getBody());
+ if (!(planNode instanceof LoadTsFilePieceNode)) {
+ return new
TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode());
+ }
+ return writeLoadTsFileNode(dataRegionId, (LoadTsFilePieceNode) planNode,
uuid);
+ } catch (final Exception e) {
+ return new
TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode());
+ }
+ }
+
public TSStatus executeLoadCommand(
LoadTsFileScheduler.LoadCommand loadCommand,
String uuid,
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java
index 1be41760e8e..43774b4ac66 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFileManager.java
@@ -23,6 +23,7 @@ import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
import org.apache.iotdb.common.rpc.thrift.TTimePartitionSlot;
import org.apache.iotdb.commons.conf.IoTDBConstant;
import org.apache.iotdb.commons.consensus.ConsensusGroupId;
+import org.apache.iotdb.commons.consensus.DataRegionId;
import org.apache.iotdb.commons.consensus.index.ProgressIndex;
import org.apache.iotdb.commons.consensus.index.impl.MinimumProgressIndex;
import org.apache.iotdb.commons.disk.FolderManager;
@@ -77,6 +78,7 @@ import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.nio.file.DirectoryNotEmptyException;
import java.nio.file.Files;
import java.nio.file.Path;
@@ -123,6 +125,9 @@ public class LoadTsFileManager {
private final Map<String, TsFileWriterManager> uuid2WriterManager = new
ConcurrentHashMap<>();
+ private final Map<String, Map<DataRegionId, LoadTsFilePieceNodeAssembler>>
+ uuid2PieceNodeAssembler = new ConcurrentHashMap<>();
+
private final Map<String, CleanupTask> uuid2CleanupTask = new
ConcurrentHashMap<>();
private final PriorityBlockingQueue<CleanupTask> cleanupTaskQueue = new
PriorityBlockingQueue<>();
@@ -145,6 +150,7 @@ public class LoadTsFileManager {
cleanupTaskQueue.clear();
}
new
HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager);
+ uuid2PieceNodeAssembler.clear();
}
private long getCleanupTaskDelayInMs() {
@@ -301,6 +307,59 @@ public class LoadTsFileManager {
}
}
+ public LoadTsFilePieceNodeAssembler.Result appendPieceNodeSlice(
+ final DataRegionId dataRegionId,
+ final String uuid,
+ final ByteBuffer body,
+ final int sliceIndex,
+ final int sliceCount,
+ final int originBodySize) {
+ createCleanupTaskIfAbsent(uuid);
+
+ final Optional<CleanupTask> cleanupTask =
Optional.ofNullable(uuid2CleanupTask.get(uuid));
+ cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning);
+ try {
+ final Map<DataRegionId, LoadTsFilePieceNodeAssembler> regionId2Assembler
=
+ uuid2PieceNodeAssembler.computeIfAbsent(uuid, key -> new
ConcurrentHashMap<>());
+ synchronized (regionId2Assembler) {
+ final LoadTsFilePieceNodeAssembler assembler;
+ if (sliceIndex == 0) {
+ assembler = new LoadTsFilePieceNodeAssembler(sliceCount,
originBodySize);
+ regionId2Assembler.put(dataRegionId, assembler);
+ } else {
+ assembler = regionId2Assembler.get(dataRegionId);
+ if (assembler == null) {
+ removePieceNodeAssemblerIfEmpty(uuid, regionId2Assembler);
+ return LoadTsFilePieceNodeAssembler.Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_MISSING_LOAD_TSFILE_ASSEMBLER_FOR_UUID_ARG_DATAREGION_ARG_SLICEINDEX_ARG_FF6EA463,
+ uuid,
+ dataRegionId,
+ sliceIndex));
+ }
+ }
+
+ final LoadTsFilePieceNodeAssembler.Result result =
+ assembler.append(body, sliceIndex, sliceCount, originBodySize);
+ if (!result.isValid() || result.isComplete()) {
+ regionId2Assembler.remove(dataRegionId, assembler);
+ removePieceNodeAssemblerIfEmpty(uuid, regionId2Assembler);
+ }
+ return result;
+ }
+ } finally {
+ cleanupTask.ifPresent(CleanupTask::markLoadTaskNotRunning);
+ }
+ }
+
+ private void removePieceNodeAssemblerIfEmpty(
+ final String uuid, final Map<DataRegionId, LoadTsFilePieceNodeAssembler>
regionId2Assembler) {
+ if (regionId2Assembler.isEmpty()) {
+ uuid2PieceNodeAssembler.remove(uuid, regionId2Assembler);
+ }
+ }
+
private FolderManager getFolderManager() throws
DiskSpaceInsufficientException {
if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
synchronized (FOLDER_MANAGER) {
@@ -333,7 +392,7 @@ public class LoadTsFileManager {
boolean isGeneratedByPipe,
Map<TTimePartitionSlot, ProgressIndex> timePartitionProgressIndexMap)
throws IOException, LoadFileException {
- if (!uuid2WriterManager.containsKey(uuid)) {
+ if (!uuid2WriterManager.containsKey(uuid) ||
uuid2PieceNodeAssembler.containsKey(uuid)) {
return false;
}
@@ -352,7 +411,9 @@ public class LoadTsFileManager {
}
public boolean deleteAll(String uuid) {
- if (!uuid2WriterManager.containsKey(uuid)) {
+ if (!uuid2WriterManager.containsKey(uuid)
+ && !uuid2PieceNodeAssembler.containsKey(uuid)
+ && !uuid2CleanupTask.containsKey(uuid)) {
return false;
}
clean(uuid);
@@ -368,6 +429,7 @@ public class LoadTsFileManager {
}
}
+ uuid2PieceNodeAssembler.remove(uuid);
forceCloseWriterManager(uuid);
}
@@ -845,6 +907,7 @@ public class LoadTsFileManager {
} else {
LOGGER.info(StorageEngineMessages.LOAD_CLEANUP_TASK_STARTS, uuid);
try {
+ uuid2PieceNodeAssembler.remove(uuid);
forceCloseWriterManager(uuid);
} catch (Exception e) {
LOGGER.warn(StorageEngineMessages.LOAD_CLEANUP_TASK_ERROR, uuid, e);
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssembler.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssembler.java
new file mode 100644
index 00000000000..db534516042
--- /dev/null
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssembler.java
@@ -0,0 +1,170 @@
+/*
+ * 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.db.storageengine.load;
+
+import org.apache.iotdb.db.i18n.StorageEngineMessages;
+
+import org.apache.tsfile.utils.PublicBAOS;
+
+import java.nio.ByteBuffer;
+
+public final class LoadTsFilePieceNodeAssembler {
+
+ private final int sliceCount;
+ private final int originBodySize;
+ private final PublicBAOS assembledBody;
+
+ private int nextSliceIndex;
+
+ LoadTsFilePieceNodeAssembler(final int sliceCount, final int originBodySize)
{
+ this.sliceCount = sliceCount;
+ this.originBodySize = originBodySize;
+ this.assembledBody = new PublicBAOS(Math.max(0, originBodySize));
+ }
+
+ synchronized Result append(
+ final ByteBuffer sliceBody,
+ final int sliceIndex,
+ final int requestSliceCount,
+ final int requestOriginBodySize) {
+ if (sliceBody == null || !sliceBody.hasRemaining()) {
+ return Result.invalid(
+
StorageEngineMessages.MESSAGE_LOAD_TSFILE_SLICE_BODY_IS_NULL_OR_EMPTY_2A65366C);
+ }
+ if (sliceCount <= 1 || originBodySize <= 0) {
+ return Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_INVALID_LOAD_TSFILE_SLICE_METADATA_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_379BF1B8,
+ sliceCount,
+ originBodySize));
+ }
+ if (sliceCount != requestSliceCount || originBodySize !=
requestOriginBodySize) {
+ return Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_SLICE_METADATA_CHANGED_SLICECOUNT_ARG_EXPECTED_ARG_ORIGINBODYSIZE_ARG_EXPECTED_ARG_B16B3122,
+ requestSliceCount,
+ sliceCount,
+ requestOriginBodySize,
+ originBodySize));
+ }
+ if (sliceIndex != nextSliceIndex || sliceIndex < 0 || sliceIndex >=
sliceCount) {
+ return Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_UNEXPECTED_LOAD_TSFILE_SLICE_INDEX_ARG_EXPECTED_ARG_SLICECOUNT_ARG_76260F62,
+ sliceIndex,
+ nextSliceIndex,
+ sliceCount));
+ }
+ if (assembledBody.size() > originBodySize - sliceBody.remaining()) {
+ return Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_SLICE_EXCEEDS_ORIGINBODYSIZE_ASSEMBLEDSIZE_ARG_SLICESIZE_ARG_ORIGINBODYSIZE_ARG_0198E0D0,
+ assembledBody.size(),
+ sliceBody.remaining(),
+ originBodySize));
+ }
+
+ final ByteBuffer duplicatedBody = sliceBody.duplicate();
+ if (duplicatedBody.hasArray()) {
+ assembledBody.write(
+ duplicatedBody.array(),
+ duplicatedBody.arrayOffset() + duplicatedBody.position(),
+ duplicatedBody.remaining());
+ } else {
+ // Bulk reads avoid a ByteBuffer get and a stream capacity check for
every byte.
+ final byte[] bytes = new byte[Math.min(duplicatedBody.remaining(),
8192)];
+ while (duplicatedBody.hasRemaining()) {
+ final int size = Math.min(duplicatedBody.remaining(), bytes.length);
+ duplicatedBody.get(bytes, 0, size);
+ assembledBody.write(bytes, 0, size);
+ }
+ }
+ nextSliceIndex++;
+
+ if (nextSliceIndex < sliceCount) {
+ return assembledBody.size() < originBodySize
+ ? Result.incomplete()
+ : Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_BODY_COMPLETED_BEFORE_THE_LAST_SLICE_RECEIVED_ARG_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_0425EA2C,
+ nextSliceIndex,
+ sliceCount,
+ originBodySize));
+ }
+ if (assembledBody.size() != originBodySize) {
+ return Result.invalid(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_BODY_SIZE_MISMATCH_ASSEMBLEDSIZE_ARG_ORIGINBODYSIZE_ARG_5733FC4B,
+ assembledBody.size(),
+ originBodySize));
+ }
+ return Result.complete(
+ ByteBuffer.wrap(assembledBody.getBuf(), 0,
assembledBody.size()).asReadOnlyBuffer());
+ }
+
+ public static final class Result {
+
+ private static final Result INCOMPLETE = new Result(true, null, null);
+
+ private final boolean valid;
+ private final ByteBuffer body;
+ private final String errorMessage;
+
+ private Result(final boolean valid, final ByteBuffer body, final String
errorMessage) {
+ this.valid = valid;
+ this.body = body;
+ this.errorMessage = errorMessage;
+ }
+
+ static Result incomplete() {
+ return INCOMPLETE;
+ }
+
+ static Result invalid(final String errorMessage) {
+ return new Result(false, null, errorMessage);
+ }
+
+ static Result complete(final ByteBuffer body) {
+ return new Result(true, body, null);
+ }
+
+ public boolean isValid() {
+ return valid;
+ }
+
+ public boolean isComplete() {
+ return body != null;
+ }
+
+ public ByteBuffer getBody() {
+ return body;
+ }
+
+ public String getErrorMessage() {
+ return errorMessage;
+ }
+ }
+}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImplLoadTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImplLoadTest.java
new file mode 100644
index 00000000000..7b54fcaec20
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImplLoadTest.java
@@ -0,0 +1,94 @@
+/*
+ * 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.db.protocol.thrift.impl;
+
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
+import org.apache.iotdb.db.i18n.DataNodeMiscMessages;
+import org.apache.iotdb.mpp.rpc.thrift.TLoadResp;
+import org.apache.iotdb.mpp.rpc.thrift.TTsFilePieceReq;
+import org.apache.iotdb.rpc.TSStatusCode;
+
+import org.junit.After;
+import org.junit.Assert;
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
+
+public class DataNodeInternalRPCServiceImplLoadTest {
+
+ private int originalDataNodeId;
+
+ @Before
+ public void setUp() {
+ originalDataNodeId =
IoTDBDescriptor.getInstance().getConfig().getDataNodeId();
+ IoTDBDescriptor.getInstance().getConfig().setDataNodeId(0);
+ }
+
+ @After
+ public void tearDown() {
+
IoTDBDescriptor.getInstance().getConfig().setDataNodeId(originalDataNodeId);
+ }
+
+ @Test
+ public void testMissingSliceMetadataReportsEveryMissingField() {
+ final DataNodeInternalRPCServiceImpl service =
+ Mockito.mock(DataNodeInternalRPCServiceImpl.class,
Mockito.CALLS_REAL_METHODS);
+ for (int mask = 1; mask < 7; mask++) {
+ final TTsFilePieceReq request =
+ new TTsFilePieceReq(
+ ByteBuffer.wrap(new byte[] {1}),
+ "test-uuid",
+ new TConsensusGroupId(TConsensusGroupType.DataRegion, 1));
+ final List<String> missingFields = new ArrayList<>();
+ if ((mask & 1) != 0) {
+ request.setSliceIndex(0);
+ } else {
+ missingFields.add("sliceIndex");
+ }
+ if ((mask & 2) != 0) {
+ request.setSliceCount(2);
+ } else {
+ missingFields.add("sliceCount");
+ }
+ if ((mask & 4) != 0) {
+ request.setOriginBodySize(2);
+ } else {
+ missingFields.add("originBodySize");
+ }
+ final TLoadResp response = service.sendTsFilePieceNode(request);
+ Assert.assertFalse(response.isAccepted());
+ Assert.assertEquals(
+ TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode(),
+ response.getStatus().getCode());
+ Assert.assertEquals(
+ String.format(
+
DataNodeMiscMessages.MESSAGE_MISSING_LOAD_TSFILE_SLICE_METADATA_ARG_DE4333DA,
+ String.join(", ", missingFields)),
+ response.getStatus().getMessage());
+ Assert.assertEquals(response.getStatus().getMessage(),
response.getMessage());
+ }
+ }
+}
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImplTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImplTest.java
index 2da982fd23d..8e36a2e6bef 100644
---
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImplTest.java
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/plan/scheduler/load/LoadTsFileDispatcherImplTest.java
@@ -20,21 +20,37 @@
package org.apache.iotdb.db.queryengine.plan.scheduler.load;
import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
+import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
import org.apache.iotdb.common.rpc.thrift.TDataNodeLocation;
import org.apache.iotdb.common.rpc.thrift.TEndPoint;
import org.apache.iotdb.common.rpc.thrift.TRegionReplicaSet;
+import org.apache.iotdb.commons.client.IClientManager;
+import org.apache.iotdb.commons.client.sync.SyncDataNodeInternalServiceClient;
import org.apache.iotdb.commons.consensus.DataRegionId;
import org.apache.iotdb.commons.partition.StorageExecutor;
import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.db.conf.IoTDBConfig;
+import org.apache.iotdb.db.conf.IoTDBDescriptor;
import org.apache.iotdb.db.queryengine.common.PlanFragmentId;
import org.apache.iotdb.db.queryengine.plan.planner.plan.FragmentInstance;
import org.apache.iotdb.db.queryengine.plan.planner.plan.PlanFragment;
import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode;
import org.apache.iotdb.db.storageengine.StorageEngine;
+import org.apache.iotdb.mpp.rpc.thrift.IDataNodeRPCService;
+import org.apache.iotdb.mpp.rpc.thrift.TLoadResp;
+import org.apache.iotdb.mpp.rpc.thrift.TTsFilePieceReq;
import org.apache.iotdb.rpc.RpcUtils;
+import org.apache.iotdb.rpc.TElasticFramedTransport;
+import org.apache.thrift.TApplicationException;
+import org.apache.thrift.TException;
+import org.apache.thrift.protocol.TBinaryProtocol;
+import org.apache.thrift.transport.TMemoryBuffer;
+import org.apache.thrift.transport.TTransportException;
+import org.junit.Assert;
import org.junit.Test;
import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.powermock.api.mockito.PowerMockito;
import org.powermock.core.classloader.annotations.PowerMockIgnore;
@@ -42,13 +58,222 @@ import
org.powermock.core.classloader.annotations.PrepareForTest;
import org.powermock.modules.junit4.PowerMockRunner;
import java.io.File;
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
@PowerMockIgnore({"com.sun.org.apache.xerces.*", "javax.xml.*", "org.xml.*",
"javax.management.*"})
@RunWith(PowerMockRunner.class)
@PrepareForTest(StorageEngine.class)
public class LoadTsFileDispatcherImplTest {
+ @Test
+ public void testLoggedOversizedFrameRequiresTwoSlices() {
+ Assert.assertEquals(
+ 2, LoadTsFileDispatcherImpl.getSliceCount(120_438_706, 64 * 1024 *
1024 - 1024));
+ }
+
+ @Test
+ public void testSplitTsFilePieceReqWithinThriftFrameSize() throws Exception {
+ final int thriftMaxFrameSize = 4096;
+ final int bodySizeLimit = thriftMaxFrameSize - 1024;
+ final byte[] body = new byte[thriftMaxFrameSize * 3];
+ for (int i = 0; i < body.length; i++) {
+ body[i] = (byte) i;
+ }
+
+ final List<TTsFilePieceReq> requests =
+ LoadTsFileDispatcherImpl.splitTsFilePieceReq(
+ ByteBuffer.wrap(body),
+ "test-uuid",
+ new TConsensusGroupId(TConsensusGroupType.DataRegion, 1),
+ bodySizeLimit);
+
+ Assert.assertEquals(4, requests.size());
+ final ByteBuffer assembledBody = ByteBuffer.allocate(body.length);
+ for (int i = 0; i < requests.size(); i++) {
+ final TTsFilePieceReq request = requests.get(i);
+ Assert.assertEquals(i, request.getSliceIndex());
+ Assert.assertEquals(requests.size(), request.getSliceCount());
+ Assert.assertEquals(body.length, request.getOriginBodySize());
+ Assert.assertTrue(request.body.remaining() <= bodySizeLimit);
+ assembledBody.put(request.body.duplicate());
+
+ final TMemoryBuffer memoryBuffer = new TMemoryBuffer(thriftMaxFrameSize);
+ final TElasticFramedTransport transport =
+ new TElasticFramedTransport(memoryBuffer, 128, thriftMaxFrameSize,
true);
+ try {
+ new IDataNodeRPCService.Client(new TBinaryProtocol(transport))
+ .send_sendTsFilePieceNode(request);
+ final int frameSize =
ByteBuffer.wrap(memoryBuffer.getArray()).getInt();
+ Assert.assertEquals(memoryBuffer.length() - Integer.BYTES, frameSize);
+ Assert.assertTrue(frameSize < thriftMaxFrameSize);
+ } finally {
+ transport.close();
+ }
+ }
+ Assert.assertArrayEquals(body, assembledBody.array());
+ }
+
+ @Test
+ public void testSmallTsFilePieceReqIsNotSliced() {
+ final List<TTsFilePieceReq> requests =
+ LoadTsFileDispatcherImpl.splitTsFilePieceReq(
+ ByteBuffer.wrap(new byte[100]),
+ "test-uuid",
+ new TConsensusGroupId(TConsensusGroupType.DataRegion, 1),
+ 1024);
+
+ Assert.assertEquals(1, requests.size());
+ Assert.assertFalse(requests.get(0).isSetSliceIndex());
+ Assert.assertFalse(requests.get(0).isSetSliceCount());
+ Assert.assertFalse(requests.get(0).isSetOriginBodySize());
+ }
+
+ @Test
+ public void testDispatchUsesEachReceiversFrameLimit() throws Exception {
+ final IoTDBConfig config = IoTDBDescriptor.getInstance().getConfig();
+ final int originalMaxFrameSize = config.getThriftMaxFrameSize();
+ final int localMaxFrameSize = 8192;
+ config.setThriftMaxFrameSize(localMaxFrameSize);
+ try {
+ final IClientManager<TEndPoint, SyncDataNodeInternalServiceClient>
clientManager =
+ Mockito.mock(IClientManager.class);
+ final byte[] body = new byte[12000];
+ for (int i = 0; i < body.length; i++) {
+ body[i] = (byte) i;
+ }
+ final LoadTsFilePieceNode pieceNode =
Mockito.mock(LoadTsFilePieceNode.class);
+
Mockito.when(pieceNode.serializeToByteBuffer()).thenReturn(ByteBuffer.wrap(body));
+ final FragmentInstance instance = createFragmentInstance(pieceNode);
+ final List<TDataNodeLocation> locations = new ArrayList<>();
+ final List<List<TTsFilePieceReq>> requestsByReceiver = new ArrayList<>();
+ final List<SyncDataNodeInternalServiceClient> clients = new
ArrayList<>();
+ for (final int receiverMaxFrameSize : new int[] {4096, 16384}) {
+ final TEndPoint endPoint = new TEndPoint("127.0.0.1",
receiverMaxFrameSize);
+ locations.add(new TDataNodeLocation().setInternalEndPoint(endPoint));
+ final SyncDataNodeInternalServiceClient client =
+ Mockito.mock(SyncDataNodeInternalServiceClient.class);
+ clients.add(client);
+ Mockito.when(clientManager.borrowClient(endPoint)).thenReturn(client);
+
Mockito.when(client.getThriftMaxFrameSize()).thenReturn(receiverMaxFrameSize);
+ final List<TTsFilePieceReq> requests = new ArrayList<>();
+ requestsByReceiver.add(requests);
+ Mockito.when(client.sendTsFilePieceNode(Mockito.any()))
+ .thenAnswer(
+ invocation -> {
+ final TTsFilePieceReq request = invocation.getArgument(0);
+ requests.add(request);
+ assertFitsBothTransports(request, localMaxFrameSize,
receiverMaxFrameSize);
+ return new TLoadResp(true);
+ });
+ }
+ instance.getRegionReplicaSet().setDataNodeLocations(locations);
+ try (LoadTsFileDispatcherImpl dispatcher =
+ new LoadTsFileDispatcherImpl(clientManager, false)) {
+ dispatcher.setUuid("test-uuid");
+ for (int attempt = 0; attempt < 2; attempt++) {
+ Assert.assertTrue(
+ dispatcher
+ .dispatch(null, Collections.singletonList(instance))
+ .get(10, TimeUnit.SECONDS)
+ .isSuccessful());
+ }
+ }
+ for (int receiver = 0; receiver < clients.size(); receiver++) {
+ Mockito.verify(clients.get(receiver),
Mockito.times(1)).getThriftMaxFrameSize();
+ final List<TTsFilePieceReq> requests =
requestsByReceiver.get(receiver);
+ final int expectedSliceCount = receiver == 0 ? 4 : 2;
+ Assert.assertEquals(expectedSliceCount * 2, requests.size());
+ for (int attempt = 0; attempt < 2; attempt++) {
+ final ByteBuffer assembled = ByteBuffer.allocate(body.length);
+ for (int slice = 0; slice < expectedSliceCount; slice++) {
+ final TTsFilePieceReq request = requests.get(attempt *
expectedSliceCount + slice);
+ Assert.assertEquals(slice, request.getSliceIndex());
+ Assert.assertEquals(expectedSliceCount, request.getSliceCount());
+ assembled.put(request.body.duplicate());
+ }
+ Assert.assertArrayEquals(body, assembled.array());
+ }
+ }
+ Mockito.verify(pieceNode, Mockito.times(2)).serializeToByteBuffer();
+ } finally {
+ config.setThriftMaxFrameSize(originalMaxFrameSize);
+ }
+ }
+
+ @Test
+ public void testOlderReceiverStillAcceptsSmallPiece() throws Exception {
+ final IClientManager<TEndPoint, SyncDataNodeInternalServiceClient>
clientManager =
+ Mockito.mock(IClientManager.class);
+ final SyncDataNodeInternalServiceClient discoveryClient =
+ Mockito.mock(SyncDataNodeInternalServiceClient.class);
+ final SyncDataNodeInternalServiceClient transferClient =
+ Mockito.mock(SyncDataNodeInternalServiceClient.class);
+ Mockito.when(clientManager.borrowClient(Mockito.any()))
+ .thenReturn(discoveryClient, transferClient);
+ Mockito.when(discoveryClient.getThriftMaxFrameSize())
+ .thenThrow(new TException(new
TApplicationException(TApplicationException.UNKNOWN_METHOD)));
+
Mockito.when(transferClient.sendTsFilePieceNode(Mockito.any())).thenReturn(new
TLoadResp(true));
+ final LoadTsFilePieceNode pieceNode =
+ new LoadTsFilePieceNode(new PlanNodeId("piece"), new
File("test.tsfile"));
+ try (LoadTsFileDispatcherImpl dispatcher = new
LoadTsFileDispatcherImpl(clientManager, false)) {
+ dispatcher.setUuid("test-uuid");
+ Assert.assertTrue(
+ dispatcher
+ .dispatch(null,
Collections.singletonList(createFragmentInstance(pieceNode)))
+ .get(10, TimeUnit.SECONDS)
+ .isSuccessful());
+ }
+ final ArgumentCaptor<TTsFilePieceReq> request =
ArgumentCaptor.forClass(TTsFilePieceReq.class);
+ Mockito.verify(transferClient).sendTsFilePieceNode(request.capture());
+ Assert.assertFalse(request.getValue().isSetSliceIndex());
+ Assert.assertEquals(pieceNode.serializeToByteBuffer(),
request.getValue().body);
+ Mockito.verify(discoveryClient,
Mockito.never()).sendTsFilePieceNode(Mockito.any());
+ }
+
+ @Test
+ public void testFrameLimitDiscoveryFailureDoesNotSendPiece() throws
Exception {
+ final IClientManager<TEndPoint, SyncDataNodeInternalServiceClient>
clientManager =
+ Mockito.mock(IClientManager.class);
+ final SyncDataNodeInternalServiceClient client =
+ Mockito.mock(SyncDataNodeInternalServiceClient.class);
+ Mockito.when(clientManager.borrowClient(Mockito.any())).thenReturn(client);
+ Mockito.when(client.getThriftMaxFrameSize()).thenThrow(new
TTransportException());
+ final LoadTsFilePieceNode pieceNode =
+ new LoadTsFilePieceNode(new PlanNodeId("piece"), new
File("test.tsfile"));
+ try (LoadTsFileDispatcherImpl dispatcher = new
LoadTsFileDispatcherImpl(clientManager, false)) {
+ dispatcher.setUuid("test-uuid");
+ Assert.assertFalse(
+ dispatcher
+ .dispatch(null,
Collections.singletonList(createFragmentInstance(pieceNode)))
+ .get(10, TimeUnit.SECONDS)
+ .isSuccessful());
+ }
+ Mockito.verify(client, Mockito.never()).sendTsFilePieceNode(Mockito.any());
+ }
+
+ private static void assertFitsBothTransports(
+ final TTsFilePieceReq request, final int senderMaxFrameSize, final int
receiverMaxFrameSize)
+ throws Exception {
+ final TMemoryBuffer wire = new TMemoryBuffer(senderMaxFrameSize);
+ try (TElasticFramedTransport sender =
+ new TElasticFramedTransport(wire, 128, senderMaxFrameSize, true);
+ TElasticFramedTransport receiver =
+ new TElasticFramedTransport(wire, 128, receiverMaxFrameSize,
true)) {
+ new IDataNodeRPCService.Client(new
TBinaryProtocol(sender)).send_sendTsFilePieceNode(request);
+ final TBinaryProtocol protocol = new TBinaryProtocol(receiver);
+ protocol.readMessageBegin();
+ final IDataNodeRPCService.sendTsFilePieceNode_args args =
+ new IDataNodeRPCService.sendTsFilePieceNode_args();
+ args.read(protocol);
+ protocol.readMessageEnd();
+ Assert.assertEquals(request, args.getReq());
+ }
+ }
+
@Test
public void testDispatchLocallyPieceNodeSkipsSerdeRoundTrip() throws
Exception {
final StorageEngine storageEngine = Mockito.mock(StorageEngine.class);
diff --git
a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssemblerTest.java
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssemblerTest.java
new file mode 100644
index 00000000000..2fa07a37673
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssemblerTest.java
@@ -0,0 +1,215 @@
+/*
+ * 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.db.storageengine.load;
+
+import org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeId;
+import
org.apache.iotdb.commons.queryengine.plan.planner.plan.node.PlanNodeType;
+import org.apache.iotdb.db.i18n.StorageEngineMessages;
+import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode;
+
+import org.junit.Assert;
+import org.junit.Test;
+
+import java.io.File;
+import java.nio.ByteBuffer;
+
+public class LoadTsFilePieceNodeAssemblerTest {
+
+ @Test
+ public void testAssembleSlices() {
+ final LoadTsFilePieceNodeAssembler assembler = new
LoadTsFilePieceNodeAssembler(3, 7);
+
+ final LoadTsFilePieceNodeAssembler.Result first =
+ assembler.append(ByteBuffer.wrap(new byte[] {0, 1, 2}), 0, 3, 7);
+ Assert.assertTrue(first.isValid());
+ Assert.assertFalse(first.isComplete());
+
+ final ByteBuffer secondBody = ByteBuffer.wrap(new byte[] {9, 3, 4, 9});
+ secondBody.position(1);
+ secondBody.limit(3);
+ final LoadTsFilePieceNodeAssembler.Result second =
assembler.append(secondBody, 1, 3, 7);
+ Assert.assertTrue(second.isValid());
+ Assert.assertFalse(second.isComplete());
+
+ final LoadTsFilePieceNodeAssembler.Result last =
+ assembler.append(ByteBuffer.wrap(new byte[] {5, 6}), 2, 3, 7);
+ Assert.assertTrue(last.isValid());
+ Assert.assertTrue(last.isComplete());
+
+ final byte[] assembled = new byte[last.getBody().remaining()];
+ last.getBody().get(assembled);
+ Assert.assertArrayEquals(new byte[] {0, 1, 2, 3, 4, 5, 6}, assembled);
+ }
+
+ @Test
+ public void testRejectOutOfOrderSlice() {
+ final LoadTsFilePieceNodeAssembler.Result result =
+ new LoadTsFilePieceNodeAssembler(2, 2).append(ByteBuffer.wrap(new
byte[] {1}), 1, 2, 2);
+
+ Assert.assertFalse(result.isValid());
+ Assert.assertFalse(result.isComplete());
+ Assert.assertEquals(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_UNEXPECTED_LOAD_TSFILE_SLICE_INDEX_ARG_EXPECTED_ARG_SLICECOUNT_ARG_76260F62,
+ 1,
+ 0,
+ 2),
+ result.getErrorMessage());
+ }
+
+ @Test
+ public void testRejectMismatchedOriginBodySize() {
+ final LoadTsFilePieceNodeAssembler assembler = new
LoadTsFilePieceNodeAssembler(2, 2);
+ Assert.assertTrue(assembler.append(ByteBuffer.wrap(new byte[] {0}), 0, 2,
2).isValid());
+
+ final LoadTsFilePieceNodeAssembler.Result result =
+ assembler.append(ByteBuffer.wrap(new byte[] {1}), 1, 2, 3);
+ Assert.assertFalse(result.isValid());
+ Assert.assertFalse(result.isComplete());
+ Assert.assertEquals(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_SLICE_METADATA_CHANGED_SLICECOUNT_ARG_EXPECTED_ARG_ORIGINBODYSIZE_ARG_EXPECTED_ARG_B16B3122,
+ 2,
+ 2,
+ 3,
+ 2),
+ result.getErrorMessage());
+ }
+
+ @Test
+ public void testAssembleDirectAndReadOnlySlicesPreservesInputPositions() {
+ final byte[] expected = new byte[20000];
+ for (int i = 0; i < expected.length; i++) {
+ expected[i] = (byte) i;
+ }
+ for (boolean direct : new boolean[] {false, true}) {
+ final LoadTsFilePieceNodeAssembler assembler =
+ new LoadTsFilePieceNodeAssembler(2, expected.length);
+ final ByteBuffer backing =
+ direct ? ByteBuffer.allocateDirect(17002) :
ByteBuffer.allocate(17002);
+ backing.position(1);
+ backing.put(expected, 0, 17000);
+ backing.flip();
+ backing.position(1);
+ final ByteBuffer slice = backing.asReadOnlyBuffer();
+ Assert.assertFalse(slice.hasArray());
+ Assert.assertTrue(assembler.append(slice, 0, 2,
expected.length).isValid());
+ Assert.assertEquals(1, slice.position());
+ Assert.assertEquals(17001, slice.limit());
+ final LoadTsFilePieceNodeAssembler.Result result =
+ assembler.append(ByteBuffer.wrap(expected, 17000, 3000), 1, 2,
expected.length);
+ Assert.assertTrue(result.isComplete());
+ Assert.assertTrue(result.getBody().isReadOnly());
+ final byte[] actual = new byte[result.getBody().remaining()];
+ result.getBody().get(actual);
+ Assert.assertArrayEquals(expected, actual);
+ }
+ }
+
+ @Test
+ public void testRejectInvalidBodyAndMetadataWithReason() {
+ for (ByteBuffer body : new ByteBuffer[] {null, ByteBuffer.allocate(0)}) {
+ final LoadTsFilePieceNodeAssembler.Result result =
+ new LoadTsFilePieceNodeAssembler(2, 2).append(body, 0, 2, 2);
+ Assert.assertFalse(result.isValid());
+ Assert.assertEquals(
+
StorageEngineMessages.MESSAGE_LOAD_TSFILE_SLICE_BODY_IS_NULL_OR_EMPTY_2A65366C,
+ result.getErrorMessage());
+ }
+ for (int[] metadata : new int[][] {{1, 2}, {0, 2}, {2, 0}, {2, -1}}) {
+ final LoadTsFilePieceNodeAssembler.Result result =
+ new LoadTsFilePieceNodeAssembler(metadata[0], metadata[1])
+ .append(ByteBuffer.wrap(new byte[] {1}), 0, metadata[0],
metadata[1]);
+ Assert.assertFalse(result.isValid());
+ Assert.assertEquals(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_INVALID_LOAD_TSFILE_SLICE_METADATA_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_379BF1B8,
+ metadata[0],
+ metadata[1]),
+ result.getErrorMessage());
+ }
+ }
+
+ @Test
+ public void testRejectBodySizeViolationsWithReason() {
+ final LoadTsFilePieceNodeAssembler.Result overflow =
+ new LoadTsFilePieceNodeAssembler(2, 2).append(ByteBuffer.allocate(3),
0, 2, 2);
+ Assert.assertFalse(overflow.isValid());
+ Assert.assertEquals(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_SLICE_EXCEEDS_ORIGINBODYSIZE_ASSEMBLEDSIZE_ARG_SLICESIZE_ARG_ORIGINBODYSIZE_ARG_0198E0D0,
+ 0,
+ 3,
+ 2),
+ overflow.getErrorMessage());
+
+ final LoadTsFilePieceNodeAssembler.Result earlyCompletion =
+ new LoadTsFilePieceNodeAssembler(2, 2).append(ByteBuffer.allocate(2),
0, 2, 2);
+ Assert.assertFalse(earlyCompletion.isValid());
+ Assert.assertEquals(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_BODY_COMPLETED_BEFORE_THE_LAST_SLICE_RECEIVED_ARG_SLICECOUNT_ARG_ORIGINBODYSIZE_ARG_0425EA2C,
+ 1,
+ 2,
+ 2),
+ earlyCompletion.getErrorMessage());
+
+ final LoadTsFilePieceNodeAssembler assembler = new
LoadTsFilePieceNodeAssembler(2, 3);
+ Assert.assertTrue(assembler.append(ByteBuffer.allocate(1), 0, 2,
3).isValid());
+ final LoadTsFilePieceNodeAssembler.Result incomplete =
+ assembler.append(ByteBuffer.allocate(1), 1, 2, 3);
+ Assert.assertFalse(incomplete.isValid());
+ Assert.assertEquals(
+ String.format(
+ StorageEngineMessages
+
.MESSAGE_LOAD_TSFILE_BODY_SIZE_MISMATCH_ASSEMBLEDSIZE_ARG_ORIGINBODYSIZE_ARG_5733FC4B,
+ 2,
+ 3),
+ incomplete.getErrorMessage());
+ }
+
+ @Test
+ public void testAssembledBodyCanDeserializeLoadTsFilePieceNode() {
+ final LoadTsFilePieceNode pieceNode =
+ new LoadTsFilePieceNode(new PlanNodeId("piece"), new
File("test.tsfile"));
+ final ByteBuffer body = pieceNode.serializeToByteBuffer();
+ final int firstSliceSize = body.remaining() / 2;
+ final LoadTsFilePieceNodeAssembler assembler =
+ new LoadTsFilePieceNodeAssembler(2, body.remaining());
+
+ final ByteBuffer firstSlice = body.duplicate();
+ firstSlice.limit(firstSlice.position() + firstSliceSize);
+ Assert.assertFalse(assembler.append(firstSlice.slice(), 0, 2,
body.remaining()).isComplete());
+
+ final ByteBuffer lastSlice = body.duplicate();
+ lastSlice.position(lastSlice.position() + firstSliceSize);
+ final LoadTsFilePieceNodeAssembler.Result result =
+ assembler.append(lastSlice.slice(), 1, 2, body.remaining());
+
+ Assert.assertTrue(result.isValid());
+ Assert.assertTrue(result.isComplete());
+ Assert.assertEquals(pieceNode, PlanNodeType.deserialize(result.getBody()));
+ }
+}
diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
index c5c2bdac4a6..c20c1099888 100644
--- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
+++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
@@ -399,6 +399,9 @@ struct TTsFilePieceReq {
1: required binary body
2: required string uuid
3: required common.TConsensusGroupId consensusGroupId
+ 4: optional i32 sliceIndex
+ 5: optional i32 sliceCount
+ 6: optional i32 originBodySize
}
struct TLoadCommandReq {
@@ -920,6 +923,8 @@ service IDataNodeRPCService {
TSchemaFetchResponse fetchSchema(TSchemaFetchRequest req);
+ i32 getThriftMaxFrameSize();
+
TLoadResp sendTsFilePieceNode(TTsFilePieceReq req);
TLoadResp sendLoadCommand(TLoadCommandReq req);
