This is an automated email from the ASF dual-hosted git repository. JackieTien97 pushed a commit to branch rc/2.0.11 in repository https://gitbox.apache.org/repos/asf/iotdb.git
commit 97031cb6903f4b91564b36e3ba901ce1ed50880b Author: Caideyipi <[email protected]> AuthorDate: Wed Jul 29 14:48:43 2026 +0800 [Pipe] Fix TsFile rate limit accounting (#18344) --- .../protocol/airgap/IoTDBDataRegionAirGapSink.java | 2 +- .../async/handler/PipeTransferTsFileHandler.java | 22 ++++-- .../thrift/sync/IoTDBDataRegionSyncSink.java | 12 ++-- .../airgap/IoTDBDataRegionAirGapSinkTest.java | 17 ++++- .../PipeTransferTsFileHandlerRateLimitTest.java | 83 ++++++++++++++++++++++ .../pipe/sink/protocol/IoTDBAirGapSink.java | 2 +- .../pipe/sink/protocol/IoTDBSslSyncSink.java | 2 +- 7 files changed, 126 insertions(+), 14 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java index 8689069334e..8eb941cad1d 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSink.java @@ -512,11 +512,11 @@ public class IoTDBDataRegionAirGapSink extends IoTDBDataNodeAirGapSink { final byte[] readBuffer = new byte[readFileBufferSize]; long position = 0; while (true) { - mayLimitRateAndRecordIO(readFileBufferSize); final int readLength = reader.read(readBuffer); if (readLength == -1) { break; } + mayLimitRateAndRecordIO(readLength); final byte[] payload = readLength == readFileBufferSize diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java index 4c25d03de53..16ffd0976ec 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandler.java @@ -176,11 +176,7 @@ public class PipeTransferTsFileHandler extends PipeTransferTrackableHandler { client.setShouldReturnSelf(false); client.setTimeoutDynamically(clientManager.getConnectionTimeout()); - PipeResourceMetrics.getInstance().recordDiskIO(readFileBufferSize); - if (sink.isEnableSendTsFileLimit()) { - TsFileSendRateLimiter.getInstance().acquire(readFileBufferSize); - } - final int readLength = reader.read(readBuffer); + final int readLength = readNextFilePiece(reader, readBuffer); if (readLength == -1) { if (currentFile == modFile) { @@ -250,6 +246,22 @@ public class PipeTransferTsFileHandler extends PipeTransferTrackableHandler { position += readLength; } + protected int readNextFilePiece(final RandomAccessFile reader, final byte[] readBuffer) + throws IOException { + final int readLength = reader.read(readBuffer); + if (readLength != -1) { + mayLimitRateAndRecordIO(readLength); + } + return readLength; + } + + protected void mayLimitRateAndRecordIO(final long requiredBytes) { + PipeResourceMetrics.getInstance().recordDiskIO(requiredBytes); + if (sink.isEnableSendTsFileLimit()) { + TsFileSendRateLimiter.getInstance().acquire(requiredBytes); + } + } + @Override public void onComplete(final TPipeTransferResp response) { try { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java index d13b9664f8a..6d205c36c65 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/sync/IoTDBDataRegionSyncSink.java @@ -630,7 +630,7 @@ public class IoTDBDataRegionSyncSink extends IoTDBDataNodeSyncSink { final byte[] readBuffer = new byte[readFileBufferSize]; long position = 0; int readLength; - while ((readLength = readNextFilePiece(reader, readBuffer, readFileBufferSize)) != -1) { + while ((readLength = readNextFilePiece(reader, readBuffer)) != -1) { position = transferFilePiece( pipe2WeightMap, @@ -645,11 +645,13 @@ public class IoTDBDataRegionSyncSink extends IoTDBDataNodeSyncSink { } } - private int readNextFilePiece( - final RandomAccessFile reader, final byte[] readBuffer, final int readFileBufferSize) + private int readNextFilePiece(final RandomAccessFile reader, final byte[] readBuffer) throws IOException { - mayLimitRateAndRecordIO(readFileBufferSize); - return reader.read(readBuffer); + final int readLength = reader.read(readBuffer); + if (readLength != -1) { + mayLimitRateAndRecordIO(readLength); + } + return readLength; } private long transferFilePiece( diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java index d49ef82ce05..1e6f3f865c1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBDataRegionAirGapSinkTest.java @@ -27,6 +27,7 @@ import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeRequestType import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent; import org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent; import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTabletBatchReqV2; +import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTsFilePieceReq; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; @@ -45,6 +46,7 @@ import java.util.Arrays; import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicLong; public class IoTDBDataRegionAirGapSinkTest { @@ -92,8 +94,14 @@ public class IoTDBDataRegionAirGapSinkTest { sink.transfer(new PipeHeartbeatEvent(-1, false)); final List<Short> requestTypes = new ArrayList<>(); + long transferredTsFileBytes = 0; for (final byte[] requestBytes : sink.sentRequests) { - requestTypes.add(toTPipeTransferReq(requestBytes).type); + final TPipeTransferReq req = toTPipeTransferReq(requestBytes); + requestTypes.add(req.type); + if (req.type == PipeRequestType.TRANSFER_TS_FILE_PIECE.getType()) { + transferredTsFileBytes += + PipeTransferTsFilePieceReq.fromTPipeTransferReq(req).getFilePiece().length; + } } Assert.assertTrue(requestTypes.contains(PipeRequestType.TRANSFER_TS_FILE_PIECE.getType())); @@ -101,6 +109,7 @@ public class IoTDBDataRegionAirGapSinkTest { requestTypes.contains(PipeRequestType.TRANSFER_TS_FILE_SEAL_WITH_MOD.getType())); Assert.assertFalse(requestTypes.contains(PipeRequestType.TRANSFER_TABLET_RAW_V2.getType())); Assert.assertFalse(requestTypes.contains(PipeRequestType.TRANSFER_TABLET_BATCH_V2.getType())); + Assert.assertEquals(transferredTsFileBytes, sink.rateLimitedBytes.get()); } } @@ -152,6 +161,7 @@ public class IoTDBDataRegionAirGapSinkTest { private static class RecordingIoTDBDataRegionAirGapSink extends IoTDBDataRegionAirGapSink { private final List<byte[]> sentRequests = new ArrayList<>(); + private final AtomicLong rateLimitedBytes = new AtomicLong(0); private void prepareSocket() { sockets.set(0, new TestingAirGapSocket()); @@ -168,6 +178,11 @@ public class IoTDBDataRegionAirGapSinkTest { return true; } + @Override + protected void mayLimitRateAndRecordIO(final long requiredBytes) { + rateLimitedBytes.addAndGet(requiredBytes); + } + private static class TestingAirGapSocket extends AirGapSocket { private TestingAirGapSocket() { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerRateLimitTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerRateLimitTest.java new file mode 100644 index 00000000000..55cfbcbd28d --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/sink/protocol/thrift/async/handler/PipeTransferTsFileHandlerRateLimitTest.java @@ -0,0 +1,83 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.iotdb.db.pipe.sink.protocol.thrift.async.handler; + +import org.apache.iotdb.db.pipe.sink.protocol.thrift.async.IoTDBDataRegionAsyncSink; + +import org.junit.Assert; +import org.junit.Test; +import org.mockito.Mockito; + +import java.io.File; +import java.io.RandomAccessFile; +import java.nio.file.Files; +import java.util.Collections; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; + +public class PipeTransferTsFileHandlerRateLimitTest { + + @Test + public void testRateLimitUsesActualReadLengthAndSkipsEndOfFile() throws Exception { + final File file = Files.createTempFile("pipe-transfer-rate-limit", ".tsfile").toFile(); + Files.write(file.toPath(), new byte[7]); + + final RecordingPipeTransferTsFileHandler handler = new RecordingPipeTransferTsFileHandler(file); + try (final RandomAccessFile reader = new RandomAccessFile(file, "r")) { + final byte[] readBuffer = new byte[4]; + + Assert.assertEquals(4, handler.readNextFilePiece(reader, readBuffer)); + Assert.assertEquals(3, handler.readNextFilePiece(reader, readBuffer)); + Assert.assertEquals(-1, handler.readNextFilePiece(reader, readBuffer)); + + Assert.assertEquals(file.length(), handler.rateLimitedBytes.get()); + Assert.assertEquals(2, handler.rateLimitInvocationCount.get()); + } finally { + handler.close(); + Assert.assertTrue(file.delete()); + } + } + + private static class RecordingPipeTransferTsFileHandler extends PipeTransferTsFileHandler { + + private final AtomicLong rateLimitedBytes = new AtomicLong(0); + private final AtomicInteger rateLimitInvocationCount = new AtomicInteger(0); + + private RecordingPipeTransferTsFileHandler(final File file) throws InterruptedException { + super( + Mockito.mock(IoTDBDataRegionAsyncSink.class), + Collections.emptyMap(), + Collections.emptyList(), + new AtomicInteger(1), + new AtomicBoolean(false), + file, + null, + false, + null); + } + + @Override + protected void mayLimitRateAndRecordIO(final long requiredBytes) { + rateLimitedBytes.addAndGet(requiredBytes); + rateLimitInvocationCount.incrementAndGet(); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java index 8bcca0315f4..6919b3b87ce 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBAirGapSink.java @@ -259,11 +259,11 @@ public abstract class IoTDBAirGapSink extends IoTDBSink { long position = 0; try (final RandomAccessFile reader = new RandomAccessFile(file, "r")) { while (true) { - mayLimitRateAndRecordIO(readFileBufferSize); final int readLength = reader.read(readBuffer); if (readLength == -1) { break; } + mayLimitRateAndRecordIO(readLength); final byte[] payload = readLength == readFileBufferSize diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java index a045082b24e..4595a0448ea 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/protocol/IoTDBSslSyncSink.java @@ -286,11 +286,11 @@ public abstract class IoTDBSslSyncSink extends IoTDBSink { long position = 0; try (final RandomAccessFile reader = new RandomAccessFile(file, "r")) { while (true) { - mayLimitRateAndRecordIO(readFileBufferSize); final int readLength = reader.read(readBuffer); if (readLength == -1) { break; } + mayLimitRateAndRecordIO(readLength); final byte[] payLoad = readLength == readFileBufferSize
