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] 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 34ef951296ce..9784a1adbc19 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 1ba39ffd0bd3..9f92886bafee 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 1ac6c15488b1..f00d786c016e 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 1be41760e8e3..f51afadb3bb7 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 000000000000..9b80c236077c --- /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 2da982fd23da..834bb43cd486 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 000000000000..918cc90d8952 --- /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 c5c2bdac4a65..4d57ffd69e78 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 {