This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new f3f7d125fe0 [Pipe] Fix TsFile rate limit accounting (#18344)
f3f7d125fe0 is described below
commit f3f7d125fe01d0e22ceb2211b942669bd845ca65
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 a535f5bb33a..1da56f5703d 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
@@ -517,11 +517,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 7f06b771a0c..e46b25d8ff2 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
@@ -179,11 +179,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) {
@@ -253,6 +249,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 46d4034cf35..c1a1b9b7e6e 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
@@ -635,7 +635,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,
@@ -650,11 +650,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
