From dc324a7fe64f00d08457194887dfaac9ab6005f6 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 11 Sep 2026 18:56:10 +0800 Subject: [PATCH 1/2] Fix oversized Load TsFile piece dispatch --- .../impl/DataNodeInternalRPCServiceImpl.java | 22 +++- .../load/LoadTsFileDispatcherImpl.java | 102 ++++++++++++--- .../iotdb/db/storageengine/StorageEngine.java | 31 +++++ .../storageengine/load/LoadTsFileManager.java | 61 ++++++++- .../load/LoadTsFilePieceNodeAssembler.java | 120 ++++++++++++++++++ .../load/LoadTsFileDispatcherImplTest.java | 72 +++++++++++ .../LoadTsFilePieceNodeAssemblerTest.java | 102 +++++++++++++++ .../src/main/thrift/datanode.thrift | 3 + 8 files changed, 490 insertions(+), 23 deletions(-) create mode 100644 iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssembler.java create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssemblerTest.java 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 34ef951296ce8..9784a1adbc19e 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 @@ -629,10 +629,30 @@ public TSchemaFetchResponse fetchSchema(final TSchemaFetchRequest req) { @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()) { + return createTLoadResp( + new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode())); + } + 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 1ba39ffd0bd3b..9f92886bafee6 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; @@ -63,6 +64,7 @@ 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; @@ -82,6 +84,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()); @@ -143,7 +146,7 @@ public Future dispatch( private void dispatchOneInstance(FragmentInstance instance) throws FragmentInstanceDispatchException { - TTsFilePieceReq loadTsFileReq = null; + List loadTsFileReqs = null; for (TDataNodeLocation dataNodeLocation : instance.getRegionReplicaSet().getDataNodeLocations()) { @@ -151,18 +154,75 @@ private void dispatchOneInstance(FragmentInstance instance) if (isDispatchedToLocal(endPoint)) { dispatchLocally(instance); } else { - if (loadTsFileReq == null) { - loadTsFileReq = - new TTsFilePieceReq( + if (loadTsFileReqs == null) { + loadTsFileReqs = + splitTsFilePieceReq( instance.getFragment().getPlanNodeTree().serializeToByteBuffer(), uuid, - instance.getRegionReplicaSet().getRegionId()); + instance.getRegionReplicaSet().getRegionId(), + getLoadTsFilePieceBodySizeLimit()); } - dispatchRemote(loadTsFileReq, endPoint); + dispatchRemote(loadTsFileReqs, endPoint); } } } + private static int getLoadTsFilePieceBodySizeLimit() { + final int thriftMaxFrameSize = + IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize(); + return Math.max(1, thriftMaxFrameSize - LOAD_TSFILE_PIECE_RPC_FRAME_RESERVED_BYTES); + } + + static List 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 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,25 +282,27 @@ public void dispatchLocally(FragmentInstance instance) throws FragmentInstanceDi } } - private void dispatchRemote(TTsFilePieceReq loadTsFileReq, TEndPoint endPoint) + private void dispatchRemote(List loadTsFileReqs, 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); - transferAttemptRecorded = true; - LOGGER.warn(loadResp.message); - throw new FragmentInstanceDispatchException(loadResp.status); + 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; @@ -253,7 +315,7 @@ private void dispatchRemote(TTsFilePieceReq loadTsFileReq, TEndPoint endPoint) 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 1ac6c15488b11..f00d786c016e8 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.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.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 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,34 @@ public TSStatus writeLoadTsFileNode( 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 new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode()); + } + 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 1be41760e8e30..f51afadb3bb73 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.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 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 uuid2WriterManager = new ConcurrentHashMap<>(); + private final Map> + uuid2PieceNodeAssembler = new ConcurrentHashMap<>(); + private final Map uuid2CleanupTask = new ConcurrentHashMap<>(); private final PriorityBlockingQueue cleanupTaskQueue = new PriorityBlockingQueue<>(); @@ -145,6 +150,7 @@ public void stop() { cleanupTaskQueue.clear(); } new HashSet<>(uuid2WriterManager.keySet()).forEach(this::forceCloseWriterManager); + uuid2PieceNodeAssembler.clear(); } private long getCleanupTaskDelayInMs() { @@ -301,6 +307,53 @@ public void writeToDataRegion(DataRegion dataRegion, LoadTsFilePieceNode pieceNo } } + 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 = Optional.ofNullable(uuid2CleanupTask.get(uuid)); + cleanupTask.ifPresent(CleanupTask::markLoadTaskRunning); + try { + final Map 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(); + } + } + + 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 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 +386,7 @@ public boolean loadAll( boolean isGeneratedByPipe, Map timePartitionProgressIndexMap) throws IOException, LoadFileException { - if (!uuid2WriterManager.containsKey(uuid)) { + if (!uuid2WriterManager.containsKey(uuid) || uuid2PieceNodeAssembler.containsKey(uuid)) { return false; } @@ -352,7 +405,9 @@ public boolean loadAll( } 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 +423,7 @@ private void clean(String uuid) { } } + uuid2PieceNodeAssembler.remove(uuid); forceCloseWriterManager(uuid); } @@ -845,6 +901,7 @@ public void run() { } 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 0000000000000..9b80c236077c7 --- /dev/null +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssembler.java @@ -0,0 +1,120 @@ +/* + * 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 = new PublicBAOS(); + + private int nextSliceIndex; + + LoadTsFilePieceNodeAssembler(final int sliceCount, final int originBodySize) { + this.sliceCount = sliceCount; + this.originBodySize = originBodySize; + } + + synchronized Result append( + final ByteBuffer sliceBody, + final int sliceIndex, + final int requestSliceCount, + final int requestOriginBodySize) { + if (sliceBody == null + || !sliceBody.hasRemaining() + || sliceCount <= 1 + || originBodySize <= 0 + || sliceCount != requestSliceCount + || originBodySize != requestOriginBodySize + || sliceIndex != nextSliceIndex + || sliceIndex < 0 + || sliceIndex >= sliceCount + || assembledBody.size() > originBodySize - sliceBody.remaining()) { + return Result.invalid(); + } + + final ByteBuffer duplicatedBody = sliceBody.duplicate(); + if (duplicatedBody.hasArray()) { + assembledBody.write( + duplicatedBody.array(), + duplicatedBody.arrayOffset() + duplicatedBody.position(), + duplicatedBody.remaining()); + } else { + 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(); + } + if (assembledBody.size() != originBodySize) { + return Result.invalid(); + } + 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); + private static final Result INVALID = new Result(false, null); + + private final boolean valid; + private final ByteBuffer body; + + private Result(final boolean valid, final ByteBuffer body) { + this.valid = valid; + this.body = body; + } + + static Result incomplete() { + return INCOMPLETE; + } + + static Result invalid() { + return INVALID; + } + + static Result complete(final ByteBuffer body) { + return new Result(true, body); + } + + public boolean isValid() { + return valid; + } + + public boolean isComplete() { + return body != null; + } + + public ByteBuffer getBody() { + return body; + } + } +} 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 2da982fd23da9..834bb43cd4864 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,6 +20,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.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; @@ -31,8 +32,14 @@ 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.TTsFilePieceReq; import org.apache.iotdb.rpc.RpcUtils; +import org.apache.iotdb.rpc.TElasticFramedTransport; +import org.apache.thrift.protocol.TBinaryProtocol; +import org.apache.thrift.transport.TMemoryBuffer; +import org.junit.Assert; import org.junit.Test; import org.junit.runner.RunWith; import org.mockito.Mockito; @@ -42,13 +49,78 @@ import org.powermock.modules.junit4.PowerMockRunner; import java.io.File; +import java.nio.ByteBuffer; import java.util.Collections; +import java.util.List; @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 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 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 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 0000000000000..918cc90d89527 --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/storageengine/load/LoadTsFilePieceNodeAssemblerTest.java @@ -0,0 +1,102 @@ +/* + * 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.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()); + } + + @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()); + } + + @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 c5c2bdac4a656..4d57ffd69e78b 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 { From fdf4ac6dd412b205de1c92e39f38d93adab06e91 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 18 Sep 2026 10:34:37 +0800 Subject: [PATCH 2/2] Address Load TsFile slicing review feedback --- .../iotdb/db/i18n/DataNodeMiscMessages.java | 3 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 5 + .../iotdb/db/i18n/StorageEngineMessages.java | 17 ++ .../iotdb/db/i18n/DataNodeMiscMessages.java | 3 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 5 + .../iotdb/db/i18n/StorageEngineMessages.java | 17 ++ .../impl/DataNodeInternalRPCServiceImpl.java | 21 ++- .../load/LoadTsFileDispatcherImpl.java | 121 +++++++++----- .../iotdb/db/storageengine/StorageEngine.java | 3 +- .../storageengine/load/LoadTsFileManager.java | 8 +- .../load/LoadTsFilePieceNodeAssembler.java | 90 ++++++++--- ...ataNodeInternalRPCServiceImplLoadTest.java | 94 +++++++++++ .../load/LoadTsFileDispatcherImplTest.java | 153 ++++++++++++++++++ .../LoadTsFilePieceNodeAssemblerTest.java | 113 +++++++++++++ .../src/main/thrift/datanode.thrift | 2 + 15 files changed, 595 insertions(+), 60 deletions(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/protocol/thrift/impl/DataNodeInternalRPCServiceImplLoadTest.java 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 12f4e374ca7a7..b571df825ab55 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 @@ /** 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 bf7a36010030a..e4cd4bff764b9 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 @@ 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 ccde46c4881b9..3d66a88869672 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 @@ 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 d809fea199259..bf15d6cde078b 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 @@ /** 编译时国际化常量 - 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 f26f2306be2fa..8dca0e0284386 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 @@ 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 d8692a3573b3a..cf7f5f01309ba 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 @@ 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 9784a1adbc19e..e7213b7822c91 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,6 +627,11 @@ public TSchemaFetchResponse fetchSchema(final TSchemaFetchRequest req) { throw new UnsupportedOperationException(); } + @Override + public int getThriftMaxFrameSize() { + return IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize(); + } + @Override public TLoadResp sendTsFilePieceNode(final TTsFilePieceReq req) { if (!req.isSetSliceIndex() || req.sliceIndex == 0) { @@ -639,8 +644,22 @@ public TLoadResp sendTsFilePieceNode(final TTsFilePieceReq req) { req.isSetSliceIndex() || req.isSetSliceCount() || req.isSetOriginBodySize(); if (isSliced) { if (!req.isSetSliceIndex() || !req.isSetSliceCount() || !req.isSetOriginBodySize()) { + final List missingFields = new ArrayList<>(3); + if (!req.isSetSliceIndex()) { + missingFields.add("sliceIndex"); + } + if (!req.isSetSliceCount()) { + missingFields.add("sliceCount"); + } + if (!req.isSetOriginBodySize()) { + missingFields.add("originBodySize"); + } return createTLoadResp( - new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode())); + 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() 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 9f92886bafee6..80f09f17e2872 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 @@ -58,6 +58,7 @@ 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; @@ -71,6 +72,7 @@ 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; @@ -95,6 +97,7 @@ public class LoadTsFileDispatcherImpl implements IFragInstanceDispatcher, AutoCl internalServiceClientManager; private ExecutorService executor; private final boolean isGeneratedByPipe; + private final Map endPoint2ThriftMaxFrameSize = new ConcurrentHashMap<>(); public LoadTsFileDispatcherImpl( IClientManager internalServiceClientManager, @@ -146,7 +149,7 @@ public Future dispatch( private void dispatchOneInstance(FragmentInstance instance) throws FragmentInstanceDispatchException { - List loadTsFileReqs = null; + ByteBuffer body = null; for (TDataNodeLocation dataNodeLocation : instance.getRegionReplicaSet().getDataNodeLocations()) { @@ -154,23 +157,55 @@ private void dispatchOneInstance(FragmentInstance instance) if (isDispatchedToLocal(endPoint)) { dispatchLocally(instance); } else { - if (loadTsFileReqs == null) { - loadTsFileReqs = - splitTsFilePieceReq( - instance.getFragment().getPlanNodeTree().serializeToByteBuffer(), - uuid, - instance.getRegionReplicaSet().getRegionId(), - getLoadTsFilePieceBodySizeLimit()); + if (body == null) { + body = instance.getFragment().getPlanNodeTree().serializeToByteBuffer(); } - dispatchRemote(loadTsFileReqs, endPoint); + dispatchRemote(body, instance.getRegionReplicaSet().getRegionId(), endPoint); } } } - private static int getLoadTsFilePieceBodySizeLimit() { - final int thriftMaxFrameSize = - IoTDBDescriptor.getInstance().getConfig().getThriftMaxFrameSize(); - return Math.max(1, thriftMaxFrameSize - LOAD_TSFILE_PIECE_RPC_FRAME_RESERVED_BYTES); + 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 splitTsFilePieceReq( @@ -282,30 +317,36 @@ public void dispatchLocally(FragmentInstance instance) throws FragmentInstanceDi } } - private void dispatchRemote(List loadTsFileReqs, 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()); - - 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); + try { + final List 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; } - recordTransferAttempt(endPoint, true, null, null); - transferAttemptRecorded = true; } catch (Exception e) { if (!transferAttemptRecorded) { recordTransferAttempt(endPoint, false, null, e); @@ -314,8 +355,11 @@ private void dispatchRemote(List loadTsFileReqs, TEndPoint endP final String exceptionMessage = String.format( - "failed to dispatch load command %s to node %s because of exception: %s", - uuid, 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() @@ -423,8 +467,11 @@ private void dispatchRemote(TLoadCommandReq loadCommandReq, TEndPoint endPoint) 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 f00d786c016e8..a06ec8c00b574 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 @@ -1080,7 +1080,8 @@ public TSStatus writeLoadTsFileNodeSlice( loadTsFileManager.appendPieceNodeSlice( dataRegionId, uuid, body, sliceIndex, sliceCount, originBodySize); if (!result.isValid()) { - return new TSStatus(TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR.getStatusCode()); + return RpcUtils.getStatus( + TSStatusCode.DESERIALIZE_PIECE_OF_TSFILE_ERROR, result.getErrorMessage()); } if (!result.isComplete()) { return RpcUtils.SUCCESS_STATUS; 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 f51afadb3bb73..43774b4ac66ef 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 @@ -330,7 +330,13 @@ public LoadTsFilePieceNodeAssembler.Result appendPieceNodeSlice( assembler = regionId2Assembler.get(dataRegionId); if (assembler == null) { removePieceNodeAssemblerIfEmpty(uuid, regionId2Assembler); - return LoadTsFilePieceNodeAssembler.Result.invalid(); + return LoadTsFilePieceNodeAssembler.Result.invalid( + String.format( + StorageEngineMessages + .MESSAGE_MISSING_LOAD_TSFILE_ASSEMBLER_FOR_UUID_ARG_DATAREGION_ARG_SLICEINDEX_ARG_FF6EA463, + uuid, + dataRegionId, + sliceIndex)); } } 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 index 9b80c236077c7..db534516042c8 100644 --- 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 @@ -19,6 +19,8 @@ package org.apache.iotdb.db.storageengine.load; +import org.apache.iotdb.db.i18n.StorageEngineMessages; + import org.apache.tsfile.utils.PublicBAOS; import java.nio.ByteBuffer; @@ -27,13 +29,14 @@ public final class LoadTsFilePieceNodeAssembler { private final int sliceCount; private final int originBodySize; - private final PublicBAOS assembledBody = new PublicBAOS(); + 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( @@ -41,17 +44,45 @@ synchronized Result append( final int sliceIndex, final int requestSliceCount, final int requestOriginBodySize) { - if (sliceBody == null - || !sliceBody.hasRemaining() - || sliceCount <= 1 - || originBodySize <= 0 - || sliceCount != requestSliceCount - || originBodySize != requestOriginBodySize - || sliceIndex != nextSliceIndex - || sliceIndex < 0 - || sliceIndex >= sliceCount - || assembledBody.size() > originBodySize - sliceBody.remaining()) { - return Result.invalid(); + 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(); @@ -61,6 +92,7 @@ synchronized Result append( 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); @@ -71,10 +103,23 @@ synchronized Result append( nextSliceIndex++; if (nextSliceIndex < sliceCount) { - return assembledBody.size() < originBodySize ? Result.incomplete() : Result.invalid(); + 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(); + 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()); @@ -82,27 +127,28 @@ synchronized Result append( public static final class Result { - private static final Result INCOMPLETE = new Result(true, null); - private static final Result INVALID = new Result(false, null); + 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) { + 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() { - return INVALID; + static Result invalid(final String errorMessage) { + return new Result(false, null, errorMessage); } static Result complete(final ByteBuffer body) { - return new Result(true, body); + return new Result(true, body, null); } public boolean isValid() { @@ -116,5 +162,9 @@ public boolean isComplete() { 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 0000000000000..7b54fcaec20a6 --- /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 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 834bb43cd4864..8e36a2e6bef97 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 @@ -24,24 +24,33 @@ 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; @@ -50,8 +59,10 @@ 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) @@ -121,6 +132,148 @@ public void testSmallTsFilePieceReqIsNotSliced() { 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 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 locations = new ArrayList<>(); + final List> requestsByReceiver = new ArrayList<>(); + final List 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 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 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 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 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 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 index 918cc90d89527..2fa07a37673fb 100644 --- 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 @@ -21,6 +21,7 @@ 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; @@ -64,6 +65,14 @@ public void testRejectOutOfOrderSlice() { 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 @@ -75,6 +84,110 @@ public void testRejectMismatchedOriginBodySize() { 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 diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index 4d57ffd69e78b..c20c1099888b5 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -923,6 +923,8 @@ service IDataNodeRPCService { TSchemaFetchResponse fetchSchema(TSchemaFetchRequest req); + i32 getThriftMaxFrameSize(); + TLoadResp sendTsFilePieceNode(TTsFilePieceReq req); TLoadResp sendLoadCommand(TLoadCommandReq req);