From 280c432188bf740ad779e020ee7cd2a6a2d9a86d Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Mon, 14 Sep 2026 14:36:43 +0800 Subject: [PATCH 01/12] draft Signed-off-by: Weihao Li <18110526956@163.com> --- .../org/apache/iotdb/db/conf/IoTDBConfig.java | 12 +++ .../apache/iotdb/db/conf/IoTDBDescriptor.java | 5 + .../exchange/MPPDataExchangeManager.java | 35 +++++++ .../execution/exchange/sink/SinkChannel.java | 32 +++++- .../exchange/source/SourceHandle.java | 98 ++++++++++++++++++- .../conf/iotdb-system.properties.template | 5 + .../src/main/thrift/datanode.thrift | 4 + 7 files changed, 189 insertions(+), 2 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index 989031fc0fd12..7d7f3652d38d9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -956,6 +956,8 @@ public class IoTDBConfig { /** Core pool size of mpp data exchange. */ private int mppDataExchangeCorePoolSize = 10; + private int mppDataExchangeMaxPayloadSizeInBytes = 8 * 1024 * 1024; + /** Max pool size of mpp data exchange. */ private int mppDataExchangeMaxPoolSize = 10; @@ -3421,6 +3423,16 @@ public void setMppDataExchangeKeepAliveTimeInMs(int mppDataExchangeKeepAliveTime this.mppDataExchangeKeepAliveTimeInMs = mppDataExchangeKeepAliveTimeInMs; } + public int getMppDataExchangeMaxPayloadSizeInBytes() { + return mppDataExchangeMaxPayloadSizeInBytes; + } + + public void setMppDataExchangeMaxPayloadSizeInBytes( + int mppDataExchangeMaxPayloadSizeInBytes) { + this.mppDataExchangeMaxPayloadSizeInBytes = + Math.max(1, Math.min(mppDataExchangeMaxPayloadSizeInBytes, thriftMaxFrameSize - 1024)); + } + public int getConnectionTimeoutInMS() { return connectionTimeoutInMS; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 9f48bac0e0cd5..8c3efb84b83b2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -3070,6 +3070,11 @@ public void loadShuffleProps(TrimProperties properties) { properties.getProperty( "mpp_data_exchange_keep_alive_time_in_ms", Integer.toString(conf.getMppDataExchangeKeepAliveTimeInMs())))); + conf.setMppDataExchangeMaxPayloadSizeInBytes( + Integer.parseInt( + properties.getProperty( + "mpp_data_exchange_max_payload_size_in_bytes", + Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes())))); conf.setPartitionCacheSize( Integer.parseInt( diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index ceba3880122aa..f8327673a8064 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -24,6 +24,7 @@ import org.apache.iotdb.commons.client.IClientManager; import org.apache.iotdb.commons.client.sync.SyncDataNodeMPPDataExchangeServiceClient; import org.apache.iotdb.commons.utils.TestOnly; +import org.apache.iotdb.db.conf.IoTDBDescriptor; import org.apache.iotdb.db.i18n.DataNodeQueryMessages; import org.apache.iotdb.db.queryengine.exception.exchange.GetTsBlockFromClosedOrAbortedChannelException; import org.apache.iotdb.db.queryengine.execution.driver.DriverContext; @@ -176,6 +177,40 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce } // index of the channel must be a SinkChannel SinkChannel sinkChannel = (SinkChannel) (sinkHandle.getChannel(req.getIndex())); + if (req.isSetOffset()) { + int remainingPayloadSize = + IoTDBDescriptor.getInstance() + .getConfig() + .getMppDataExchangeMaxPayloadSizeInBytes(); + long offset = req.getOffset(); + for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { + try { + ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); + long blockOffset = i == req.getStartSequenceId() ? offset : 0L; + long remainingBlockSize = serializedTsBlock.remaining() - blockOffset; + if (remainingBlockSize <= remainingPayloadSize) { + resp.addToTsBlocks( + sinkChannel.getSerializedTsBlockFragment( + i, blockOffset, Math.toIntExact(remainingBlockSize))); + remainingPayloadSize -= Math.toIntExact(remainingBlockSize); + if (remainingPayloadSize == 0) { + break; + } + } else { + resp.addToTsBlocks( + sinkChannel.getSerializedTsBlockFragment( + i, blockOffset, remainingPayloadSize)); + resp.setOffset(blockOffset + remainingPayloadSize); + break; + } + } catch (GetTsBlockFromClosedOrAbortedChannelException e) { + return new TGetDataBlockResponse(new ArrayList<>()); + } catch (IllegalArgumentException | IllegalStateException | IOException e) { + throw new TException(e); + } + } + return resp; + } for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { try { ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index daaff9fdbcd29..533387afdf01a 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -96,6 +96,10 @@ public class SinkChannel implements ISinkChannel { private final LinkedHashMap> sequenceIdToTsBlock = new LinkedHashMap<>(); + /** Serialized blocks are cached so fragmented requests do not serialize the same block again. */ + private final LinkedHashMap sequenceIdToSerializedTsBlock = + new LinkedHashMap<>(); + // size for current TsBlock to reserve and free private long currentTsBlockSize; @@ -305,6 +309,7 @@ public synchronized boolean abort() { return false; } sequenceIdToTsBlock.clear(); + sequenceIdToSerializedTsBlock.clear(); if (blocked != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked); } @@ -335,6 +340,7 @@ public synchronized boolean close() { return false; } sequenceIdToTsBlock.clear(); + sequenceIdToSerializedTsBlock.clear(); if (blocked != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked); } @@ -407,6 +413,10 @@ public synchronized ByteBuffer getSerializedTsBlock(int sequenceId) throws IOExc throw new GetTsBlockFromClosedOrAbortedChannelException( DataNodeQueryMessages.SINKCHANNEL_IS_ABORTED_OR_CLOSED); } + ByteBuffer serializedTsBlock = sequenceIdToSerializedTsBlock.get(sequenceId); + if (serializedTsBlock != null) { + return serializedTsBlock.duplicate(); + } Pair pair = sequenceIdToTsBlock.get(sequenceId); if (pair == null || pair.left == null) { LOGGER.warn( @@ -416,7 +426,26 @@ public synchronized ByteBuffer getSerializedTsBlock(int sequenceId) throws IOExc throw new IllegalStateException( DataNodeQueryMessages.THE_DATA_BLOCK_DOESN_T_EXIST_SEQUENCE_ID + sequenceId); } - return serde.serialize(pair.left); + serializedTsBlock = serde.serialize(pair.left); + sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock.asReadOnlyBuffer()); + return serializedTsBlock.duplicate(); + } + + public synchronized ByteBuffer getSerializedTsBlockFragment( + int sequenceId, long offset, int maxBytes) throws IOException { + ByteBuffer serializedTsBlock = getSerializedTsBlock(sequenceId); + if (offset < 0 || offset > serializedTsBlock.remaining() || maxBytes <= 0) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_ARG_ARG_2946DBE5, + "serialized TsBlock", + "fragment range")); + } + int length = (int) Math.min(maxBytes, serializedTsBlock.remaining() - offset); + ByteBuffer fragment = serializedTsBlock.duplicate(); + fragment.position(Math.toIntExact(offset)); + fragment.limit(Math.toIntExact(offset + length)); + return fragment.slice(); } public void acknowledgeTsBlock(int startSequenceId, int endSequenceId) { @@ -439,6 +468,7 @@ public void acknowledgeTsBlock(int startSequenceId, int endSequenceId) { freedBytes += entry.getValue().right; bufferRetainedSizeInBytes -= entry.getValue().right; iterator.remove(); + sequenceIdToSerializedTsBlock.remove(entry.getKey()); if (LOGGER.isDebugEnabled()) { LOGGER.debug(DataNodeQueryMessages.ACK_TSBLOCK, entry.getKey()); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 67a5defb09ac3..1deba48e8ad32 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -51,6 +51,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import java.io.ByteArrayOutputStream; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; @@ -640,6 +641,8 @@ public void run() { startSequenceId, endSequenceId, indexOfUpstreamSinkHandle); + DataBlockFetchProgress fetchProgress = + new DataBlockFetchProgress(startSequenceId, endSequenceId); int attempt = 0; while (attempt < MAX_ATTEMPT_TIMES) { attempt += 1; @@ -648,7 +651,8 @@ public void run() { boolean transferAttemptRecorded = false; try (SyncDataNodeMPPDataExchangeServiceClient client = mppDataExchangeServiceClientManager.borrowClient(remoteEndpoint)) { - TGetDataBlockResponse resp = client.getDataBlock(req); + TGetDataBlockResponse resp = + getDataBlockWithFragments(client, req, fetchProgress); int tsBlockNum = resp.getTsBlocks().size(); if (tsBlockNum != endSequenceId - startSequenceId) { recordTransferAttempt( @@ -767,6 +771,98 @@ private void fail(Throwable t) { sourceHandleListener.onFailure(SourceHandle.this, t); } } + + private TGetDataBlockResponse getDataBlockWithFragments( + SyncDataNodeMPPDataExchangeServiceClient client, + TGetDataBlockRequest request, + DataBlockFetchProgress fetchProgress) + throws TException { + while (!fetchProgress.isFinished()) { + TGetDataBlockRequest fragmentRequest = request.deepCopy(); + fragmentRequest.setStartSequenceId(fetchProgress.nextSequenceId); + fragmentRequest.setOffset(fetchProgress.offset); + TGetDataBlockResponse response = client.getDataBlock(fragmentRequest); + if (response.getTsBlocks().isEmpty()) { + return response; + } + fetchProgress.addResponse(response); + } + return new TGetDataBlockResponse(fetchProgress.tsBlocks); + } + + private class DataBlockFetchProgress { + private final int endSequenceId; + private final List tsBlocks; + private int nextSequenceId; + private long offset; + private ByteArrayOutputStream partialTsBlock; + + private DataBlockFetchProgress(int startSequenceId, int endSequenceId) { + this.nextSequenceId = startSequenceId; + this.endSequenceId = endSequenceId; + this.tsBlocks = new ArrayList<>(endSequenceId - startSequenceId); + } + + private boolean isFinished() { + return nextSequenceId == endSequenceId && partialTsBlock == null; + } + + private void addResponse(TGetDataBlockResponse response) throws TException { + List responseBlocks = response.getTsBlocks(); + boolean lastBlockIsFragment = response.isSetOffset(); + int blockIndex = 0; + + if (partialTsBlock != null) { + appendFragment(responseBlocks.get(blockIndex++)); + if (lastBlockIsFragment && blockIndex == responseBlocks.size()) { + updateOffset(response.getOffset()); + return; + } + tsBlocks.add(ByteBuffer.wrap(partialTsBlock.toByteArray())); + partialTsBlock = null; + offset = 0; + nextSequenceId++; + } + + int lastCompleteBlockIndex = + lastBlockIsFragment ? responseBlocks.size() - 1 : responseBlocks.size(); + while (blockIndex < lastCompleteBlockIndex) { + tsBlocks.add(responseBlocks.get(blockIndex++)); + nextSequenceId++; + } + + if (lastBlockIsFragment) { + partialTsBlock = new ByteArrayOutputStream(); + appendFragment(responseBlocks.get(blockIndex)); + updateOffset(response.getOffset()); + } + + if (nextSequenceId > endSequenceId + || (!lastBlockIsFragment && nextSequenceId == endSequenceId && partialTsBlock != null)) { + throw new TException( + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + } + } + + private void appendFragment(ByteBuffer fragment) throws TException { + ByteBuffer duplicate = fragment.duplicate(); + if (!duplicate.hasRemaining()) { + throw new TException( + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + } + byte[] bytes = new byte[duplicate.remaining()]; + duplicate.get(bytes); + partialTsBlock.writeBytes(bytes); + } + + private void updateOffset(long nextOffset) throws TException { + if (nextOffset <= offset || nextOffset != partialTsBlock.size()) { + throw new TException( + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + } + offset = nextOffset; + } + } } class SendAcknowledgeDataBlockEventTask implements Runnable { diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index c53c95fd016ff..f7665c3009872 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -1194,6 +1194,11 @@ mpp_data_exchange_max_pool_size=10 # Datatype: int mpp_data_exchange_keep_alive_time_in_ms=1000 +# The maximum payload size of one MPP data exchange RPC response +# effectiveMode: restart +# Datatype: int, Unit: byte +mpp_data_exchange_max_payload_size_in_bytes=8388608 + # The max execution time of a DriverTask # effectiveMode: restart # Datatype: int, Unit: ms diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index c5c2bdac4a656..77dacf07a6bc2 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -90,10 +90,14 @@ struct TGetDataBlockRequest { 3: required i32 endSequenceId // Index of upstream SinkChannel 4: required i32 index + // Optional byte range for fetching one serialized TsBlock in fragments. + 5: optional i64 offset } struct TGetDataBlockResponse { 1: required list tsBlocks + // The start offset of the next fragment. It is set only when the last element in tsBlocks is a fragment. + 2: optional i64 offset } struct TAcknowledgeDataBlockEvent { From ee86f1c393c82f7b186e125aa61376d4258188ce Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Mon, 14 Sep 2026 15:46:59 +0800 Subject: [PATCH 02/12] make offset int Signed-off-by: Weihao Li <18110526956@163.com> --- .../execution/exchange/MPPDataExchangeManager.java | 10 +++++----- .../execution/exchange/sink/SinkChannel.java | 8 ++++---- .../execution/exchange/source/SourceHandle.java | 4 ++-- .../thrift-datanode/src/main/thrift/datanode.thrift | 4 ++-- 4 files changed, 13 insertions(+), 13 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index f8327673a8064..ea3dce68281cb 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -182,17 +182,17 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce IoTDBDescriptor.getInstance() .getConfig() .getMppDataExchangeMaxPayloadSizeInBytes(); - long offset = req.getOffset(); + int offset = req.getOffset(); for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { try { ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); - long blockOffset = i == req.getStartSequenceId() ? offset : 0L; - long remainingBlockSize = serializedTsBlock.remaining() - blockOffset; + int blockOffset = i == req.getStartSequenceId() ? offset : 0; + int remainingBlockSize = serializedTsBlock.remaining() - blockOffset; if (remainingBlockSize <= remainingPayloadSize) { resp.addToTsBlocks( sinkChannel.getSerializedTsBlockFragment( - i, blockOffset, Math.toIntExact(remainingBlockSize))); - remainingPayloadSize -= Math.toIntExact(remainingBlockSize); + i, blockOffset, remainingBlockSize)); + remainingPayloadSize -= remainingBlockSize; if (remainingPayloadSize == 0) { break; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index 533387afdf01a..86c66b6677e9e 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -432,7 +432,7 @@ public synchronized ByteBuffer getSerializedTsBlock(int sequenceId) throws IOExc } public synchronized ByteBuffer getSerializedTsBlockFragment( - int sequenceId, long offset, int maxBytes) throws IOException { + int sequenceId, int offset, int maxBytes) throws IOException { ByteBuffer serializedTsBlock = getSerializedTsBlock(sequenceId); if (offset < 0 || offset > serializedTsBlock.remaining() || maxBytes <= 0) { throw new IllegalArgumentException( @@ -441,10 +441,10 @@ public synchronized ByteBuffer getSerializedTsBlockFragment( "serialized TsBlock", "fragment range")); } - int length = (int) Math.min(maxBytes, serializedTsBlock.remaining() - offset); + int length = Math.min(maxBytes, serializedTsBlock.remaining() - offset); ByteBuffer fragment = serializedTsBlock.duplicate(); - fragment.position(Math.toIntExact(offset)); - fragment.limit(Math.toIntExact(offset + length)); + fragment.position(offset); + fragment.limit(offset + length); return fragment.slice(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 1deba48e8ad32..db4f1dfd40843 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -794,7 +794,7 @@ private class DataBlockFetchProgress { private final int endSequenceId; private final List tsBlocks; private int nextSequenceId; - private long offset; + private int offset; private ByteArrayOutputStream partialTsBlock; private DataBlockFetchProgress(int startSequenceId, int endSequenceId) { @@ -855,7 +855,7 @@ private void appendFragment(ByteBuffer fragment) throws TException { partialTsBlock.writeBytes(bytes); } - private void updateOffset(long nextOffset) throws TException { + private void updateOffset(int nextOffset) throws TException { if (nextOffset <= offset || nextOffset != partialTsBlock.size()) { throw new TException( DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index 77dacf07a6bc2..cf9aca63b17b9 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -91,13 +91,13 @@ struct TGetDataBlockRequest { // Index of upstream SinkChannel 4: required i32 index // Optional byte range for fetching one serialized TsBlock in fragments. - 5: optional i64 offset + 5: optional i32 offset } struct TGetDataBlockResponse { 1: required list tsBlocks // The start offset of the next fragment. It is set only when the last element in tsBlocks is a fragment. - 2: optional i64 offset + 2: optional i32 offset } struct TAcknowledgeDataBlockEvent { From 62fd96b325217c61b08dd59a04132823e5268a0c Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Mon, 14 Sep 2026 16:45:25 +0800 Subject: [PATCH 03/12] optimize some Signed-off-by: Weihao Li <18110526956@163.com> --- .../org/apache/iotdb/db/conf/IoTDBConfig.java | 3 +- .../exchange/MPPDataExchangeManager.java | 30 ++++++++++++------- .../exchange/source/SourceHandle.java | 7 +++-- 3 files changed, 24 insertions(+), 16 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index 7d7f3652d38d9..ff141a8ed85bf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -3427,8 +3427,7 @@ public int getMppDataExchangeMaxPayloadSizeInBytes() { return mppDataExchangeMaxPayloadSizeInBytes; } - public void setMppDataExchangeMaxPayloadSizeInBytes( - int mppDataExchangeMaxPayloadSizeInBytes) { + public void setMppDataExchangeMaxPayloadSizeInBytes(int mppDataExchangeMaxPayloadSizeInBytes) { this.mppDataExchangeMaxPayloadSizeInBytes = Math.max(1, Math.min(mppDataExchangeMaxPayloadSizeInBytes, thriftMaxFrameSize - 1024)); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index ea3dce68281cb..e60a7b3d89fb6 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -179,27 +179,35 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce SinkChannel sinkChannel = (SinkChannel) (sinkHandle.getChannel(req.getIndex())); if (req.isSetOffset()) { int remainingPayloadSize = - IoTDBDescriptor.getInstance() - .getConfig() - .getMppDataExchangeMaxPayloadSizeInBytes(); + IoTDBDescriptor.getInstance().getConfig().getMppDataExchangeMaxPayloadSizeInBytes(); int offset = req.getOffset(); for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { try { - ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); + ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i).asReadOnlyBuffer(); int blockOffset = i == req.getStartSequenceId() ? offset : 0; - int remainingBlockSize = serializedTsBlock.remaining() - blockOffset; + int serializedTsBlockSize = serializedTsBlock.remaining(); + if (blockOffset < 0 || blockOffset > serializedTsBlockSize) { + throw new IllegalArgumentException( + String.format( + DataNodeQueryMessages.EXCEPTION_INVALID_ARG_ARG_2946DBE5, + "serialized TsBlock", + "fragment range")); + } + int remainingBlockSize = serializedTsBlockSize - blockOffset; if (remainingBlockSize <= remainingPayloadSize) { - resp.addToTsBlocks( - sinkChannel.getSerializedTsBlockFragment( - i, blockOffset, remainingBlockSize)); + ByteBuffer fragment = serializedTsBlock.duplicate(); + fragment.position(blockOffset); + fragment.limit(blockOffset + remainingBlockSize); + resp.addToTsBlocks(fragment.slice()); remainingPayloadSize -= remainingBlockSize; if (remainingPayloadSize == 0) { break; } } else { - resp.addToTsBlocks( - sinkChannel.getSerializedTsBlockFragment( - i, blockOffset, remainingPayloadSize)); + ByteBuffer fragment = serializedTsBlock.duplicate(); + fragment.position(blockOffset); + fragment.limit(blockOffset + remainingPayloadSize); + resp.addToTsBlocks(fragment.slice()); resp.setOffset(blockOffset + remainingPayloadSize); break; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index db4f1dfd40843..3afc800c2afa4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -651,8 +651,7 @@ public void run() { boolean transferAttemptRecorded = false; try (SyncDataNodeMPPDataExchangeServiceClient client = mppDataExchangeServiceClientManager.borrowClient(remoteEndpoint)) { - TGetDataBlockResponse resp = - getDataBlockWithFragments(client, req, fetchProgress); + TGetDataBlockResponse resp = getDataBlockWithFragments(client, req, fetchProgress); int tsBlockNum = resp.getTsBlocks().size(); if (tsBlockNum != endSequenceId - startSequenceId) { recordTransferAttempt( @@ -838,7 +837,9 @@ private void addResponse(TGetDataBlockResponse response) throws TException { } if (nextSequenceId > endSequenceId - || (!lastBlockIsFragment && nextSequenceId == endSequenceId && partialTsBlock != null)) { + || (!lastBlockIsFragment + && nextSequenceId == endSequenceId + && partialTsBlock != null)) { throw new TException( DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); } From 95c4176f640d30bb4e1b01762a1b8bb7e733a111 Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Tue, 15 Sep 2026 12:59:16 +0800 Subject: [PATCH 04/12] config hotReload Signed-off-by: Weihao Li <18110526956@163.com> --- .../iotdb/db/i18n/DataNodeMiscMessages.java | 6 ++++ .../iotdb/db/i18n/DataNodeMiscMessages.java | 6 ++++ .../org/apache/iotdb/db/conf/IoTDBConfig.java | 5 ++- .../apache/iotdb/db/conf/IoTDBDescriptor.java | 35 ++++++++++++++++--- .../conf/iotdb-system.properties.template | 10 +++--- 5 files changed, 50 insertions(+), 12 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index 12f4e374ca7a7..0cdcce0121e12 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 @@ -1488,5 +1488,11 @@ private DataNodeMiscMessages() {} public static final String LOG_TABLE_QUERY_DEVICE_ENTRY_BATCH_SIZE_IN_BYTES_ARG_EXCEEDS_DN_THRIFT_MAX_FRAME_SIZE_ARG_USING_ARG_AS_THE_EFFECTIVE_VALUE_2AE1BEDA = "table_query_device_entry_batch_size_in_bytes (%d) exceeds the maximum RPC payload (dn_thrift_max_frame_size %d minus 1024 bytes); using %d as the effective value"; + public static final String + LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_NOT_POSITIVE_USING_DEFAULT_VALUE_ARG_1AA821B2 = + "mpp_data_exchange_max_payload_size_in_bytes (%d) is not positive, using default value %d"; + public static final String + LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC = + "mpp_data_exchange_max_payload_size_in_bytes (%d) exceeds the maximum allowed value %d, using %d"; } 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..155ca9f0c4d00 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 @@ -1468,5 +1468,11 @@ private DataNodeMiscMessages() {} public static final String LOG_TABLE_QUERY_DEVICE_ENTRY_BATCH_SIZE_IN_BYTES_ARG_EXCEEDS_DN_THRIFT_MAX_FRAME_SIZE_ARG_USING_ARG_AS_THE_EFFECTIVE_VALUE_2AE1BEDA = "table_query_device_entry_batch_size_in_bytes（%d）超过最大 RPC payload（dn_thrift_max_frame_size %d 减去 1024 字节），将使用 %d 作为生效值"; + public static final String + LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_NOT_POSITIVE_USING_DEFAULT_VALUE_ARG_1AA821B2 = + "mpp_data_exchange_max_payload_size_in_bytes（%d）不是正数，将使用默认值 %d"; + public static final String + LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC = + "mpp_data_exchange_max_payload_size_in_bytes（%d）超过允许的最大值 %d，将使用 %d"; } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java index ff141a8ed85bf..6d169eda42cf9 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBConfig.java @@ -956,7 +956,7 @@ public class IoTDBConfig { /** Core pool size of mpp data exchange. */ private int mppDataExchangeCorePoolSize = 10; - private int mppDataExchangeMaxPayloadSizeInBytes = 8 * 1024 * 1024; + private int mppDataExchangeMaxPayloadSizeInBytes = 4 * 1024 * 1024; /** Max pool size of mpp data exchange. */ private int mppDataExchangeMaxPoolSize = 10; @@ -3428,8 +3428,7 @@ public int getMppDataExchangeMaxPayloadSizeInBytes() { } public void setMppDataExchangeMaxPayloadSizeInBytes(int mppDataExchangeMaxPayloadSizeInBytes) { - this.mppDataExchangeMaxPayloadSizeInBytes = - Math.max(1, Math.min(mppDataExchangeMaxPayloadSizeInBytes, thriftMaxFrameSize - 1024)); + this.mppDataExchangeMaxPayloadSizeInBytes = mppDataExchangeMaxPayloadSizeInBytes; } public int getConnectionTimeoutInMS() { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 8c3efb84b83b2..83cdf3cac1070 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -2268,6 +2268,7 @@ public synchronized void loadHotModifiedProps(TrimProperties properties) "enable_topk_runtime_filter")))); loadTableQueryDeviceEntryBatchSize(properties); + loadMppDataExchangeMaxPayloadSize(properties); // update wal config long prevDeleteWalFilesPeriodInMs = conf.getDeleteWalFilesPeriodInMs(); @@ -3070,11 +3071,7 @@ public void loadShuffleProps(TrimProperties properties) { properties.getProperty( "mpp_data_exchange_keep_alive_time_in_ms", Integer.toString(conf.getMppDataExchangeKeepAliveTimeInMs())))); - conf.setMppDataExchangeMaxPayloadSizeInBytes( - Integer.parseInt( - properties.getProperty( - "mpp_data_exchange_max_payload_size_in_bytes", - Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes())))); + loadMppDataExchangeMaxPayloadSize(properties); conf.setPartitionCacheSize( Integer.parseInt( @@ -3088,6 +3085,34 @@ public void loadShuffleProps(TrimProperties properties) { Integer.toString(commonConfig.getDriverTaskExecutionTimeSliceInMs())))); } + private void loadMppDataExchangeMaxPayloadSize(TrimProperties properties) { + int configuredSize = + Integer.parseInt( + properties.getProperty( + "mpp_data_exchange_max_payload_size_in_bytes", + Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes()))); + int defaultSize = 4 * 1024 * 1024; + if (configuredSize <= 0) { + LOGGER.warn( + DataNodeMiscMessages + .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_NOT_POSITIVE_USING_DEFAULT_VALUE_ARG_1AA821B2, + configuredSize, + defaultSize); + configuredSize = defaultSize; + } + int maxAllowedSize = Math.max(1, conf.getThriftMaxFrameSize() - 1024); + if (configuredSize > maxAllowedSize) { + LOGGER.warn( + DataNodeMiscMessages + .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC, + configuredSize, + maxAllowedSize, + maxAllowedSize); + configuredSize = maxAllowedSize; + } + conf.setMppDataExchangeMaxPayloadSizeInBytes(configuredSize); + } + /** Get default encode algorithm by data type */ public TSEncoding getDefaultEncodingByType(TSDataType dataType) { switch (dataType) { diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index f7665c3009872..2a51650c0cc51 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -71,7 +71,7 @@ cn_consensus_port=10720 # Used for connection of IoTDB native clients(Session) # Could set 127.0.0.1(for local test), ipv4/ipv6 address, or hostname. -# effectiveMode: restart +# effectiveMode: hot_reload # Datatype: String dn_rpc_address=127.0.0.1 @@ -1194,10 +1194,12 @@ mpp_data_exchange_max_pool_size=10 # Datatype: int mpp_data_exchange_keep_alive_time_in_ms=1000 -# The maximum payload size of one MPP data exchange RPC response -# effectiveMode: restart +# The maximum payload size of one MPP data exchange RPC response. +# <=0 use default value +# The effective value is capped by dn_thrift_max_frame_size minus 1024 bytes reserved for the RPC response envelope. +# effectiveMode: hot_reload # Datatype: int, Unit: byte -mpp_data_exchange_max_payload_size_in_bytes=8388608 +mpp_data_exchange_max_payload_size_in_bytes=4194304 # The max execution time of a DriverTask # effectiveMode: restart From a2beae65cccb484841328f60f1a50033367f76ae Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Tue, 15 Sep 2026 16:51:39 +0800 Subject: [PATCH 05/12] modify some Signed-off-by: Weihao Li <18110526956@163.com> --- .../exchange/MPPDataExchangeManager.java | 5 ++- .../execution/exchange/sink/SinkChannel.java | 3 +- .../exchange/source/SourceHandle.java | 39 ++++++------------- .../execution/exchange/SinkChannelTest.java | 16 +++++++- .../queryengine/execution/exchange/Utils.java | 7 ++++ .../src/main/thrift/datanode.thrift | 2 + 6 files changed, 42 insertions(+), 30 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index e60a7b3d89fb6..284aa1976ce00 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -183,7 +183,7 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce int offset = req.getOffset(); for (int i = req.getStartSequenceId(); i < req.getEndSequenceId(); i++) { try { - ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i).asReadOnlyBuffer(); + ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); int blockOffset = i == req.getStartSequenceId() ? offset : 0; int serializedTsBlockSize = serializedTsBlock.remaining(); if (blockOffset < 0 || blockOffset > serializedTsBlockSize) { @@ -209,6 +209,9 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce fragment.limit(blockOffset + remainingPayloadSize); resp.addToTsBlocks(fragment.slice()); resp.setOffset(blockOffset + remainingPayloadSize); + if (blockOffset == 0) { + resp.setTotalLength(serializedTsBlockSize); + } break; } } catch (GetTsBlockFromClosedOrAbortedChannelException e) { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index 86c66b6677e9e..c98be7d55bd61 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -427,7 +427,8 @@ public synchronized ByteBuffer getSerializedTsBlock(int sequenceId) throws IOExc DataNodeQueryMessages.THE_DATA_BLOCK_DOESN_T_EXIST_SEQUENCE_ID + sequenceId); } serializedTsBlock = serde.serialize(pair.left); - sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock.asReadOnlyBuffer()); + sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock); + pair.left = null; return serializedTsBlock.duplicate(); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 3afc800c2afa4..780c523b007b2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -47,11 +47,11 @@ import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.apache.tsfile.utils.Pair; +import org.apache.tsfile.utils.PublicBAOS; import org.apache.tsfile.utils.RamUsageEstimator; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -import java.io.ByteArrayOutputStream; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.HashMap; @@ -60,6 +60,7 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; +import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.util.concurrent.Futures.nonCancellationPropagating; import static org.apache.iotdb.db.queryengine.execution.exchange.MPPDataExchangeManager.createFullIdFrom; import static org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet.GET_DATA_BLOCK_TASK_CALLER; @@ -794,7 +795,7 @@ private class DataBlockFetchProgress { private final List tsBlocks; private int nextSequenceId; private int offset; - private ByteArrayOutputStream partialTsBlock; + private PublicBAOS partialTsBlock; private DataBlockFetchProgress(int startSequenceId, int endSequenceId) { this.nextSequenceId = startSequenceId; @@ -806,7 +807,7 @@ private boolean isFinished() { return nextSequenceId == endSequenceId && partialTsBlock == null; } - private void addResponse(TGetDataBlockResponse response) throws TException { + private void addResponse(TGetDataBlockResponse response) { List responseBlocks = response.getTsBlocks(); boolean lastBlockIsFragment = response.isSetOffset(); int blockIndex = 0; @@ -817,7 +818,7 @@ private void addResponse(TGetDataBlockResponse response) throws TException { updateOffset(response.getOffset()); return; } - tsBlocks.add(ByteBuffer.wrap(partialTsBlock.toByteArray())); + tsBlocks.add(ByteBuffer.wrap(partialTsBlock.getBuf())); partialTsBlock = null; offset = 0; nextSequenceId++; @@ -831,36 +832,20 @@ private void addResponse(TGetDataBlockResponse response) throws TException { } if (lastBlockIsFragment) { - partialTsBlock = new ByteArrayOutputStream(); + checkArgument(response.isSetTotalLength(), "xxx"); + partialTsBlock = new PublicBAOS(response.getTotalLength()); appendFragment(responseBlocks.get(blockIndex)); updateOffset(response.getOffset()); } - - if (nextSequenceId > endSequenceId - || (!lastBlockIsFragment - && nextSequenceId == endSequenceId - && partialTsBlock != null)) { - throw new TException( - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); - } } - private void appendFragment(ByteBuffer fragment) throws TException { - ByteBuffer duplicate = fragment.duplicate(); - if (!duplicate.hasRemaining()) { - throw new TException( - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); - } - byte[] bytes = new byte[duplicate.remaining()]; - duplicate.get(bytes); - partialTsBlock.writeBytes(bytes); + private void appendFragment(ByteBuffer fragment) { + checkArgument(fragment.hasRemaining(), "xxx"); + partialTsBlock.writeBytes(fragment.array()); } - private void updateOffset(int nextOffset) throws TException { - if (nextOffset <= offset || nextOffset != partialTsBlock.size()) { - throw new TException( - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); - } + private void updateOffset(int nextOffset) { + checkArgument(nextOffset > offset && nextOffset == partialTsBlock.size(), "xxx"); offset = nextOffset; } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java index f59dd9925aad1..f22725c0feab7 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SinkChannelTest.java @@ -37,6 +37,7 @@ import org.apache.thrift.TException; import org.apache.tsfile.common.conf.TSFileDescriptor; import org.apache.tsfile.read.common.block.TsBlock; +import org.apache.tsfile.read.common.block.column.TsBlockSerde; import org.junit.Assert; import org.junit.Test; import org.mockito.Mockito; @@ -87,6 +88,7 @@ public void testOneTimeNotBlockedSend() { SinkListener mockSinkListener = Mockito.mock(SinkListener.class); // Construct several mock TsBlock(s). List mockTsBlocks = Utils.createMockTsBlocks(numOfMockTsBlock, mockTsBlockSize); + TsBlockSerde mockTsBlockSerde = Utils.createMockTsBlockSerde(mockTsBlockSize); // Construct SinkChannel. SinkChannel sinkChannel = @@ -98,7 +100,7 @@ public void testOneTimeNotBlockedSend() { localFragmentInstanceId, mockLocalMemoryManager, Executors.newSingleThreadExecutor(), - Utils.createMockTsBlockSerde(mockTsBlockSize), + mockTsBlockSerde, mockSinkListener, mockClientManager); sinkChannel.open(); @@ -145,6 +147,8 @@ public void testOneTimeNotBlockedSend() { for (int i = 0; i < numOfMockTsBlock; i++) { try { sinkChannel.getSerializedTsBlock(i); + sinkChannel.getSerializedTsBlock(i); + Mockito.verify(mockTsBlockSerde, Mockito.times(1)).serialize(mockTsBlocks.get(i)); } catch (IOException e) { e.printStackTrace(); Assert.fail(); @@ -166,6 +170,16 @@ public void testOneTimeNotBlockedSend() { Assert.assertTrue(sinkChannel.isFinished()); Assert.assertFalse(sinkChannel.isAborted()); Assert.assertEquals(mockTsBlockSize, sinkChannel.getBufferRetainedSizeInBytes()); + for (int i = 0; i < numOfMockTsBlock; i++) { + try { + sinkChannel.getSerializedTsBlock(i); + Assert.fail("The acknowledged serialized TsBlock should have been released"); + } catch (IllegalStateException expected) { + // Both the original entry and serialized cache must be removed by acknowledgement. + } catch (IOException e) { + Assert.fail(e.getMessage()); + } + } Mockito.verify(mockMemoryPool, Mockito.timeout(10_0000).times(1)) .free( queryId, diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java index b09498ad949dd..1a701b9eded95 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/Utils.java @@ -28,6 +28,7 @@ import org.mockito.Mockito; import org.mockito.stubbing.Answer; +import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; @@ -155,6 +156,12 @@ public static TsBlockSerde createMockTsBlockSerde(long mockTsBlockSize) { TsBlock mockTsBlock = Mockito.mock(TsBlock.class); Mockito.when(mockTsBlock.getRetainedSizeInBytes()).thenReturn(mockTsBlockSize); Mockito.when(mockTsBlock.getSizeInBytes()).thenReturn(mockTsBlockSize); + try { + Mockito.when(mockTsBlockSerde.serialize(Mockito.any(TsBlock.class))) + .thenReturn(ByteBuffer.allocate(Math.toIntExact(mockTsBlockSize))); + } catch (IOException e) { + throw new AssertionError(e); + } Mockito.when(mockTsBlockSerde.deserialize(Mockito.any(ByteBuffer.class))) .thenReturn(mockTsBlock); return mockTsBlockSerde; diff --git a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift index cf9aca63b17b9..ceb794989ce70 100644 --- a/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift +++ b/iotdb-protocol/thrift-datanode/src/main/thrift/datanode.thrift @@ -98,6 +98,8 @@ struct TGetDataBlockResponse { 1: required list tsBlocks // The start offset of the next fragment. It is set only when the last element in tsBlocks is a fragment. 2: optional i32 offset + // Total serialized length of the TsBlock when the response starts its first fragment. + 3: optional i32 totalLength } struct TAcknowledgeDataBlockEvent { From 5f73d3f2e5019928e996a96c4b1fe7a79d272e55 Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Tue, 15 Sep 2026 18:10:17 +0800 Subject: [PATCH 06/12] merge map Signed-off-by: Weihao Li <18110526956@163.com> --- .../execution/exchange/sink/SinkChannel.java | 72 +++++++++---------- 1 file changed, 33 insertions(+), 39 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java index c98be7d55bd61..d24a1552a5b22 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/sink/SinkChannel.java @@ -44,7 +44,6 @@ import org.apache.tsfile.external.commons.lang3.Validate; import org.apache.tsfile.read.common.block.TsBlock; import org.apache.tsfile.read.common.block.column.TsBlockSerde; -import org.apache.tsfile.utils.Pair; import org.apache.tsfile.utils.RamUsageEstimator; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -67,6 +66,19 @@ public class SinkChannel implements ISinkChannel { + private static class TsBlockInfo { + private TsBlock tsBlock; + private ByteBuffer serializedTsBlock; + private final long tsBlockSize; + + // private Class, no need to do access control + private TsBlockInfo(TsBlock tsBlock, ByteBuffer serializedTsBlock, long tsBlockSize) { + this.tsBlock = tsBlock; + this.serializedTsBlock = serializedTsBlock; + this.tsBlockSize = tsBlockSize; + } + } + private static final Logger LOGGER = LoggerFactory.getLogger(SinkChannel.class); public static final int MAX_ATTEMPT_TIMES = 3; @@ -93,12 +105,7 @@ public class SinkChannel implements ISinkChannel { // Use LinkedHashMap to meet 2 needs, // 1. Predictable iteration order so that removing buffered TsBlocks can be efficient. // 2. Fast lookup. - private final LinkedHashMap> sequenceIdToTsBlock = - new LinkedHashMap<>(); - - /** Serialized blocks are cached so fragmented requests do not serialize the same block again. */ - private final LinkedHashMap sequenceIdToSerializedTsBlock = - new LinkedHashMap<>(); + private final LinkedHashMap sequenceIdToTsBlock = new LinkedHashMap<>(); // size for current TsBlock to reserve and free private long currentTsBlockSize; @@ -278,7 +285,7 @@ public synchronized void send(TsBlock tsBlock) { blocked = reserveResult.getFuture(); bufferRetainedSizeInBytes += reserveResult.getReservedBytes(); - sequenceIdToTsBlock.put(nextSequenceId, new Pair<>(tsBlock, currentTsBlockSize)); + sequenceIdToTsBlock.put(nextSequenceId, new TsBlockInfo(tsBlock, null, currentTsBlockSize)); nextSequenceId += 1; currentTsBlockSize = reserveResult.getReservedBytes(); @@ -309,7 +316,6 @@ public synchronized boolean abort() { return false; } sequenceIdToTsBlock.clear(); - sequenceIdToSerializedTsBlock.clear(); if (blocked != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked); } @@ -340,7 +346,6 @@ public synchronized boolean close() { return false; } sequenceIdToTsBlock.clear(); - sequenceIdToSerializedTsBlock.clear(); if (blocked != null) { bufferRetainedSizeInBytes -= localMemoryManager.getQueryPool().tryCancel(blocked); } @@ -413,12 +418,20 @@ public synchronized ByteBuffer getSerializedTsBlock(int sequenceId) throws IOExc throw new GetTsBlockFromClosedOrAbortedChannelException( DataNodeQueryMessages.SINKCHANNEL_IS_ABORTED_OR_CLOSED); } - ByteBuffer serializedTsBlock = sequenceIdToSerializedTsBlock.get(sequenceId); + TsBlockInfo tsBlockInfo = sequenceIdToTsBlock.get(sequenceId); + if (tsBlockInfo == null) { + LOGGER.warn( + DataNodeQueryMessages.THE_TSBLOCK_DOESNT_EXIST_SEQUENCE_ID_REMAINING, + sequenceId, + sequenceIdToTsBlock.entrySet()); + throw new IllegalStateException( + DataNodeQueryMessages.THE_DATA_BLOCK_DOESN_T_EXIST_SEQUENCE_ID + sequenceId); + } + ByteBuffer serializedTsBlock = tsBlockInfo.serializedTsBlock; if (serializedTsBlock != null) { return serializedTsBlock.duplicate(); } - Pair pair = sequenceIdToTsBlock.get(sequenceId); - if (pair == null || pair.left == null) { + if (tsBlockInfo.tsBlock == null) { LOGGER.warn( DataNodeQueryMessages.THE_TSBLOCK_DOESNT_EXIST_SEQUENCE_ID_REMAINING, sequenceId, @@ -426,39 +439,21 @@ public synchronized ByteBuffer getSerializedTsBlock(int sequenceId) throws IOExc throw new IllegalStateException( DataNodeQueryMessages.THE_DATA_BLOCK_DOESN_T_EXIST_SEQUENCE_ID + sequenceId); } - serializedTsBlock = serde.serialize(pair.left); - sequenceIdToSerializedTsBlock.put(sequenceId, serializedTsBlock); - pair.left = null; + serializedTsBlock = serde.serialize(tsBlockInfo.tsBlock); + tsBlockInfo.serializedTsBlock = serializedTsBlock; + tsBlockInfo.tsBlock = null; return serializedTsBlock.duplicate(); } - public synchronized ByteBuffer getSerializedTsBlockFragment( - int sequenceId, int offset, int maxBytes) throws IOException { - ByteBuffer serializedTsBlock = getSerializedTsBlock(sequenceId); - if (offset < 0 || offset > serializedTsBlock.remaining() || maxBytes <= 0) { - throw new IllegalArgumentException( - String.format( - DataNodeQueryMessages.EXCEPTION_INVALID_ARG_ARG_2946DBE5, - "serialized TsBlock", - "fragment range")); - } - int length = Math.min(maxBytes, serializedTsBlock.remaining() - offset); - ByteBuffer fragment = serializedTsBlock.duplicate(); - fragment.position(offset); - fragment.limit(offset + length); - return fragment.slice(); - } - public void acknowledgeTsBlock(int startSequenceId, int endSequenceId) { long freedBytes = 0L; synchronized (this) { if (aborted || closed) { return; } - Iterator>> iterator = - sequenceIdToTsBlock.entrySet().iterator(); + Iterator> iterator = sequenceIdToTsBlock.entrySet().iterator(); while (iterator.hasNext()) { - Entry> entry = iterator.next(); + Entry entry = iterator.next(); if (entry.getKey() < startSequenceId) { continue; } @@ -466,10 +461,9 @@ public void acknowledgeTsBlock(int startSequenceId, int endSequenceId) { break; } - freedBytes += entry.getValue().right; - bufferRetainedSizeInBytes -= entry.getValue().right; + freedBytes += entry.getValue().tsBlockSize; + bufferRetainedSizeInBytes -= entry.getValue().tsBlockSize; iterator.remove(); - sequenceIdToSerializedTsBlock.remove(entry.getKey()); if (LOGGER.isDebugEnabled()) { LOGGER.debug(DataNodeQueryMessages.ACK_TSBLOCK, entry.getKey()); } From 36cc05dc6a6649d14ab3a837a7130bee6cf35ba4 Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Wed, 16 Sep 2026 10:05:34 +0800 Subject: [PATCH 07/12] fix UT Signed-off-by: Weihao Li <18110526956@163.com> --- .../queryengine/execution/exchange/SourceHandleTest.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index 03220a3d10bf9..2e8c8a8c9d366 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -638,9 +638,11 @@ public void testShortResponseRetriesAndFails() { for (int i = 0; i < request.getEndSequenceId() - request.getStartSequenceId() - 1; i++) { - shortResponse.add(ByteBuffer.allocate(0)); + shortResponse.add(ByteBuffer.allocate(1)); } - return new TGetDataBlockResponse(shortResponse); + TGetDataBlockResponse response = new TGetDataBlockResponse(shortResponse); + response.setOffset(1); + return response; }) .when(mockClient) .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); @@ -676,7 +678,7 @@ public void testShortResponseRetriesAndFails() { Assert.fail(e.getMessage()); } Mockito.verify(mockSourceHandleListener, Mockito.timeout(10_000).times(1)) - .onFailure(Mockito.eq(sourceHandle), Mockito.any(TException.class)); + .onFailure(Mockito.eq(sourceHandle), Mockito.any(Throwable.class)); Assert.assertFalse(blocked.isDone()); Assert.assertEquals(0L, sourceHandle.getBufferRetainedSizeInBytes()); From e8f1916b502c488167e862095b782ad465330693 Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Wed, 16 Sep 2026 15:57:24 +0800 Subject: [PATCH 08/12] modify some Signed-off-by: Weihao Li <18110526956@163.com> --- .../iotdb/db/i18n/DataNodeMiscMessages.java | 3 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../iotdb/db/i18n/DataNodeMiscMessages.java | 3 + .../iotdb/db/i18n/DataNodeQueryMessages.java | 2 + .../apache/iotdb/db/conf/IoTDBDescriptor.java | 22 ++- .../exchange/MPPDataExchangeManager.java | 12 +- .../exchange/source/SourceHandle.java | 36 ++++- .../apache/iotdb/db/conf/PropertiesTest.java | 28 ++++ .../execution/exchange/SourceHandleTest.java | 144 +++++++++++++++++- .../conf/iotdb-system.properties.template | 4 +- 10 files changed, 233 insertions(+), 23 deletions(-) diff --git a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodeMiscMessages.java index 0cdcce0121e12..dc71d3f22a941 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 @@ -1494,5 +1494,8 @@ private DataNodeMiscMessages() {} public static final String LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC = "mpp_data_exchange_max_payload_size_in_bytes (%d) exceeds the maximum allowed value %d, using %d"; + public static final String + LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_BELOW_MINIMUM_ALLOWED_VALUE_ARG_USING_ARG_794ABC76 = + "mpp_data_exchange_max_payload_size_in_bytes (%d) is below the minimum allowed value %d, using %d"; } 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 2c25ba7460643..946e39c4874c6 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 @@ -1412,6 +1412,8 @@ public final class DataNodeQueryMessages { "failed to get data block [{}, {}), attempt times: {}"; public static final String EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33 = "Unexpected data block response size."; + public static final String EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_RANGE_0672002F = + "Invalid serialized TsBlock fragment range"; public static final String FAILED_TO_SEND_ACK_DATA_BLOCK_EVENT = "failed to send ack data block event [{}, {}), attempt times: {}"; public static final String SEND_CLOSE_SINK_CHANNEL_EVENT_FAILED = 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 155ca9f0c4d00..618554527e316 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 @@ -1474,5 +1474,8 @@ private DataNodeMiscMessages() {} public static final String LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC = "mpp_data_exchange_max_payload_size_in_bytes（%d）超过允许的最大值 %d，将使用 %d"; + public static final String + LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_BELOW_MINIMUM_ALLOWED_VALUE_ARG_USING_ARG_794ABC76 = + "mpp_data_exchange_max_payload_size_in_bytes（%d）低于允许的最小值 %d，将使用 %d"; } 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 c6a55ebad3e06..940a19cc86036 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 @@ -1394,6 +1394,8 @@ public final class DataNodeQueryMessages { "获取数据块 [{}, {}) 失败，尝试次数：{}"; public static final String EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33 = "数据块响应数量异常。"; + public static final String EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_RANGE_0672002F = + "无效的序列化 TsBlock 分片范围"; public static final String FAILED_TO_SEND_ACK_DATA_BLOCK_EVENT = "发送数据块确认事件 [{}, {}) 失败，尝试次数：{}"; public static final String SEND_CLOSE_SINK_CHANNEL_EVENT_FAILED = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index 83cdf3cac1070..f46e02fc18a6f 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -121,6 +121,10 @@ public class IoTDBDescriptor { private static final long DEVICE_ENTRY_RPC_FRAME_RESERVED_BYTES = 1024; + private static final int DEFAULT_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES = 4 * 1024 * 1024; + + private static final int MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES = 128 * 1024; + private static final String[] DEFAULT_WAL_THRESHOLD_NAME = { "iot_consensus_throttle_threshold_in_byte", "wal_throttle_threshold_in_byte" }; @@ -2447,6 +2451,9 @@ private void overlayEffectiveConfigurationValues() { ConfigurationFileUtils.updateAppliedProperties( "table_query_device_entry_batch_size_in_bytes", Long.toString(conf.getTableQueryDeviceEntryBatchSizeInBytes())); + ConfigurationFileUtils.updateAppliedProperties( + "mpp_data_exchange_max_payload_size_in_bytes", + Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes())); ConfigurationFileUtils.updateAppliedProperties( DEFAULT_WAL_THRESHOLD_NAME[1], Long.toString(conf.getThrottleThreshold())); } @@ -3091,16 +3098,23 @@ private void loadMppDataExchangeMaxPayloadSize(TrimProperties properties) { properties.getProperty( "mpp_data_exchange_max_payload_size_in_bytes", Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes()))); - int defaultSize = 4 * 1024 * 1024; if (configuredSize <= 0) { LOGGER.warn( DataNodeMiscMessages .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_NOT_POSITIVE_USING_DEFAULT_VALUE_ARG_1AA821B2, configuredSize, - defaultSize); - configuredSize = defaultSize; + DEFAULT_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES); + configuredSize = DEFAULT_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES; + } else if (configuredSize < MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES) { + LOGGER.warn( + DataNodeMiscMessages + .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_BELOW_MINIMUM_ALLOWED_VALUE_ARG_USING_ARG_794ABC76, + configuredSize, + MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES, + MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES); + configuredSize = MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES; } - int maxAllowedSize = Math.max(1, conf.getThriftMaxFrameSize() - 1024); + int maxAllowedSize = conf.getThriftMaxFrameSize() - 1024; if (configuredSize > maxAllowedSize) { LOGGER.warn( DataNodeMiscMessages diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index 284aa1976ce00..4f9653ccd0718 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -80,6 +80,7 @@ import java.util.function.Supplier; import java.util.stream.Collectors; +import static com.google.common.base.Preconditions.checkArgument; import static org.apache.iotdb.db.queryengine.common.DataNodeEndPoints.isSameNode; import static org.apache.iotdb.db.queryengine.common.FragmentInstanceId.createFullId; import static org.apache.iotdb.db.queryengine.metric.DataExchangeCostMetricSet.GET_DATA_BLOCK_TASK_SERVER; @@ -186,13 +187,10 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce ByteBuffer serializedTsBlock = sinkChannel.getSerializedTsBlock(i); int blockOffset = i == req.getStartSequenceId() ? offset : 0; int serializedTsBlockSize = serializedTsBlock.remaining(); - if (blockOffset < 0 || blockOffset > serializedTsBlockSize) { - throw new IllegalArgumentException( - String.format( - DataNodeQueryMessages.EXCEPTION_INVALID_ARG_ARG_2946DBE5, - "serialized TsBlock", - "fragment range")); - } + checkArgument( + blockOffset >= 0 && blockOffset <= serializedTsBlockSize, + DataNodeQueryMessages + .EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_RANGE_0672002F); int remainingBlockSize = serializedTsBlockSize - blockOffset; if (remainingBlockSize <= remainingPayloadSize) { ByteBuffer fragment = serializedTsBlock.duplicate(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 780c523b007b2..a48da4f066e14 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -40,6 +40,7 @@ import org.apache.iotdb.mpp.rpc.thrift.TGetDataBlockRequest; import org.apache.iotdb.mpp.rpc.thrift.TGetDataBlockResponse; +import com.google.common.base.Preconditions; import com.google.common.util.concurrent.ListenableFuture; import com.google.common.util.concurrent.SettableFuture; import org.apache.thrift.TException; @@ -713,7 +714,14 @@ public void run() { return; } break; - } catch (Throwable e) { + } catch (IllegalArgumentException | IllegalStateException e) { + if (!transferAttemptRecorded) { + recordTransferAttempt( + false, UserDataTransferErrorCode.UNEXPECTED_RESPONSE_SIZE.name(), e); + } + fail(e); + return; + } catch (Exception e) { if (!transferAttemptRecorded) { recordTransferAttempt(false, null, e); @@ -796,6 +804,7 @@ private class DataBlockFetchProgress { private int nextSequenceId; private int offset; private PublicBAOS partialTsBlock; + private int partialTsBlockTotalLength; private DataBlockFetchProgress(int startSequenceId, int endSequenceId) { this.nextSequenceId = startSequenceId; @@ -818,8 +827,12 @@ private void addResponse(TGetDataBlockResponse response) { updateOffset(response.getOffset()); return; } - tsBlocks.add(ByteBuffer.wrap(partialTsBlock.getBuf())); + Preconditions.checkState( + partialTsBlock.size() == partialTsBlockTotalLength, + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + tsBlocks.add(ByteBuffer.wrap(partialTsBlock.getBuf(), 0, partialTsBlock.size())); partialTsBlock = null; + partialTsBlockTotalLength = 0; offset = 0; nextSequenceId++; } @@ -832,20 +845,33 @@ private void addResponse(TGetDataBlockResponse response) { } if (lastBlockIsFragment) { - checkArgument(response.isSetTotalLength(), "xxx"); + checkArgument( + response.isSetTotalLength(), + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); partialTsBlock = new PublicBAOS(response.getTotalLength()); appendFragment(responseBlocks.get(blockIndex)); updateOffset(response.getOffset()); } + + Preconditions.checkState( + nextSequenceId <= endSequenceId + && (lastBlockIsFragment + || nextSequenceId != endSequenceId + || partialTsBlock == null), + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); } private void appendFragment(ByteBuffer fragment) { - checkArgument(fragment.hasRemaining(), "xxx"); + checkArgument( + fragment.hasRemaining(), + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); partialTsBlock.writeBytes(fragment.array()); } private void updateOffset(int nextOffset) { - checkArgument(nextOffset > offset && nextOffset == partialTsBlock.size(), "xxx"); + checkArgument( + nextOffset > offset && nextOffset == partialTsBlock.size(), + DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); offset = nextOffset; } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java index b842318199ffc..844357ed82e92 100755 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/conf/PropertiesTest.java @@ -271,6 +271,34 @@ public void testHotReloadRestRuntimeLimitProperties() throws Exception { } } + @Test + public void testMppDataExchangeMaxPayloadSizeHotReload() throws Exception { + IoTDBDescriptor descriptor = IoTDBDescriptor.getInstance(); + int originalPayloadSize = descriptor.getConfig().getMppDataExchangeMaxPayloadSizeInBytes(); + try { + TrimProperties properties = new TrimProperties(); + + properties.setProperty("mpp_data_exchange_max_payload_size_in_bytes", "0"); + descriptor.loadHotModifiedProps(properties); + Assert.assertEquals( + 4 * 1024 * 1024, descriptor.getConfig().getMppDataExchangeMaxPayloadSizeInBytes()); + + properties.setProperty("mpp_data_exchange_max_payload_size_in_bytes", "1"); + descriptor.loadHotModifiedProps(properties); + Assert.assertEquals( + 128 * 1024, descriptor.getConfig().getMppDataExchangeMaxPayloadSizeInBytes()); + + int maximumPayloadSize = descriptor.getConfig().getThriftMaxFrameSize() - 1024; + properties.setProperty( + "mpp_data_exchange_max_payload_size_in_bytes", Integer.toString(Integer.MAX_VALUE)); + descriptor.loadHotModifiedProps(properties); + Assert.assertEquals( + maximumPayloadSize, descriptor.getConfig().getMppDataExchangeMaxPayloadSizeInBytes()); + } finally { + descriptor.getConfig().setMppDataExchangeMaxPayloadSizeInBytes(originalPayloadSize); + } + } + @Test public void PropertiesWithSpace() { IoTDBDescriptor descriptor = IoTDBDescriptor.getInstance(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index 2e8c8a8c9d366..d959d5ef16128 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -41,6 +41,7 @@ import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import java.nio.ByteBuffer; @@ -48,6 +49,7 @@ import java.util.List; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -70,6 +72,137 @@ public static void afterClass() { IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(maxBytesPerFI); } + @Test + public void testFragmentAssemblyUsesOnlyByteBufferRemainingBytes() { + final String queryId = "q0"; + final TEndPoint remoteEndpoint = + new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); + final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); + final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); + + LocalMemoryManager localMemoryManager = Mockito.mock(LocalMemoryManager.class); + MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool(); + Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool); + SourceHandleListener sourceHandleListener = Mockito.mock(SourceHandleListener.class); + TsBlockSerde serde = Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE); + IClientManager clientManager = + Mockito.mock(IClientManager.class); + SyncDataNodeMPPDataExchangeServiceClient client = + Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); + AtomicInteger rpcCount = new AtomicInteger(); + try { + Mockito.when(clientManager.borrowClient(remoteEndpoint)).thenReturn(client); + Mockito.doAnswer( + invocation -> { + if (rpcCount.getAndIncrement() == 0) { + ByteBuffer firstFragment = ByteBuffer.wrap(new byte[] {99, 1, 2, 98}, 1, 2); + return new TGetDataBlockResponse(List.of(firstFragment)) + .setOffset(2) + .setTotalLength(4); + } + ByteBuffer secondFragment = + ByteBuffer.wrap(new byte[] {97, 3, 4, 96}, 1, 2).slice(); + return new TGetDataBlockResponse(List.of(secondFragment)); + }) + .when(client) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + } catch (ClientManagerException | TException e) { + Assert.fail(e.getMessage()); + } + + SourceHandle sourceHandle = + new SourceHandle( + remoteEndpoint, + remoteFragmentInstanceId, + localFragmentInstanceId, + "exchange_0", + 0, + localMemoryManager, + Executors.newSingleThreadExecutor(), + serde, + sourceHandleListener, + clientManager); + Assert.assertFalse(sourceHandle.isBlocked().isDone()); + sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); + try { + Mockito.verify(client, Mockito.timeout(10_000).times(2)) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + } catch (TException e) { + Assert.fail(e.getMessage()); + } + + sourceHandle.receive(); + ArgumentCaptor serializedBlock = ArgumentCaptor.forClass(ByteBuffer.class); + Mockito.verify(serde).deserialize(serializedBlock.capture()); + ByteBuffer actual = serializedBlock.getValue().duplicate(); + Assert.assertEquals(4, actual.remaining()); + byte[] actualBytes = new byte[actual.remaining()]; + actual.get(actualBytes); + Assert.assertArrayEquals(new byte[] {1, 2, 3, 4}, actualBytes); + sourceHandle.abort(); + } + + @Test + public void testFragmentAssemblyRejectsUnusedPublicBaosCapacity() { + final String queryId = "q0"; + final TEndPoint remoteEndpoint = + new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); + final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); + final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); + + LocalMemoryManager localMemoryManager = Mockito.mock(LocalMemoryManager.class); + MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool(); + Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool); + SourceHandleListener sourceHandleListener = Mockito.mock(SourceHandleListener.class); + TsBlockSerde serde = Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE); + IClientManager clientManager = + Mockito.mock(IClientManager.class); + SyncDataNodeMPPDataExchangeServiceClient client = + Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); + AtomicInteger rpcCount = new AtomicInteger(); + try { + Mockito.when(clientManager.borrowClient(remoteEndpoint)).thenReturn(client); + Mockito.doAnswer( + invocation -> { + if (rpcCount.getAndIncrement() == 0) { + return new TGetDataBlockResponse(List.of(ByteBuffer.wrap(new byte[] {1, 2}))) + .setOffset(2) + .setTotalLength(5); + } + return new TGetDataBlockResponse(List.of(ByteBuffer.wrap(new byte[] {3, 4}))); + }) + .when(client) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + } catch (ClientManagerException | TException e) { + Assert.fail(e.getMessage()); + } + + SourceHandle sourceHandle = + new SourceHandle( + remoteEndpoint, + remoteFragmentInstanceId, + localFragmentInstanceId, + "exchange_0", + 0, + localMemoryManager, + Executors.newSingleThreadExecutor(), + serde, + sourceHandleListener, + clientManager); + Assert.assertFalse(sourceHandle.isBlocked().isDone()); + sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); + try { + Mockito.verify(client, Mockito.timeout(10_000).times(2)) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + } catch (TException e) { + Assert.fail(e.getMessage()); + } + Mockito.verify(sourceHandleListener, Mockito.timeout(10_000)) + .onFailure(Mockito.eq(sourceHandle), Mockito.any(IllegalArgumentException.class)); + Mockito.verify(serde, Mockito.never()).deserialize(Mockito.any(ByteBuffer.class)); + sourceHandle.abort(); + } + @Test public void testNonBlockedOneTimeReceive() { final String queryId = "q0"; @@ -634,15 +767,16 @@ public void testShortResponseRetriesAndFails() { Mockito.doAnswer( invocation -> { final TGetDataBlockRequest request = invocation.getArgument(0); + if (request.getStartSequenceId() > 0) { + throw new TException("mock RPC failure"); + } final List shortResponse = new ArrayList<>(); for (int i = 0; i < request.getEndSequenceId() - request.getStartSequenceId() - 1; i++) { shortResponse.add(ByteBuffer.allocate(1)); } - TGetDataBlockResponse response = new TGetDataBlockResponse(shortResponse); - response.setOffset(1); - return response; + return new TGetDataBlockResponse(shortResponse); }) .when(mockClient) .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); @@ -672,13 +806,13 @@ public void testShortResponseRetriesAndFails() { .collect(Collectors.toList())); try { - Mockito.verify(mockClient, Mockito.timeout(10_000).times(SourceHandle.MAX_ATTEMPT_TIMES)) + Mockito.verify(mockClient, Mockito.timeout(10_000).times(SourceHandle.MAX_ATTEMPT_TIMES + 1)) .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); } catch (TException e) { Assert.fail(e.getMessage()); } Mockito.verify(mockSourceHandleListener, Mockito.timeout(10_000).times(1)) - .onFailure(Mockito.eq(sourceHandle), Mockito.any(Throwable.class)); + .onFailure(Mockito.eq(sourceHandle), Mockito.any(TException.class)); Assert.assertFalse(blocked.isDone()); Assert.assertEquals(0L, sourceHandle.getBufferRetainedSizeInBytes()); diff --git a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template index 2a51650c0cc51..161e0380dc176 100644 --- a/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template +++ b/iotdb-core/node-commons/src/assembly/resources/conf/iotdb-system.properties.template @@ -71,7 +71,7 @@ cn_consensus_port=10720 # Used for connection of IoTDB native clients(Session) # Could set 127.0.0.1(for local test), ipv4/ipv6 address, or hostname. -# effectiveMode: hot_reload +# effectiveMode: restart # Datatype: String dn_rpc_address=127.0.0.1 @@ -1196,7 +1196,7 @@ mpp_data_exchange_keep_alive_time_in_ms=1000 # The maximum payload size of one MPP data exchange RPC response. # <=0 use default value -# The effective value is capped by dn_thrift_max_frame_size minus 1024 bytes reserved for the RPC response envelope. +# The effective value is between 131072(128KB) and dn_thrift_max_frame_size minus 1024 bytes reserved for the RPC response envelope. # effectiveMode: hot_reload # Datatype: int, Unit: byte mpp_data_exchange_max_payload_size_in_bytes=4194304 From c2cd98518b5be55353717d97b32f8b7ab3954b36 Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Thu, 17 Sep 2026 11:24:15 +0800 Subject: [PATCH 09/12] clear check message Signed-off-by: Weihao Li <18110526956@163.com> --- .../iotdb/db/i18n/DataNodeQueryMessages.java | 22 +++++- .../iotdb/db/i18n/DataNodeQueryMessages.java | 22 +++++- .../exchange/MPPDataExchangeManager.java | 4 +- .../exchange/source/SourceHandle.java | 30 +++++--- .../execution/exchange/SourceHandleTest.java | 75 +------------------ 5 files changed, 66 insertions(+), 87 deletions(-) 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 946e39c4874c6..5f9750bd2a5c1 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 @@ -1412,8 +1412,26 @@ public final class DataNodeQueryMessages { "failed to get data block [{}, {}), attempt times: {}"; public static final String EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33 = "Unexpected data block response size."; - public static final String EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_RANGE_0672002F = - "Invalid serialized TsBlock fragment range"; + public static final String + EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_OFFSET_ARG_FOR_BLOCK_SIZE_ARG_53BC0284 = + "Invalid serialized TsBlock fragment offset %s for block size %s."; + public static final String + EXCEPTION_ACCUMULATED_TSBLOCK_FRAGMENT_LENGTH_ARG_DOES_NOT_MATCH_TOTALLENGTH_ARG_1B784303 = + "Accumulated TsBlock fragment length %s does not match totalLength %s."; + public static final String + EXCEPTION_THE_FIRST_FRAGMENTED_DATA_BLOCK_RESPONSE_MUST_INCLUDE_TOTALLENGTH_C5C79BC2 = + "The first fragmented data block response must include totalLength."; + public static final String + EXCEPTION_NEXT_SEQUENCE_ID_ARG_EXCEEDS_REQUESTED_END_SEQUENCE_ID_ARG_30B1726E = + "Next sequence ID %s exceeds requested end sequence ID %s."; + public static final String + EXCEPTION_A_COMPLETED_DATA_BLOCK_RESPONSE_RANGE_MUST_NOT_RETAIN_A_PARTIAL_TSBLOCK_85E5C287 = + "A completed data block response range must not retain a partial TsBlock."; + public static final String EXCEPTION_TSBLOCK_FRAGMENT_MUST_NOT_BE_EMPTY_C7D19863 = + "TsBlock fragment must not be empty."; + public static final String + EXCEPTION_NEXT_FRAGMENT_OFFSET_ARG_MUST_BE_GREATER_THAN_CURRENT_OFFSET_ARG_AND_MATCH_ACCUMULATED_FRAGMENT_LENGTH_ARG_ECC31047 = + "Next fragment offset %s must be greater than current offset %s and match accumulated fragment length %s."; public static final String FAILED_TO_SEND_ACK_DATA_BLOCK_EVENT = "failed to send ack data block event [{}, {}), attempt times: {}"; public static final String SEND_CLOSE_SINK_CHANNEL_EVENT_FAILED = 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 940a19cc86036..eba001e39cce3 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 @@ -1394,8 +1394,26 @@ public final class DataNodeQueryMessages { "获取数据块 [{}, {}) 失败，尝试次数：{}"; public static final String EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33 = "数据块响应数量异常。"; - public static final String EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_RANGE_0672002F = - "无效的序列化 TsBlock 分片范围"; + public static final String + EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_OFFSET_ARG_FOR_BLOCK_SIZE_ARG_53BC0284 = + "序列化 TsBlock 分片偏移量 %s 无效，块大小为 %s。"; + public static final String + EXCEPTION_ACCUMULATED_TSBLOCK_FRAGMENT_LENGTH_ARG_DOES_NOT_MATCH_TOTALLENGTH_ARG_1B784303 = + "TsBlock 分片累计长度 %s 与 totalLength %s 不一致。"; + public static final String + EXCEPTION_THE_FIRST_FRAGMENTED_DATA_BLOCK_RESPONSE_MUST_INCLUDE_TOTALLENGTH_C5C79BC2 = + "首个分片数据块响应必须包含 totalLength。"; + public static final String + EXCEPTION_NEXT_SEQUENCE_ID_ARG_EXCEEDS_REQUESTED_END_SEQUENCE_ID_ARG_30B1726E = + "下一个 sequence ID %s 超过请求的结束 sequence ID %s。"; + public static final String + EXCEPTION_A_COMPLETED_DATA_BLOCK_RESPONSE_RANGE_MUST_NOT_RETAIN_A_PARTIAL_TSBLOCK_85E5C287 = + "已完成的数据块响应区间不能保留未完成的 TsBlock。"; + public static final String EXCEPTION_TSBLOCK_FRAGMENT_MUST_NOT_BE_EMPTY_C7D19863 = + "TsBlock 分片不能为空。"; + public static final String + EXCEPTION_NEXT_FRAGMENT_OFFSET_ARG_MUST_BE_GREATER_THAN_CURRENT_OFFSET_ARG_AND_MATCH_ACCUMULATED_FRAGMENT_LENGTH_ARG_ECC31047 = + "下一分片偏移量 %s 必须大于当前偏移量 %s，并且等于分片累计长度 %s。"; public static final String FAILED_TO_SEND_ACK_DATA_BLOCK_EVENT = "发送数据块确认事件 [{}, {}) 失败，尝试次数：{}"; public static final String SEND_CLOSE_SINK_CHANNEL_EVENT_FAILED = diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index 4f9653ccd0718..6f5f6088c09bf 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -190,7 +190,9 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce checkArgument( blockOffset >= 0 && blockOffset <= serializedTsBlockSize, DataNodeQueryMessages - .EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_RANGE_0672002F); + .EXCEPTION_INVALID_SERIALIZED_TSBLOCK_FRAGMENT_OFFSET_ARG_FOR_BLOCK_SIZE_ARG_53BC0284, + blockOffset, + serializedTsBlockSize); int remainingBlockSize = serializedTsBlockSize - blockOffset; if (remainingBlockSize <= remainingPayloadSize) { ByteBuffer fragment = serializedTsBlock.duplicate(); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index a48da4f066e14..2f02b4fb9d9c2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -829,7 +829,10 @@ private void addResponse(TGetDataBlockResponse response) { } Preconditions.checkState( partialTsBlock.size() == partialTsBlockTotalLength, - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + DataNodeQueryMessages + .EXCEPTION_ACCUMULATED_TSBLOCK_FRAGMENT_LENGTH_ARG_DOES_NOT_MATCH_TOTALLENGTH_ARG_1B784303, + partialTsBlock.size(), + partialTsBlockTotalLength); tsBlocks.add(ByteBuffer.wrap(partialTsBlock.getBuf(), 0, partialTsBlock.size())); partialTsBlock = null; partialTsBlockTotalLength = 0; @@ -847,31 +850,40 @@ private void addResponse(TGetDataBlockResponse response) { if (lastBlockIsFragment) { checkArgument( response.isSetTotalLength(), - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + DataNodeQueryMessages + .EXCEPTION_THE_FIRST_FRAGMENTED_DATA_BLOCK_RESPONSE_MUST_INCLUDE_TOTALLENGTH_C5C79BC2); partialTsBlock = new PublicBAOS(response.getTotalLength()); appendFragment(responseBlocks.get(blockIndex)); updateOffset(response.getOffset()); } Preconditions.checkState( - nextSequenceId <= endSequenceId - && (lastBlockIsFragment - || nextSequenceId != endSequenceId - || partialTsBlock == null), - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + nextSequenceId <= endSequenceId, + DataNodeQueryMessages + .EXCEPTION_NEXT_SEQUENCE_ID_ARG_EXCEEDS_REQUESTED_END_SEQUENCE_ID_ARG_30B1726E, + nextSequenceId, + endSequenceId); + Preconditions.checkState( + lastBlockIsFragment || nextSequenceId != endSequenceId || partialTsBlock == null, + DataNodeQueryMessages + .EXCEPTION_A_COMPLETED_DATA_BLOCK_RESPONSE_RANGE_MUST_NOT_RETAIN_A_PARTIAL_TSBLOCK_85E5C287); } private void appendFragment(ByteBuffer fragment) { checkArgument( fragment.hasRemaining(), - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + DataNodeQueryMessages.EXCEPTION_TSBLOCK_FRAGMENT_MUST_NOT_BE_EMPTY_C7D19863); partialTsBlock.writeBytes(fragment.array()); } private void updateOffset(int nextOffset) { checkArgument( nextOffset > offset && nextOffset == partialTsBlock.size(), - DataNodeQueryMessages.EXCEPTION_UNEXPECTED_DATA_BLOCK_RESPONSE_SIZE_A7DD7E33); + DataNodeQueryMessages + .EXCEPTION_NEXT_FRAGMENT_OFFSET_ARG_MUST_BE_GREATER_THAN_CURRENT_OFFSET_ARG_AND_MATCH_ACCUMULATED_FRAGMENT_LENGTH_ARG_ECC31047, + nextOffset, + offset, + partialTsBlock.size()); offset = nextOffset; } } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index d959d5ef16128..047fe5162fedb 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -41,7 +41,6 @@ import org.junit.Assert; import org.junit.BeforeClass; import org.junit.Test; -import org.mockito.ArgumentCaptor; import org.mockito.Mockito; import java.nio.ByteBuffer; @@ -73,77 +72,7 @@ public static void afterClass() { } @Test - public void testFragmentAssemblyUsesOnlyByteBufferRemainingBytes() { - final String queryId = "q0"; - final TEndPoint remoteEndpoint = - new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); - final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); - final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); - - LocalMemoryManager localMemoryManager = Mockito.mock(LocalMemoryManager.class); - MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool(); - Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool); - SourceHandleListener sourceHandleListener = Mockito.mock(SourceHandleListener.class); - TsBlockSerde serde = Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE); - IClientManager clientManager = - Mockito.mock(IClientManager.class); - SyncDataNodeMPPDataExchangeServiceClient client = - Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); - AtomicInteger rpcCount = new AtomicInteger(); - try { - Mockito.when(clientManager.borrowClient(remoteEndpoint)).thenReturn(client); - Mockito.doAnswer( - invocation -> { - if (rpcCount.getAndIncrement() == 0) { - ByteBuffer firstFragment = ByteBuffer.wrap(new byte[] {99, 1, 2, 98}, 1, 2); - return new TGetDataBlockResponse(List.of(firstFragment)) - .setOffset(2) - .setTotalLength(4); - } - ByteBuffer secondFragment = - ByteBuffer.wrap(new byte[] {97, 3, 4, 96}, 1, 2).slice(); - return new TGetDataBlockResponse(List.of(secondFragment)); - }) - .when(client) - .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); - } catch (ClientManagerException | TException e) { - Assert.fail(e.getMessage()); - } - - SourceHandle sourceHandle = - new SourceHandle( - remoteEndpoint, - remoteFragmentInstanceId, - localFragmentInstanceId, - "exchange_0", - 0, - localMemoryManager, - Executors.newSingleThreadExecutor(), - serde, - sourceHandleListener, - clientManager); - Assert.assertFalse(sourceHandle.isBlocked().isDone()); - sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); - try { - Mockito.verify(client, Mockito.timeout(10_000).times(2)) - .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); - } catch (TException e) { - Assert.fail(e.getMessage()); - } - - sourceHandle.receive(); - ArgumentCaptor serializedBlock = ArgumentCaptor.forClass(ByteBuffer.class); - Mockito.verify(serde).deserialize(serializedBlock.capture()); - ByteBuffer actual = serializedBlock.getValue().duplicate(); - Assert.assertEquals(4, actual.remaining()); - byte[] actualBytes = new byte[actual.remaining()]; - actual.get(actualBytes); - Assert.assertArrayEquals(new byte[] {1, 2, 3, 4}, actualBytes); - sourceHandle.abort(); - } - - @Test - public void testFragmentAssemblyRejectsUnusedPublicBaosCapacity() { + public void testFragmentAssemblyFailsWhenTotalLengthDoesNotMatch() { final String queryId = "q0"; final TEndPoint remoteEndpoint = new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); @@ -198,7 +127,7 @@ public void testFragmentAssemblyRejectsUnusedPublicBaosCapacity() { Assert.fail(e.getMessage()); } Mockito.verify(sourceHandleListener, Mockito.timeout(10_000)) - .onFailure(Mockito.eq(sourceHandle), Mockito.any(IllegalArgumentException.class)); + .onFailure(Mockito.eq(sourceHandle), Mockito.any(IllegalStateException.class)); Mockito.verify(serde, Mockito.never()).deserialize(Mockito.any(ByteBuffer.class)); sourceHandle.abort(); } From 4b404aed1654722b7273acef31fc877e4895d88f Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Thu, 17 Sep 2026 11:39:09 +0800 Subject: [PATCH 10/12] remove some Signed-off-by: Weihao Li <18110526956@163.com> --- .../execution/exchange/MPPDataExchangeManager.java | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java index 6f5f6088c09bf..61e776d7d9e89 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/MPPDataExchangeManager.java @@ -194,9 +194,9 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce blockOffset, serializedTsBlockSize); int remainingBlockSize = serializedTsBlockSize - blockOffset; + ByteBuffer fragment = serializedTsBlock; + fragment.position(blockOffset); if (remainingBlockSize <= remainingPayloadSize) { - ByteBuffer fragment = serializedTsBlock.duplicate(); - fragment.position(blockOffset); fragment.limit(blockOffset + remainingBlockSize); resp.addToTsBlocks(fragment.slice()); remainingPayloadSize -= remainingBlockSize; @@ -204,8 +204,6 @@ public TGetDataBlockResponse getDataBlock(TGetDataBlockRequest req) throws TExce break; } } else { - ByteBuffer fragment = serializedTsBlock.duplicate(); - fragment.position(blockOffset); fragment.limit(blockOffset + remainingPayloadSize); resp.addToTsBlocks(fragment.slice()); resp.setOffset(blockOffset + remainingPayloadSize); From e794159c67216f63461d71e073b39e7aea76396d Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Thu, 17 Sep 2026 12:29:31 +0800 Subject: [PATCH 11/12] fix some and add IT Signed-off-by: Weihao Li <18110526956@163.com> --- .../env/cluster/config/MppDataNodeConfig.java | 9 ++ .../remote/config/RemoteDataNodeConfig.java | 6 ++ .../iotdb/itbase/env/DataNodeConfig.java | 2 + ...actIoTDBMPPDataExchangeLargeTsBlockIT.java | 98 +++++++++++++++++++ ...oTDBMPPDataExchangeLargeTsBlock128KIT.java | 51 ++++++++++ .../exchange/source/SourceHandle.java | 3 +- 6 files changed, 168 insertions(+), 1 deletion(-) create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/AbstractIoTDBMPPDataExchangeLargeTsBlockIT.java create mode 100644 integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBMPPDataExchangeLargeTsBlock128KIT.java diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java index 8c5f56ad9f1f0..f7dd625704cee 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/cluster/config/MppDataNodeConfig.java @@ -186,4 +186,13 @@ public DataNodeConfig setTableQueryDeviceEntryBatchSizeInBytes(long batchSizeInB setProperty("table_query_device_entry_batch_size_in_bytes", String.valueOf(batchSizeInBytes)); return this; } + + @Override + public DataNodeConfig setMppDataExchangeMaxPayloadSizeInBytes( + int mppDataExchangeMaxPayloadSizeInBytes) { + setProperty( + "mpp_data_exchange_max_payload_size_in_bytes", + String.valueOf(mppDataExchangeMaxPayloadSizeInBytes)); + return this; + } } diff --git a/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java b/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java index c97e2b5065af6..13eeb34391920 100644 --- a/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java +++ b/integration-test/src/main/java/org/apache/iotdb/it/env/remote/config/RemoteDataNodeConfig.java @@ -130,4 +130,10 @@ public DataNodeConfig setDnMultiDirStrategy(String multiDirStrategy) { public DataNodeConfig setTableQueryDeviceEntryBatchSizeInBytes(long batchSizeInBytes) { return this; } + + @Override + public DataNodeConfig setMppDataExchangeMaxPayloadSizeInBytes( + int mppDataExchangeMaxPayloadSizeInBytes) { + return this; + } } diff --git a/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java b/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java index fb969c778ea8f..d321053f2787e 100644 --- a/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java +++ b/integration-test/src/main/java/org/apache/iotdb/itbase/env/DataNodeConfig.java @@ -67,4 +67,6 @@ DataNodeConfig setLoadActiveListeningCheckIntervalSeconds( DataNodeConfig setDnMultiDirStrategy(String multiDirStrategy); DataNodeConfig setTableQueryDeviceEntryBatchSizeInBytes(long batchSizeInBytes); + + DataNodeConfig setMppDataExchangeMaxPayloadSizeInBytes(int mppDataExchangeMaxPayloadSizeInBytes); } diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/AbstractIoTDBMPPDataExchangeLargeTsBlockIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/AbstractIoTDBMPPDataExchangeLargeTsBlockIT.java new file mode 100644 index 0000000000000..34e7798d73aa1 --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/AbstractIoTDBMPPDataExchangeLargeTsBlockIT.java @@ -0,0 +1,98 @@ +/* + * 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.relational.it.query.recent; + +import org.apache.iotdb.isession.ITableSession; +import org.apache.iotdb.it.env.EnvFactory; + +import org.apache.tsfile.enums.ColumnCategory; +import org.apache.tsfile.enums.TSDataType; +import org.apache.tsfile.utils.Binary; +import org.apache.tsfile.write.record.Tablet; +import org.apache.tsfile.write.schema.IMeasurementSchema; +import org.apache.tsfile.write.schema.MeasurementSchema; +import org.junit.Test; + +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.Statement; +import java.util.List; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public abstract class AbstractIoTDBMPPDataExchangeLargeTsBlockIT { + + protected static final String DATABASE_NAME = "large_tsblock"; + protected static final int PAYLOAD_SIZE_IN_BYTES = 128 * 1024; + private static final int BLOB_SIZE_IN_BYTES = 2 * PAYLOAD_SIZE_IN_BYTES + 1; + private static final byte[] EXPECTED_BLOB = createBlob(); + + protected static void prepareData() throws Exception { + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("CREATE DATABASE " + DATABASE_NAME); + statement.execute("USE " + DATABASE_NAME); + statement.execute("CREATE TABLE large_blob(payload BLOB FIELD)"); + } + + List schemas = List.of(new MeasurementSchema("payload", TSDataType.BLOB)); + Tablet tablet = + new Tablet( + "large_blob", + IMeasurementSchema.getMeasurementNameList(schemas), + IMeasurementSchema.getDataTypeList(schemas), + List.of(ColumnCategory.FIELD), + 1); + tablet.addTimestamp(0, 1); + tablet.addValue("payload", 0, new Binary(EXPECTED_BLOB)); + + try (ITableSession session = + EnvFactory.getEnv().getTableSessionConnectionWithDB(DATABASE_NAME)) { + session.insert(tablet); + } + } + + @Test + public void testLargeBlobTransferredInFragments() throws Exception { + assertTrue(EXPECTED_BLOB.length > PAYLOAD_SIZE_IN_BYTES); + + try (Connection connection = EnvFactory.getEnv().getTableConnection(); + Statement statement = connection.createStatement()) { + statement.execute("USE " + DATABASE_NAME); + try (ResultSet resultSet = statement.executeQuery("SELECT time, payload FROM large_blob")) { + assertTrue(resultSet.next()); + assertEquals(1, resultSet.getLong("time")); + assertArrayEquals(EXPECTED_BLOB, resultSet.getBytes("payload")); + assertFalse(resultSet.next()); + } + } + } + + private static byte[] createBlob() { + byte[] blob = new byte[BLOB_SIZE_IN_BYTES]; + for (int i = 0; i < blob.length; i++) { + blob[i] = (byte) (i * 31 + 7); + } + return blob; + } +} diff --git a/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBMPPDataExchangeLargeTsBlock128KIT.java b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBMPPDataExchangeLargeTsBlock128KIT.java new file mode 100644 index 0000000000000..672859429f39a --- /dev/null +++ b/integration-test/src/test/java/org/apache/iotdb/relational/it/query/recent/IoTDBMPPDataExchangeLargeTsBlock128KIT.java @@ -0,0 +1,51 @@ +/* + * 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.relational.it.query.recent; + +import org.apache.iotdb.it.env.EnvFactory; +import org.apache.iotdb.it.framework.IoTDBTestRunner; +import org.apache.iotdb.itbase.category.TableClusterIT; + +import org.junit.AfterClass; +import org.junit.BeforeClass; +import org.junit.experimental.categories.Category; +import org.junit.runner.RunWith; + +@RunWith(IoTDBTestRunner.class) +@Category({TableClusterIT.class}) +public class IoTDBMPPDataExchangeLargeTsBlock128KIT + extends AbstractIoTDBMPPDataExchangeLargeTsBlockIT { + + @BeforeClass + public static void setUp() throws Exception { + EnvFactory.getEnv() + .getConfig() + .getDataNodeConfig() + .setMppDataExchangeMaxPayloadSizeInBytes(PAYLOAD_SIZE_IN_BYTES); + EnvFactory.getEnv().getConfig().getCommonConfig().setDataReplicationFactor(1); + EnvFactory.getEnv().initClusterEnvironment(); + prepareData(); + } + + @AfterClass + public static void tearDown() { + EnvFactory.getEnv().cleanClusterEnvironment(); + } +} diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index 2f02b4fb9d9c2..dab35c18d05f2 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -852,7 +852,8 @@ private void addResponse(TGetDataBlockResponse response) { response.isSetTotalLength(), DataNodeQueryMessages .EXCEPTION_THE_FIRST_FRAGMENTED_DATA_BLOCK_RESPONSE_MUST_INCLUDE_TOTALLENGTH_C5C79BC2); - partialTsBlock = new PublicBAOS(response.getTotalLength()); + partialTsBlockTotalLength = response.getTotalLength(); + partialTsBlock = new PublicBAOS(partialTsBlockTotalLength); appendFragment(responseBlocks.get(blockIndex)); updateOffset(response.getOffset()); } From c02cd116b7b13fbc06b60e2afee5d442a334ff71 Mon Sep 17 00:00:00 2001 From: Weihao Li <18110526956@163.com> Date: Thu, 17 Sep 2026 16:56:49 +0800 Subject: [PATCH 12/12] fix according review Signed-off-by: Weihao Li <18110526956@163.com> --- .../apache/iotdb/db/conf/IoTDBDescriptor.java | 31 ++--- .../exchange/source/SourceHandle.java | 39 +++++- .../execution/exchange/SourceHandleTest.java | 116 ++++++++++++++++++ .../audit/UserDataTransferErrorCode.java | 3 +- 4 files changed, 171 insertions(+), 18 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java index f46e02fc18a6f..8be34fc8ec711 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/conf/IoTDBDescriptor.java @@ -3100,28 +3100,31 @@ private void loadMppDataExchangeMaxPayloadSize(TrimProperties properties) { Integer.toString(conf.getMppDataExchangeMaxPayloadSizeInBytes()))); if (configuredSize <= 0) { LOGGER.warn( - DataNodeMiscMessages - .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_NOT_POSITIVE_USING_DEFAULT_VALUE_ARG_1AA821B2, - configuredSize, - DEFAULT_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES); + String.format( + DataNodeMiscMessages + .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_NOT_POSITIVE_USING_DEFAULT_VALUE_ARG_1AA821B2, + configuredSize, + DEFAULT_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES)); configuredSize = DEFAULT_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES; } else if (configuredSize < MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES) { LOGGER.warn( - DataNodeMiscMessages - .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_BELOW_MINIMUM_ALLOWED_VALUE_ARG_USING_ARG_794ABC76, - configuredSize, - MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES, - MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES); + String.format( + DataNodeMiscMessages + .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_IS_BELOW_MINIMUM_ALLOWED_VALUE_ARG_USING_ARG_794ABC76, + configuredSize, + MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES, + MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES)); configuredSize = MIN_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_IN_BYTES; } int maxAllowedSize = conf.getThriftMaxFrameSize() - 1024; if (configuredSize > maxAllowedSize) { LOGGER.warn( - DataNodeMiscMessages - .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC, - configuredSize, - maxAllowedSize, - maxAllowedSize); + String.format( + DataNodeMiscMessages + .LOG_MPP_DATA_EXCHANGE_MAX_PAYLOAD_SIZE_ARG_EXCEEDS_MAXIMUM_ALLOWED_VALUE_ARG_USING_ARG_D9BF0BBC, + configuredSize, + maxAllowedSize, + maxAllowedSize)); configuredSize = maxAllowedSize; } conf.setMppDataExchangeMaxPayloadSizeInBytes(configuredSize); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java index dab35c18d05f2..aed3f9533fb1d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/queryengine/execution/exchange/source/SourceHandle.java @@ -654,6 +654,9 @@ public void run() { try (SyncDataNodeMPPDataExchangeServiceClient client = mppDataExchangeServiceClientManager.borrowClient(remoteEndpoint)) { TGetDataBlockResponse resp = getDataBlockWithFragments(client, req, fetchProgress); + if (resp == null) { + return; + } int tsBlockNum = resp.getTsBlocks().size(); if (tsBlockNum != endSequenceId - startSequenceId) { recordTransferAttempt( @@ -721,6 +724,17 @@ public void run() { } fail(e); return; + } catch (Error e) { + if (!transferAttemptRecorded) { + recordTransferAttempt( + false, + e instanceof OutOfMemoryError + ? UserDataTransferErrorCode.OUT_OF_MEMORY.name() + : null, + e); + } + fail(e); + throw e; } catch (Exception e) { if (!transferAttemptRecorded) { @@ -786,14 +800,26 @@ private TGetDataBlockResponse getDataBlockWithFragments( DataBlockFetchProgress fetchProgress) throws TException { while (!fetchProgress.isFinished()) { + synchronized (SourceHandle.this) { + if (aborted || closed) { + fetchProgress.discard(); + return null; + } + } TGetDataBlockRequest fragmentRequest = request.deepCopy(); fragmentRequest.setStartSequenceId(fetchProgress.nextSequenceId); fragmentRequest.setOffset(fetchProgress.offset); TGetDataBlockResponse response = client.getDataBlock(fragmentRequest); - if (response.getTsBlocks().isEmpty()) { - return response; + synchronized (SourceHandle.this) { + if (aborted || closed) { + fetchProgress.discard(); + return null; + } + if (response.getTsBlocks().isEmpty()) { + return response; + } + fetchProgress.addResponse(response); } - fetchProgress.addResponse(response); } return new TGetDataBlockResponse(fetchProgress.tsBlocks); } @@ -816,6 +842,13 @@ private boolean isFinished() { return nextSequenceId == endSequenceId && partialTsBlock == null; } + private void discard() { + tsBlocks.clear(); + partialTsBlock = null; + partialTsBlockTotalLength = 0; + offset = 0; + } + private void addResponse(TGetDataBlockResponse response) { List responseBlocks = response.getTsBlocks(); boolean lastBlockIsFragment = response.isSetOffset(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java index 047fe5162fedb..af96360b12e7d 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/queryengine/execution/exchange/SourceHandleTest.java @@ -46,8 +46,10 @@ import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.List; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import java.util.stream.Collectors; import java.util.stream.Stream; @@ -71,6 +73,120 @@ public static void afterClass() { IoTDBDescriptor.getInstance().getMemoryConfig().setMaxBytesPerFragmentInstance(maxBytesPerFI); } + @Test + public void testCloseStopsFragmentFetchAfterInFlightRpc() throws Exception { + testCancellationStopsFragmentFetchAfterInFlightRpc(false); + } + + @Test + public void testAbortStopsFragmentFetchAfterInFlightRpc() throws Exception { + testCancellationStopsFragmentFetchAfterInFlightRpc(true); + } + + private void testCancellationStopsFragmentFetchAfterInFlightRpc(boolean abort) throws Exception { + final String queryId = "q0"; + final TEndPoint remoteEndpoint = + new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); + final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); + final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); + final LocalMemoryManager localMemoryManager = Mockito.mock(LocalMemoryManager.class); + final MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool(); + Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool); + final SourceHandleListener sourceHandleListener = Mockito.mock(SourceHandleListener.class); + final IClientManager clientManager = + Mockito.mock(IClientManager.class); + final SyncDataNodeMPPDataExchangeServiceClient client = + Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); + final CountDownLatch rpcStarted = new CountDownLatch(1); + final CountDownLatch returnResponse = new CountDownLatch(1); + Mockito.when(clientManager.borrowClient(remoteEndpoint)).thenReturn(client); + Mockito.doAnswer( + invocation -> { + rpcStarted.countDown(); + Assert.assertTrue(returnResponse.await(10, TimeUnit.SECONDS)); + return new TGetDataBlockResponse(List.of(ByteBuffer.wrap(new byte[] {1, 2}))) + .setOffset(2) + .setTotalLength(6); + }) + .when(client) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + + final SourceHandle sourceHandle = + new SourceHandle( + remoteEndpoint, + remoteFragmentInstanceId, + localFragmentInstanceId, + "exchange_0", + 0, + localMemoryManager, + Executors.newSingleThreadExecutor(), + Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE), + sourceHandleListener, + clientManager); + sourceHandle.isBlocked(); + sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); + Assert.assertTrue(rpcStarted.await(10, TimeUnit.SECONDS)); + + if (abort) { + sourceHandle.abort(); + } else { + sourceHandle.close(); + } + returnResponse.countDown(); + + Mockito.verify(client, Mockito.timeout(10_000).times(1)) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + Mockito.verify(client, Mockito.after(500).times(1)) + .getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + Mockito.verify(sourceHandleListener, Mockito.never()) + .onFailure(Mockito.eq(sourceHandle), Mockito.any()); + Assert.assertEquals(0L, sourceHandle.getBufferRetainedSizeInBytes()); + } + + @Test + public void testFragmentAllocationErrorFailsWithoutRetry() throws Exception { + final String queryId = "q0"; + final TEndPoint remoteEndpoint = + new TEndPoint("remote", IoTDBDescriptor.getInstance().getConfig().getMppDataExchangePort()); + final TFragmentInstanceId remoteFragmentInstanceId = new TFragmentInstanceId(queryId, 1, "0"); + final TFragmentInstanceId localFragmentInstanceId = new TFragmentInstanceId(queryId, 0, "0"); + final LocalMemoryManager localMemoryManager = Mockito.mock(LocalMemoryManager.class); + final MemoryPool memoryPool = Utils.createMockNonBlockedMemoryPool(); + Mockito.when(localMemoryManager.getQueryPool()).thenReturn(memoryPool); + final SourceHandleListener sourceHandleListener = Mockito.mock(SourceHandleListener.class); + final IClientManager clientManager = + Mockito.mock(IClientManager.class); + final SyncDataNodeMPPDataExchangeServiceClient client = + Mockito.mock(SyncDataNodeMPPDataExchangeServiceClient.class); + Mockito.when(clientManager.borrowClient(remoteEndpoint)).thenReturn(client); + Mockito.when(client.getDataBlock(Mockito.any(TGetDataBlockRequest.class))) + .thenReturn( + new TGetDataBlockResponse(List.of(ByteBuffer.wrap(new byte[] {1}))) + .setOffset(1) + .setTotalLength(Integer.MAX_VALUE)); + + final SourceHandle sourceHandle = + new SourceHandle( + remoteEndpoint, + remoteFragmentInstanceId, + localFragmentInstanceId, + "exchange_0", + 0, + localMemoryManager, + Executors.newSingleThreadExecutor(), + Utils.createMockTsBlockSerde(MOCK_TSBLOCK_SIZE), + sourceHandleListener, + clientManager); + sourceHandle.isBlocked(); + sourceHandle.updatePendingDataBlockInfo(0, List.of(MOCK_TSBLOCK_SIZE)); + + Mockito.verify(sourceHandleListener, Mockito.timeout(10_000)) + .onFailure(Mockito.eq(sourceHandle), Mockito.any(OutOfMemoryError.class)); + Mockito.verify(client, Mockito.times(1)).getDataBlock(Mockito.any(TGetDataBlockRequest.class)); + Assert.assertEquals(0L, sourceHandle.getBufferRetainedSizeInBytes()); + sourceHandle.abort(); + } + @Test public void testFragmentAssemblyFailsWhenTotalLengthDoesNotMatch() { final String queryId = "q0"; diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/UserDataTransferErrorCode.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/UserDataTransferErrorCode.java index 2bf187999edad..18e9b5f933733 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/UserDataTransferErrorCode.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/audit/UserDataTransferErrorCode.java @@ -23,5 +23,6 @@ public enum UserDataTransferErrorCode { EMPTY_RESPONSE, UNEXPECTED_RESPONSE_SIZE, RECEIVER_CLOSED, - REMOTE_REJECTED + REMOTE_REJECTED, + OUT_OF_MEMORY }