This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch dev/1.3
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/dev/1.3 by this push:
new 6c08c62aeb2 [Pipe] Fix TsFile rate limit accounting (#18344) (#18347)
6c08c62aeb2 is described below
commit 6c08c62aeb2dcc7f6e5dd36df01db495ebee7e1d
Author: Caideyipi <[email protected]>
AuthorDate: Thu Jul 30 10:51:25 2026 +0800
[Pipe] Fix TsFile rate limit accounting (#18344) (#18347)
---
.../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 8226eeac0f7..4b088b0aefd 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
@@ -450,11 +450,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 b6936485770..fc6dfefc24f 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
@@ -173,11 +173,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 7a1fe1463f6..49d52fac3a7 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
@@ -547,7 +547,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,
@@ -562,11 +562,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 ef947f13c28..db6169e9831 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.PipeTransferTabletBatchReq;
+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;
@@ -44,6 +45,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 {
@@ -91,14 +93,21 @@ 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()));
Assert.assertTrue(requestTypes.contains(PipeRequestType.TRANSFER_TS_FILE_SEAL.getType()));
Assert.assertFalse(requestTypes.contains(PipeRequestType.TRANSFER_TABLET_RAW.getType()));
Assert.assertFalse(requestTypes.contains(PipeRequestType.TRANSFER_TABLET_BATCH.getType()));
+ Assert.assertEquals(transferredTsFileBytes, sink.rateLimitedBytes.get());
}
}
@@ -140,6 +149,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());
@@ -156,6 +166,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 49caa1fc5ca..2dcb1af27f0 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
@@ -268,11 +268,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 96b7346d025..33276098ff6 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
@@ -180,11 +180,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
