This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch dev/1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/dev/1.3 by this push:
new 46541cccdfa Fix oversized Load TsFile piece dispatch (#18627) (#18674)
46541cccdfa is described below
commit 46541cccdfa83c5e58a3fd691bf7d381986d3679
Author: Caideyipi <[email protected]>
AuthorDate: Mon Sep 21 11:54:18 2026 +0800
Fix oversized Load TsFile piece dispatch (#18627) (#18674)
* Fix oversized Load TsFile piece dispatch
* Address Load TsFile slicing review feedback
---
.../impl/DataNodeInternalRPCServiceImpl.java | 44 +++-
.../scheduler/load/LoadTsFileDispatcherImpl.java | 136 +++++++++++--
.../iotdb/db/storageengine/StorageEngine.java | 32 +++
.../db/storageengine/load/LoadTsFileManager.java | 64 +++++-
.../load/LoadTsFilePieceNodeAssembler.java | 149 ++++++++++++++
.../DataNodeInternalRPCServiceImplLoadTest.java | 91 +++++++++
.../load/LoadTsFileDispatcherImplTest.java | 225 +++++++++++++++++++++
.../load/LoadTsFilePieceNodeAssemblerTest.java | 184 +++++++++++++++++
.../src/main/thrift/datanode.thrift | 5 +
9 files changed, 907 insertions(+), 23 deletions(-)
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 458f2279355..55ee38253ea 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
@@ -498,12 +498,50 @@ public class DataNodeInternalRPCServiceImpl implements
IDataNodeRPCService.Iface
}
@Override
- public TLoadResp sendTsFilePieceNode(TTsFilePieceReq req) {
- LOGGER.info("Receive load node from uuid {}.", req.uuid);
+ public int getThriftMaxFrameSize() {
+ return IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize();
+ }
+
+ @Override
+ public TLoadResp sendTsFilePieceNode(final TTsFilePieceReq req) {
+ if (!req.isSetSliceIndex() || req.sliceIndex == 0) {
+ LOGGER.info("Receive load node from uuid {}.", req.uuid);
+ }
ConsensusGroupId groupId =
ConsensusGroupId.Factory.createFromTConsensusGroupId(req.consensusGroupId);
- LoadTsFilePieceNode pieceNode = (LoadTsFilePieceNode)
PlanNodeType.deserialize(req.body);
+ 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(
+ "Missing Load TsFile slice metadata: %s", 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(
new
TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode()));
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 b734412ca9c..542d726186f 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;
@@ -51,17 +52,20 @@ import org.apache.iotdb.rpc.RpcUtils;
import org.apache.iotdb.rpc.TSStatusCode;
import io.airlift.concurrent.SetThreadName;
+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.Set;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeoutException;
@@ -75,6 +79,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());
@@ -85,6 +90,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,
@@ -134,7 +140,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()) {
@@ -142,18 +148,104 @@ 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(
+ "Invalid Thrift maximum frame size %d from %s",
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("Receive load node from uuid {}.", uuid);
@@ -211,16 +303,24 @@ public class LoadTsFileDispatcherImpl implements
IFragInstanceDispatcher, AutoCl
}
}
- private void dispatchRemote(TTsFilePieceReq loadTsFileReq, TEndPoint
endPoint)
+ private void dispatchRemote(
+ ByteBuffer body, TConsensusGroupId consensusGroupId, TEndPoint endPoint)
throws FragmentInstanceDispatchException {
- try (SyncDataNodeInternalServiceClient client =
- internalServiceClientManager.borrowClient(endPoint)) {
- client.setTimeout(CONNECTION_TIMEOUT_MS.get());
-
- final TLoadResp loadResp = client.sendTsFilePieceNode(loadTsFileReq);
- if (!loadResp.isAccepted()) {
- LOGGER.warn(loadResp.message);
- throw new FragmentInstanceDispatchException(loadResp.status);
+ 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()) {
+ LOGGER.warn(loadResp.message);
+ throw new FragmentInstanceDispatchException(loadResp.status);
+ }
+ }
}
} catch (Exception e) {
adjustTimeoutIfNecessary(e);
@@ -228,7 +328,7 @@ 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);
+ uuid, 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 67d4abf5907..a9d6490d5f4 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
@@ -54,6 +54,7 @@ import
org.apache.iotdb.db.exception.WriteProcessRejectException;
import org.apache.iotdb.db.exception.load.LoadReadOnlyException;
import org.apache.iotdb.db.exception.runtime.StorageEngineFailureException;
import
org.apache.iotdb.db.queryengine.plan.analyze.cache.schema.DataNodeTTLCache;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeType;
import
org.apache.iotdb.db.queryengine.plan.planner.plan.node.load.LoadTsFilePieceNode;
import org.apache.iotdb.db.queryengine.plan.scheduler.load.LoadTsFileScheduler;
import org.apache.iotdb.db.service.metrics.FileMetrics;
@@ -73,6 +74,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.memory.SystemInfo;
import org.apache.iotdb.db.utils.ThreadUtils;
@@ -88,6 +90,7 @@ import org.slf4j.LoggerFactory;
import java.io.File;
import java.io.IOException;
import java.net.URL;
+import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
@@ -976,6 +979,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 671444352f9..abbce7022a3 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.file.SystemFileFactory;
@@ -71,6 +72,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;
@@ -117,6 +119,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<>();
@@ -139,6 +144,7 @@ public class LoadTsFileManager {
cleanupTaskQueue.clear();
}
new
HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager);
+ uuid2PieceNodeAssembler.clear();
}
private long getCleanupTaskDelayInMs() {
@@ -285,6 +291,56 @@ 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(
+ "Missing Load TsFile assembler for uuid %s, DataRegion %s,
sliceIndex=%d",
+ 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 String getNextFolder() throws DiskSpaceInsufficientException {
if (CONFIG.getLoadTsFileDirs() != LOAD_BASE_DIRS.get()) {
synchronized (FOLDER_MANAGER) {
@@ -317,7 +373,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;
}
@@ -336,7 +392,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);
@@ -352,6 +410,7 @@ public class LoadTsFileManager {
}
}
+ uuid2PieceNodeAssembler.remove(uuid);
forceCloseWriterManager(uuid);
}
@@ -798,6 +857,7 @@ public class LoadTsFileManager {
} else {
LOGGER.info("Load cleanup task {} starts.", uuid);
try {
+ uuid2PieceNodeAssembler.remove(uuid);
forceCloseWriterManager(uuid);
} catch (Exception e) {
LOGGER.warn("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..ccf9d77ac8f
--- /dev/null
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssembler.java
@@ -0,0 +1,149 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.iotdb.db.storageengine.load;
+
+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("Load TsFile slice body is null or empty");
+ }
+ if (sliceCount <= 1 || originBodySize <= 0) {
+ return Result.invalid(
+ String.format(
+ "Invalid Load TsFile slice metadata: sliceCount=%d,
originBodySize=%d",
+ sliceCount, originBodySize));
+ }
+ if (sliceCount != requestSliceCount || originBodySize !=
requestOriginBodySize) {
+ return Result.invalid(
+ String.format(
+ "Load TsFile slice metadata changed: sliceCount=%d (expected
%d), originBodySize=%d (expected %d)",
+ requestSliceCount, sliceCount, requestOriginBodySize,
originBodySize));
+ }
+ if (sliceIndex != nextSliceIndex || sliceIndex < 0 || sliceIndex >=
sliceCount) {
+ return Result.invalid(
+ String.format(
+ "Unexpected Load TsFile slice index %d: expected %d,
sliceCount=%d",
+ sliceIndex, nextSliceIndex, sliceCount));
+ }
+ if (assembledBody.size() > originBodySize - sliceBody.remaining()) {
+ return Result.invalid(
+ String.format(
+ "Load TsFile slice exceeds originBodySize: assembledSize=%d,
sliceSize=%d, originBodySize=%d",
+ 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(
+ "Load TsFile body completed before the last slice:
received=%d, sliceCount=%d, originBodySize=%d",
+ nextSliceIndex, sliceCount, originBodySize));
+ }
+ if (assembledBody.size() != originBodySize) {
+ return Result.invalid(
+ String.format(
+ "Load TsFile body size mismatch: assembledSize=%d,
originBodySize=%d",
+ assembledBody.size(), originBodySize));
+ }
+ return Result.complete(ByteBuffer.wrap(assembledBody.getBuf(), 0,
assembledBody.size()));
+ }
+
+ 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..97eb4510b24
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImplLoadTest.java
@@ -0,0 +1,91 @@
+/*
+ * 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.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("Missing Load TsFile slice metadata: %s",
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 4fc0da17bb9..b9082a54c75 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.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.PlanNodeId;
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..79c572a14ed
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssemblerTest.java
@@ -0,0 +1,184 @@
+/*
+ * 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.queryengine.plan.planner.plan.node.PlanNodeId;
+import org.apache.iotdb.db.queryengine.plan.planner.plan.node.PlanNodeType;
+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(
+ "Unexpected Load TsFile slice index 1: expected 0, sliceCount=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(
+ "Load TsFile slice metadata changed: sliceCount=2 (expected 2),
originBodySize=3 (expected 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.assertFalse(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("Load TsFile slice body is null or empty",
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(
+ "Invalid Load TsFile slice metadata: sliceCount=%d,
originBodySize=%d",
+ 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(
+ "Load TsFile slice exceeds originBodySize: assembledSize=0,
sliceSize=3, originBodySize=2",
+ overflow.getErrorMessage());
+
+ final LoadTsFilePieceNodeAssembler.Result earlyCompletion =
+ new LoadTsFilePieceNodeAssembler(2, 2).append(ByteBuffer.allocate(2),
0, 2, 2);
+ Assert.assertFalse(earlyCompletion.isValid());
+ Assert.assertEquals(
+ "Load TsFile body completed before the last slice: received=1,
sliceCount=2, originBodySize=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(
+ "Load TsFile body size mismatch: assembledSize=2, originBodySize=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 4a8b5e48f1d..9be30e94fa3 100644
--- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
+++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift
@@ -333,6 +333,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 {
@@ -730,6 +733,8 @@ service IDataNodeRPCService {
TSchemaFetchResponse fetchSchema(TSchemaFetchRequest req);
+ i32 getThriftMaxFrameSize();
+
TLoadResp sendTsFilePieceNode(TTsFilePieceReq req);
TLoadResp sendLoadCommand(TLoadCommandReq req);
