From f03440874c1b404808b4f512795d64c092bea9f1 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 30 Oct 2026 15:28:38 +0800 Subject: [PATCH 1/5] Support strictly unidirectional UDP air-gap transfer --- .../protocol/airgap/IoTDBAirGapReceiver.java | 4 +- .../airgap/IoTDBAirGapReceiverTest.java | 56 +++++++ .../airgap/IoTDBDataRegionAirGapSinkTest.java | 137 ++++++++++++++++++ .../iotdb/commons/i18n/PipeMessages.java | 6 + .../iotdb/commons/i18n/PipeMessages.java | 6 + .../config/constant/PipeSinkConstant.java | 6 + .../payload/airgap/AirGapUdpProtocol.java | 29 ++++ .../pipe/sink/protocol/IoTDBAirGapSink.java | 64 +++++++- 8 files changed, 303 insertions(+), 5 deletions(-) create mode 100644 iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapUdpProtocol.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java index e3413b009b2ad..1a9cb1616d082 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java @@ -25,6 +25,7 @@ import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapELanguageConstant; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; +import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapUdpProtocol; import org.apache.iotdb.db.i18n.DataNodePipeMessages; import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.pipe.receiver.protocol.thrift.IoTDBDataNodeReceiverAgent; @@ -296,7 +297,8 @@ private void fail() throws IOException { private void respond(final byte[] response) throws IOException { if (currentDatagramSocket != null && currentRemoteSocketAddress != null) { - if (currentDatagramRequestId == null) { + if (currentDatagramRequestId == null + || (currentDatagramRequestId & AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK) != 0) { return; } final byte[] responseBytes = diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java index 57cacb3879906..86481ce3517b1 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java @@ -27,6 +27,7 @@ import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapELanguageConstant; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; +import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapUdpProtocol; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion; import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent; import org.apache.iotdb.db.pipe.receiver.protocol.thrift.IoTDBDataNodeReceiverAgent; @@ -202,6 +203,61 @@ public void testUdpResponseContainsRequestId() throws Exception { } } + @Test + public void testStrictUnidirectionalUdpNeverRespondsToSuccessOrFailure() throws Exception { + for (final TSStatusCode status : + new TSStatusCode[] { + TSStatusCode.SUCCESS_STATUS, TSStatusCode.PIPE_RECEIVER_USER_CONFLICT_EXCEPTION + }) { + try (final DatagramSocket receiverSocket = new NoResponseDatagramSocket()) { + final byte[] request = + toUdpDatagram(AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK | 73L, toRawRequestBytes()); + final DatagramPacket packet = + new DatagramPacket(request, request.length, InetAddress.getLoopbackAddress(), 12345); + final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(new RecordingSocket(), 5L); + final StubIoTDBDataNodeReceiverAgent stubAgent = new StubIoTDBDataNodeReceiverAgent(); + final AtomicInteger receivedRequests = new AtomicInteger(); + stubAgent.setStubReceiver( + "udp-client", + new StubReceiver(new TSStatus(status.getStatusCode())) { + @Override + public TPipeTransferResp receive(final TPipeTransferReq req) { + receivedRequests.incrementAndGet(); + return super.receive(req); + } + }); + setField(receiver, "agent", stubAgent); + receiver.receiveUdp(receiverSocket, packet, "udp-client", request, () -> {}); + Assert.assertEquals(1, receivedRequests.get()); + } + } + } + + @Test + public void testStrictUnidirectionalUdpDoesNotRespondToCorruptRequest() throws Exception { + try (final DatagramSocket receiverSocket = new NoResponseDatagramSocket()) { + final byte[] request = + toUdpDatagram(AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK | 74L, toRawRequestBytes()); + request[request.length - 1] ^= 1; + final DatagramPacket packet = + new DatagramPacket(request, request.length, InetAddress.getLoopbackAddress(), 12345); + final AtomicBoolean registered = new AtomicBoolean(); + new IoTDBAirGapReceiver(new RecordingSocket(), 6L) + .receiveUdp(receiverSocket, packet, "udp-client", request, () -> registered.set(true)); + Assert.assertFalse(registered.get()); + } + } + + private static class NoResponseDatagramSocket extends DatagramSocket { + + private NoResponseDatagramSocket() throws SocketException {} + + @Override + public void send(final DatagramPacket packet) { + Assert.fail("Strictly unidirectional requests must never send an ACK or NACK"); + } + } + @Test public void testTemporaryUnavailableRetryUsesFreshRequestBody() throws Exception { assertTemporaryUnavailableRetryUsesFreshRequestBody(null); 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 82000564cfef5..cdc0a300067af 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 @@ -24,6 +24,7 @@ import org.apache.iotdb.commons.pipe.config.plugin.configuraion.PipeTaskRuntimeConfiguration; import org.apache.iotdb.commons.pipe.config.plugin.env.PipeTaskSinkRuntimeEnvironment; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; +import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapUdpProtocol; 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; @@ -32,6 +33,7 @@ import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameterValidator; import org.apache.iotdb.pipe.api.customizer.parameter.PipeParameters; import org.apache.iotdb.pipe.api.exception.PipeConnectionException; +import org.apache.iotdb.pipe.api.exception.PipeParameterNotValidException; import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; import org.apache.tsfile.enums.TSDataType; @@ -60,6 +62,137 @@ public class IoTDBDataRegionAirGapSinkTest { + @Test + public void testStrictUnidirectionalRequiresUdpAndBooleanValue() throws Exception { + try (final IoTDBDataRegionAirGapSink sink = new IoTDBDataRegionAirGapSink()) { + for (final String key : + Arrays.asList( + PipeSinkConstant.CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, + PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY)) { + final PipeParameters parameters = buildParameters(false); + parameters.getAttribute().put(key, "true"); + Assert.assertThrows( + PipeParameterNotValidException.class, + () -> sink.validate(new PipeParameterValidator(parameters))); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + sink.validate(new PipeParameterValidator(parameters)); + parameters.getAttribute().put(key, "invalid"); + Assert.assertThrows( + PipeParameterNotValidException.class, + () -> sink.validate(new PipeParameterValidator(parameters))); + } + } + } + + @Test + public void testStrictUnidirectionalUdpSendsAllSlicesWithoutAcknowledgements() throws Exception { + for (final String key : + Arrays.asList( + PipeSinkConstant.CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, + PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY)) { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = buildParameters(false); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + parameters + .getAttribute() + .put(PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, "1024"); + parameters.getAttribute().put(key, "true"); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + + final byte[] request = new byte[2500]; + request[0] = 1; + request[2] = (byte) PipeRequestType.TRANSFER_TABLET_RAW.getType(); + // Nothing receives or acknowledges until sendUdp has returned. + Assert.assertTrue(sink.sendUdp(receiverSocket.getLocalPort(), request)); + + receiverSocket.setSoTimeout(1000); + for (int i = 0; i < 3; i++) { + final byte[] buffer = new byte[1025]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(packet); + Assert.assertTrue(packet.getLength() <= 1024); + final long requestId = ByteBuffer.wrap(buffer).getLong(); + Assert.assertEquals(AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK, requestId & Long.MIN_VALUE); + Assert.assertEquals(i + 1, requestId & Long.MAX_VALUE); + Assert.assertEquals( + PipeRequestType.TRANSFER_SLICE.getType(), + decodeAirGapDatagram(buffer, packet.getLength()).type); + } + } + } + } + + @Test + public void testStrictUnidirectionalHandshakeIsRefreshedBeforeEachEvent() throws Exception { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = buildParameters(false); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + parameters + .getAttribute() + .put(PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, "true"); + parameters + .getAttribute() + .put( + PipeSinkConstant.CONNECTOR_IOTDB_NODE_URLS_KEY, + "127.0.0.1:" + receiverSocket.getLocalPort()); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + sink.handshake(); + final byte[] request = new byte[] {1, 0, 11}; + sink.sendEvent(request); + sink.sendEvent(request); + + receiverSocket.setSoTimeout(1000); + for (final short type : + new short[] { + PipeRequestType.HANDSHAKE_DATANODE_V2.getType(), + PipeRequestType.HANDSHAKE_DATANODE_V2.getType(), + 11, + PipeRequestType.HANDSHAKE_DATANODE_V2.getType(), + 11 + }) { + final byte[] buffer = new byte[65_507]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(packet); + Assert.assertEquals(type, decodeAirGapDatagram(buffer, packet.getLength()).type); + Assert.assertTrue(ByteBuffer.wrap(buffer).getLong() < 0); + } + } + } + + @Test + public void testUdpDefaultAndExplicitFalseRequireAcknowledgement() throws Exception { + for (final String strictValue : Arrays.asList(null, "false")) { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = buildParameters(false); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + if (strictValue != null) { + parameters + .getAttribute() + .put(PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, strictValue); + } + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + Assert.assertThrows( + SocketTimeoutException.class, + () -> sink.sendUdp(receiverSocket.getLocalPort(), new byte[] {1, 0, 11})); + } + } + } + @Test public void testCompletionBarrierRetriesUnacknowledgedBatch() throws Exception { try (final RecordingIoTDBDataRegionAirGapSink sink = new RecordingIoTDBDataRegionAirGapSink()) { @@ -455,6 +588,10 @@ public synchronized void setSoTimeout(final int timeout) { private static class UdpTestingIoTDBDataRegionAirGapSink extends IoTDBDataRegionAirGapSink { + private void sendEvent(final byte[] request) throws IOException { + Assert.assertTrue(sendBytes(sockets.get(nextSocketIndex()), request)); + } + private boolean sendUdp(final int port, final byte[] request) throws Exception { final AirGapSocket socket = new AirGapSocket("127.0.0.1", port); try (final Socket ignored = socket) { diff --git a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java index 437e4dd68de55..7cdf30ca6e696 100644 --- a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java @@ -20,6 +20,12 @@ package org.apache.iotdb.commons.i18n; public final class PipeMessages { + public static final String EXCEPTION_STRICTLY_UNIDIRECTIONAL_AIR_GAP_TRANSPORT_REQUIRES_UDP_78921AF9 = + "Strictly unidirectional air-gap transport requires UDP."; + public static final String LOG_STRICTLY_UNIDIRECTIONAL_UDP_TRANSPORT_IS_ENABLED_REQUESTS_ARE_SENT_WITHOUT_ACKNOWLEDGEMENTS_DELIVERY_AND_RECEIVER_ERRORS_CANNOT_BE_CONFIRMED_0A2F4382 = + "Strictly unidirectional UDP transport is enabled. Requests are sent without acknowledgements; delivery and receiver errors cannot be confirmed."; + public static final String LOG_AIR_GAP_HANDSHAKE_REQUEST_SENT_ON_SOCKET_ARG_WITHOUT_WAITING_FOR_ACKNOWLEDGEMENT_E1ECD582 = + "Air gap handshake request sent on socket {} without waiting for acknowledgement."; // ===================== Common / Shared ===================== diff --git a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java index f67f0a351ac77..7b57a22fa3783 100644 --- a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java @@ -20,6 +20,12 @@ package org.apache.iotdb.commons.i18n; public final class PipeMessages { + public static final String EXCEPTION_STRICTLY_UNIDIRECTIONAL_AIR_GAP_TRANSPORT_REQUIRES_UDP_78921AF9 = + "严格单向的 air-gap 传输必须使用 UDP。"; + public static final String LOG_STRICTLY_UNIDIRECTIONAL_UDP_TRANSPORT_IS_ENABLED_REQUESTS_ARE_SENT_WITHOUT_ACKNOWLEDGEMENTS_DELIVERY_AND_RECEIVER_ERRORS_CANNOT_BE_CONFIRMED_0A2F4382 = + "已启用严格单向 UDP 传输。请求发送后不等待确认；无法确认送达情况和接收端错误。"; + public static final String LOG_AIR_GAP_HANDSHAKE_REQUEST_SENT_ON_SOCKET_ARG_WITHOUT_WAITING_FOR_ACKNOWLEDGEMENT_E1ECD582 = + "已通过 socket {} 发送 air-gap 握手请求，不等待确认。"; // ===================== Common / Shared ===================== diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java index e7a36577a0623..c04f6b87c79a4 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java @@ -198,6 +198,12 @@ private static String getDefaultConnectorOrSinkName(final PipeParameters paramet Arrays.asList( CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE, CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE))); + public static final String CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY = + "connector.air-gap.strict-unidirectional"; + public static final String SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY = + "sink.air-gap.strict-unidirectional"; + public static final boolean CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_DEFAULT_VALUE = false; + public static final String CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY = "connector.air-gap.udp.packet-size-bytes"; public static final String SINK_AIR_GAP_UDP_PACKET_SIZE_KEY = diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapUdpProtocol.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapUdpProtocol.java new file mode 100644 index 0000000000000..783c6af94ff07 --- /dev/null +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapUdpProtocol.java @@ -0,0 +1,29 @@ +/* + * 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.commons.pipe.sink.payload.airgap; + +public final class AirGapUdpProtocol { + + // Reserve the request ID's sign bit for requests that must not produce any reverse traffic. + // The remaining 63 bits identify the request; acknowledged requests keep their existing envelope. + public static final long NO_ACK_REQUEST_ID_MASK = Long.MIN_VALUE; + + private AirGapUdpProtocol() {} +} 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 bdffcfe473485..cd3c151d36f1e 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 @@ -26,6 +26,7 @@ import org.apache.iotdb.commons.pipe.resource.log.PipeLogger; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapELanguageConstant; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapOneByteResponse; +import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapUdpProtocol; import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferSliceReqBuilder; import org.apache.iotdb.commons.utils.ErrorHandlingCommonUtils; import org.apache.iotdb.pipe.api.annotation.TableModel; @@ -69,6 +70,8 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_HANDSHAKE_TIMEOUT_MS_DEFAULT_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_HANDSHAKE_TIMEOUT_MS_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_DEFAULT_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_SET; @@ -82,6 +85,7 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_ROUND_ROBIN_STRATEGY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_E_LANGUAGE_ENABLE_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_HANDSHAKE_TIMEOUT_MS_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_TRANSPORT_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_UDP_PACKET_SIZE_KEY; import static org.apache.iotdb.commons.utils.BasicStructureSerDeUtil.LONG_LEN; @@ -182,6 +186,7 @@ public String toString() { private boolean eLanguageEnable; private boolean useUdpTransport; + private boolean strictUnidirectional; private int udpPacketSizeInBytes; // The air gap connector does not use clientManager thus we put handshake type here @@ -209,6 +214,21 @@ public void validate(final PipeParameterValidator validator) throws Exception { CONNECTOR_AIR_GAP_TRANSPORT_SET, airGapTransport), airGapTransport); + validator + .validateAttributeValueRange( + CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, true, "true", "false") + .validateAttributeValueRange(SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, true, "true", "false") + .validate( + arg -> + !parameters.getBooleanOrDefault( + Arrays.asList( + CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, + SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY), + CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_DEFAULT_VALUE) + || CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE.equals(airGapTransport), + PipeMessages.EXCEPTION_STRICTLY_UNIDIRECTIONAL_AIR_GAP_TRANSPORT_REQUIRES_UDP_78921AF9, + airGapTransport); + final int packetSize = parameters.getIntOrDefault( Arrays.asList(CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, SINK_AIR_GAP_UDP_PACKET_SIZE_KEY), @@ -273,6 +293,17 @@ public void customize( CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE) .trim() .toLowerCase()); + strictUnidirectional = + parameters.getBooleanOrDefault( + Arrays.asList( + CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, + SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY), + CONNECTOR_AIR_GAP_STRICT_UNIDIRECTIONAL_DEFAULT_VALUE); + if (strictUnidirectional) { + LOGGER.warn( + PipeMessages + .LOG_STRICTLY_UNIDIRECTIONAL_UDP_TRANSPORT_IS_ENABLED_REQUESTS_ARE_SENT_WITHOUT_ACKNOWLEDGEMENTS_DELIVERY_AND_RECEIVER_ERRORS_CANNOT_BE_CONFIRMED_0A2F4382); + } udpPacketSizeInBytes = parameters.getIntOrDefault( Arrays.asList(CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, SINK_AIR_GAP_UDP_PACKET_SIZE_KEY), @@ -370,7 +401,14 @@ protected void sendHandshakeReq(final AirGapSocket socket) throws IOException { supportModsIfIsDataNodeReceiver = true; } socket.setSoTimeout(PIPE_CONFIG.getPipeSinkTransferTimeoutMs()); - LOGGER.info(PipeMessages.HANDSHAKE_SUCCESS_SOCKET, socket); + if (strictUnidirectional) { + LOGGER.debug( + PipeMessages + .LOG_AIR_GAP_HANDSHAKE_REQUEST_SENT_ON_SOCKET_ARG_WITHOUT_WAITING_FOR_ACKNOWLEDGEMENT_E1ECD582, + socket); + } else { + LOGGER.info(PipeMessages.HANDSHAKE_SUCCESS_SOCKET, socket); + } } protected abstract byte[] generateHandShakeV1Payload() throws IOException; @@ -444,8 +482,19 @@ protected abstract byte[] getTransferSingleFilePieceBytes( protected abstract byte[] getTransferMultiFilePieceBytes( final String fileName, final long position, final byte[] payLoad) throws IOException; - protected int nextSocketIndex() { - return loadBalancer.nextSocketIndex(); + protected int nextSocketIndex() throws IOException { + final int socketIndex = loadBalancer.nextSocketIndex(); + if (strictUnidirectional) { + // There is no feedback when a receiver restarts or its idle session expires. Re-establish + // the session before each event, never between the pieces/slices of the same file/request. + try { + sendHandshakeReq(sockets.get(socketIndex)); + } catch (final IOException e) { + isSocketAlive.set(socketIndex, false); + throw e; + } + } + return socketIndex; } protected boolean send( @@ -488,7 +537,9 @@ private boolean sendBytesByUdp(final AirGapSocket socket, final byte[] bytes) th private boolean sendOneDatagram(final AirGapSocket socket, final TPipeTransferReq request) throws IOException { - final long requestId = udpRequestIdGenerator.incrementAndGet(); + final long requestId = + (udpRequestIdGenerator.incrementAndGet() & Long.MAX_VALUE) + | (strictUnidirectional ? AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK : 0); final ByteBuffer body = request.body.duplicate(); final int requestSize = Byte.BYTES + Short.BYTES + body.remaining(); final int eLanguageSize = @@ -536,6 +587,11 @@ private boolean sendOneDatagram(final AirGapSocket socket, final TPipeTransferRe socket.getDatagramAddress(), socket.getEndPoint().getPort())); + if (strictUnidirectional) { + // This only confirms a local send. A one-way link cannot report delivery or remote errors. + return true; + } + final int originalTimeoutMs = datagramSocket.getSoTimeout(); final long deadlineNanos = originalTimeoutMs == 0 From 256dce46cb63300710ddb3c544c48f5bbabd2762 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 30 Oct 2026 17:17:37 +0800 Subject: [PATCH 2/5] Avoid refreshing idle air-gap heartbeat sessions --- .../airgap/IoTDBDataRegionAirGapSink.java | 7 +++- .../airgap/IoTDBSchemaRegionAirGapSink.java | 9 ++++- .../airgap/IoTDBDataRegionAirGapSinkTest.java | 38 +++++++++++++++++++ .../pipe/sink/protocol/IoTDBAirGapSink.java | 4 ++ 4 files changed, 56 insertions(+), 2 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 f6c96fe33bac1..673824139e14c 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 @@ -224,7 +224,12 @@ public void transfer(final Event event) throws Exception { return; } - final int socketIndex = nextSocketIndex(); + final boolean hasBatchToTransfer = isTabletBatchModeEnabled && !tabletBatchBuilder.isEmpty(); + final boolean shouldRefreshReceiverSession = + !(event instanceof PipeHeartbeatEvent || event instanceof PipeTerminateEvent) + || hasBatchToTransfer; + final int socketIndex = + shouldRefreshReceiverSession ? nextSocketIndex() : nextSocketIndex(false); final AirGapSocket socket = sockets.get(socketIndex); try { diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBSchemaRegionAirGapSink.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBSchemaRegionAirGapSink.java index 70dd5dfb7cf2c..547a32a29dcca 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBSchemaRegionAirGapSink.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/sink/protocol/airgap/IoTDBSchemaRegionAirGapSink.java @@ -83,7 +83,14 @@ public void transfer(final TsFileInsertionEvent tsFileInsertionEvent) throws Exc @Override public void transfer(final Event event) throws Exception { - final int socketIndex = nextSocketIndex(); + final boolean hasBatchToTransfer = + isTabletBatchModeEnabled + && Objects.nonNull(schemaRegionWritePlanEventBatch) + && !schemaRegionWritePlanEventBatch.isEmpty(); + final boolean shouldRefreshReceiverSession = + !(event instanceof PipeHeartbeatEvent) || hasBatchToTransfer; + final int socketIndex = + shouldRefreshReceiverSession ? nextSocketIndex() : nextSocketIndex(false); final AirGapSocket socket = sockets.get(socketIndex); try { 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 cdc0a300067af..c162a20d90b62 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 @@ -169,6 +169,44 @@ public void testStrictUnidirectionalHandshakeIsRefreshedBeforeEachEvent() throws } } + @Test + public void testStrictUnidirectionalDoesNotRefreshSessionForEmptyHeartbeat() throws Exception { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = buildParameters(false); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + parameters + .getAttribute() + .put(PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, "true"); + parameters + .getAttribute() + .put( + PipeSinkConstant.CONNECTOR_IOTDB_NODE_URLS_KEY, + "127.0.0.1:" + receiverSocket.getLocalPort()); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + sink.handshake(); + + receiverSocket.setSoTimeout(1000); + final byte[] buffer = new byte[65_507]; + final DatagramPacket handshakePacket = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(handshakePacket); + Assert.assertEquals( + PipeRequestType.HANDSHAKE_DATANODE_V2.getType(), + decodeAirGapDatagram(buffer, handshakePacket.getLength()).type); + + sink.transfer(new PipeHeartbeatEvent(-1, false)); + + receiverSocket.setSoTimeout(200); + Assert.assertThrows( + SocketTimeoutException.class, + () -> receiverSocket.receive(new DatagramPacket(new byte[65_507], 65_507))); + } + } + @Test public void testUdpDefaultAndExplicitFalseRequireAcknowledgement() throws Exception { for (final String strictValue : Arrays.asList(null, "false")) { 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 cd3c151d36f1e..a7d59713d5109 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 @@ -497,6 +497,10 @@ protected int nextSocketIndex() throws IOException { return socketIndex; } + protected int nextSocketIndex(final boolean shouldRefreshReceiverSession) throws IOException { + return shouldRefreshReceiverSession ? nextSocketIndex() : loadBalancer.nextSocketIndex(); + } + protected boolean send( final String pipeName, final long creationTime, final AirGapSocket socket, byte[] bytes) throws IOException { From 18846dcdeac89662be76f28f2cd97e8e73e3cfc9 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:08:48 +0800 Subject: [PATCH 3/5] Preserve UDP receiver state across repeated unidirectional handshakes --- .../protocol/airgap/IoTDBAirGapReceiver.java | 4 + .../thrift/IoTDBDataNodeReceiver.java | 39 ++++ .../airgap/IoTDBAirGapReceiverTest.java | 2 + ...TDBDataNodeReceiverUnidirectionalTest.java | 209 ++++++++++++++++++ .../AirGapPseudoTPipeTransferRequest.java | 15 +- 5 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java index 1a9cb1616d082..8c3bfe9d8b454 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java @@ -233,6 +233,9 @@ private AirGapPseudoTPipeTransferRequest toTPipeTransferReq(final byte[] data) { // Pseudo request, to reuse logic in IoTDBThriftReceiverAgent return (AirGapPseudoTPipeTransferRequest) new AirGapPseudoTPipeTransferRequest() + .setStrictUnidirectional( + currentDatagramRequestId != null + && (currentDatagramRequestId & AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK) != 0) .setVersion(ReadWriteIOUtils.readByte(byteBuffer)) .setType(ReadWriteIOUtils.readShort(byteBuffer)) .setBody(byteBuffer.slice()); @@ -282,6 +285,7 @@ private AirGapPseudoTPipeTransferRequest duplicateReq( final AirGapPseudoTPipeTransferRequest req) { return (AirGapPseudoTPipeTransferRequest) new AirGapPseudoTPipeTransferRequest() + .setStrictUnidirectional(req.isStrictUnidirectional()) .setVersion(req.getVersion()) .setType(req.getType()) .setBody(req.body.duplicate()); diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java index a7eef8a11f0b1..0c5d2e09f9e24 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java @@ -196,6 +196,9 @@ public class IoTDBDataNodeReceiver extends IoTDBFileReceiver { private final PipeTransferSliceReqHandler sliceReqHandler = new PipeTransferSliceReqHandler(); + private TPipeTransferReq lastSuccessfulUnidirectionalHandshake; + private TPipeTransferResp lastSuccessfulUnidirectionalHandshakeResp; + private final SqlParser tableSqlParser = new SqlParser(); private static final SessionManager SESSION_MANAGER = SessionManager.getInstance(); @@ -227,6 +230,40 @@ private enum TreeDatabaseCreationResult { @Override public synchronized TPipeTransferResp receive(final TPipeTransferReq req) { + final boolean isUnidirectionalHandshake = + req instanceof AirGapPseudoTPipeTransferRequest + && ((AirGapPseudoTPipeTransferRequest) req).isStrictUnidirectional() + && req.getType() == PipeRequestType.HANDSHAKE_DATANODE_V2.getType(); + // Snapshot before decoding, which consumes the request body. Repeated no-ACK handshakes + // refresh a live session without discarding its partially received file or request slices. + final TPipeTransferReq handshake = isUnidirectionalHandshake ? req.deepCopy() : null; + if (isUnidirectionalHandshake + && hasPipeHandshakeCredential + && handshake.equals(lastSuccessfulUnidirectionalHandshake)) { + final TSStatus status = login(); + if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + return lastSuccessfulUnidirectionalHandshakeResp.deepCopy(); + } + hasPipeHandshakeCredential = false; + lastSuccessfulUnidirectionalHandshake = null; + lastSuccessfulUnidirectionalHandshakeResp = null; + return new TPipeTransferResp(status); + } + if (req.getType() == PipeRequestType.HANDSHAKE_DATANODE_V1.getType() + || req.getType() == PipeRequestType.HANDSHAKE_DATANODE_V2.getType()) { + lastSuccessfulUnidirectionalHandshake = null; + lastSuccessfulUnidirectionalHandshakeResp = null; + } + final TPipeTransferResp resp = receiveInternal(req); + if (isUnidirectionalHandshake + && resp.getStatus().getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) { + lastSuccessfulUnidirectionalHandshake = handshake; + lastSuccessfulUnidirectionalHandshakeResp = resp.deepCopy(); + } + return resp; + } + + private TPipeTransferResp receiveInternal(final TPipeTransferReq req) { try { final long startTime = System.nanoTime(); final short rawRequestType = req.getType(); @@ -1938,6 +1975,8 @@ private TSStatus executeStatementForTableModelWithPermissionCheck( @Override public synchronized void handleExit() { + lastSuccessfulUnidirectionalHandshake = null; + lastSuccessfulUnidirectionalHandshakeResp = null; PipeReceiverRuntimeRegistry.getInstance() .deregister(configPipeReceiverRuntimeSessionKey.getAndSet(null)); clearSliceReqHandler(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java index 86481ce3517b1..f9933260fb984 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java @@ -222,6 +222,8 @@ public void testStrictUnidirectionalUdpNeverRespondsToSuccessOrFailure() throws new StubReceiver(new TSStatus(status.getStatusCode())) { @Override public TPipeTransferResp receive(final TPipeTransferReq req) { + Assert.assertTrue( + ((AirGapPseudoTPipeTransferRequest) req).isStrictUnidirectional()); receivedRequests.incrementAndGet(); return super.receive(req); } diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java new file mode 100644 index 0000000000000..dd2d9dce6898c --- /dev/null +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java @@ -0,0 +1,209 @@ +/* + * 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.receiver.protocol.thrift; + +import org.apache.iotdb.common.rpc.thrift.TSStatus; +import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV2Req; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferSliceReq; +import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferDataNodeHandshakeV2Req; +import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferTsFilePieceReq; +import org.apache.iotdb.rpc.RpcUtils; +import org.apache.iotdb.rpc.TSStatusCode; +import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; +import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp; + +import org.junit.Assert; +import org.junit.Test; + +import java.io.File; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.HashMap; +import java.util.Map; + +public class IoTDBDataNodeReceiverUnidirectionalTest { + + @Test + public void testRepeatedHandshakePreservesFileAndPendingSlices() throws Exception { + final Path baseDir = Files.createTempDirectory("unidirectional-receiver"); + final TestReceiver receiver = new TestReceiver(baseDir); + try { + assertSuccess(receiver.receive(handshake(true, true))); + assertSuccess( + receiver.receive( + PipeTransferTsFilePieceReq.toTPipeTransferReq( + "test.tsfile", 0, new byte[] {1, 2, 3}))); + final File receiverDir = receiver.getReceiverDir(); + final TPipeTransferReq piece = + PipeTransferTsFilePieceReq.toTPipeTransferReq("test.tsfile", 3, new byte[] {4, 5, 6}); + final int split = piece.body.remaining() / 2; + assertSuccess(receiver.receive(slice(piece, 0, 0, split))); + + assertSuccess(receiver.receive(handshake(true, true))); + + Assert.assertEquals(receiverDir, receiver.getReceiverDir()); + Assert.assertTrue(receiverDir.exists()); + assertSuccess(receiver.receive(slice(piece, 1, split, piece.body.limit()))); + Assert.assertArrayEquals( + new byte[] {1, 2, 3, 4, 5, 6}, + Files.readAllBytes(receiverDir.toPath().resolve("test.tsfile"))); + Assert.assertEquals(1, receiver.handshakeCount); + } finally { + receiver.handleExit(); + Files.deleteIfExists(baseDir); + } + } + + @Test + public void testAcknowledgedAndChangedHandshakesStillResetSession() throws Exception { + final Path baseDir = Files.createTempDirectory("unidirectional-receiver"); + final TestReceiver receiver = new TestReceiver(baseDir); + try { + assertSuccess(receiver.receive(handshake(true, true))); + File oldDir = receiver.getReceiverDir(); + + assertSuccess(receiver.receive(handshake(true, false))); + Assert.assertFalse(oldDir.exists()); + oldDir = receiver.getReceiverDir(); + + assertSuccess(receiver.receive(handshake(false, false))); + Assert.assertFalse(oldDir.exists()); + oldDir = receiver.getReceiverDir(); + + assertSuccess(receiver.receive(handshake(true, false))); + Assert.assertFalse(oldDir.exists()); + Assert.assertEquals(4, receiver.handshakeCount); + } finally { + receiver.handleExit(); + Files.deleteIfExists(baseDir); + } + } + + @Test + public void testFailedHandshakeIsRetriedAndAuthenticationIsRechecked() throws Exception { + final Path baseDir = Files.createTempDirectory("unidirectional-receiver"); + final TestReceiver receiver = new TestReceiver(baseDir); + try { + receiver.loginStatus = TSStatusCode.NOT_LOGIN; + Assert.assertEquals( + TSStatusCode.NOT_LOGIN.getStatusCode(), + receiver.receive(handshake(true, true)).getStatus().getCode()); + receiver.loginStatus = TSStatusCode.SUCCESS_STATUS; + assertSuccess(receiver.receive(handshake(true, true))); + Assert.assertEquals(2, receiver.handshakeCount); + + receiver.loginStatus = TSStatusCode.NOT_LOGIN; + Assert.assertEquals( + TSStatusCode.NOT_LOGIN.getStatusCode(), + receiver.receive(handshake(true, true)).getStatus().getCode()); + receiver.loginStatus = TSStatusCode.SUCCESS_STATUS; + final File oldDir = receiver.getReceiverDir(); + assertSuccess(receiver.receive(handshake(true, true))); + Assert.assertFalse(oldDir.exists()); + Assert.assertEquals(3, receiver.handshakeCount); + } finally { + receiver.handleExit(); + Files.deleteIfExists(baseDir); + } + } + + private static TPipeTransferReq handshake( + final boolean unidirectional, final boolean validateTsFile) throws IOException { + final Map params = new HashMap<>(); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_CLUSTER_ID, "sender-cluster"); + params.put( + PipeTransferHandshakeConstant.HANDSHAKE_KEY_TIME_PRECISION, + CommonDescriptor.getInstance().getConfig().getTimestampPrecision()); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_USERNAME, "root"); + params.put(PipeTransferHandshakeConstant.HANDSHAKE_KEY_PASSWORD, "root"); + params.put( + PipeTransferHandshakeConstant.HANDSHAKE_KEY_VALIDATE_TSFILE, + Boolean.toString(validateTsFile)); + final TPipeTransferReq req = PipeTransferDataNodeHandshakeV2Req.toTPipeTransferReq(params); + return new AirGapPseudoTPipeTransferRequest() + .setStrictUnidirectional(unidirectional) + .setVersion(req.version) + .setType(req.type) + .setBody(req.body); + } + + private static TPipeTransferReq slice( + final TPipeTransferReq req, final int index, final int start, final int end) + throws IOException { + final ByteBuffer body = req.body.duplicate(); + return PipeTransferSliceReq.toTPipeTransferReq(1, req.type, index, 2, body, start, end); + } + + private static void assertSuccess(final TPipeTransferResp resp) { + Assert.assertEquals(TSStatusCode.SUCCESS_STATUS.getStatusCode(), resp.getStatus().getCode()); + } + + private static class TestReceiver extends IoTDBDataNodeReceiver { + + private final Path baseDir; + private int handshakeCount; + private TSStatusCode loginStatus = TSStatusCode.SUCCESS_STATUS; + + private TestReceiver(final Path baseDir) { + this.baseDir = baseDir; + } + + @Override + protected TPipeTransferResp handleTransferHandshakeV2(final PipeTransferHandshakeV2Req req) + throws IOException { + handshakeCount++; + return super.handleTransferHandshakeV2(req); + } + + @Override + protected String getClusterId() { + return "receiver-cluster"; + } + + @Override + protected String getReceiverFileBaseDir() { + return baseDir.toString(); + } + + @Override + protected String getSenderHost() { + return "127.0.0.1"; + } + + @Override + protected String getSenderPort() { + return "9999"; + } + + @Override + protected TSStatus login() { + return RpcUtils.getStatus(loginStatus); + } + + private File getReceiverDir() { + return receiverFileDirWithIdSuffix.get(); + } + } +} diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapPseudoTPipeTransferRequest.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapPseudoTPipeTransferRequest.java index ed350d5cf0b65..f5758bb620069 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapPseudoTPipeTransferRequest.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/sink/payload/airgap/AirGapPseudoTPipeTransferRequest.java @@ -21,4 +21,17 @@ import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq; -public class AirGapPseudoTPipeTransferRequest extends TPipeTransferReq {} +public class AirGapPseudoTPipeTransferRequest extends TPipeTransferReq { + + private boolean strictUnidirectional; + + public boolean isStrictUnidirectional() { + return strictUnidirectional; + } + + public AirGapPseudoTPipeTransferRequest setStrictUnidirectional( + final boolean strictUnidirectional) { + this.strictUnidirectional = strictUnidirectional; + return this; + } +} From 343708dba8b5fc52cf79fa4d8cbc0511f01e5467 Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 2 Nov 2026 10:23:41 +0800 Subject: [PATCH 4/5] Pace one-way UDP packets and preserve nested receiver context --- .../airgap/IoTDBAirGapReceiverAgent.java | 130 +++++++++++++----- .../thrift/IoTDBDataNodeReceiver.java | 34 ++++- .../airgap/IoTDBAirGapReceiverAgentTest.java | 93 +++++++++++++ ...TDBDataNodeReceiverUnidirectionalTest.java | 37 ++++- .../airgap/IoTDBDataRegionAirGapSinkTest.java | 42 ++++++ .../iotdb/commons/i18n/PipeMessages.java | 2 + .../iotdb/commons/i18n/PipeMessages.java | 2 + .../iotdb/commons/conf/CommonConfig.java | 13 ++ .../iotdb/commons/pipe/config/PipeConfig.java | 4 + .../commons/pipe/config/PipeDescriptor.java | 5 + .../config/constant/PipeSinkConstant.java | 5 + .../pipe/sink/protocol/IoTDBAirGapSink.java | 47 +++++++ .../pipe/config/PipeDescriptorTest.java | 16 +++ 13 files changed, 392 insertions(+), 38 deletions(-) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java index 87683081d9be5..da056fcc2c58b 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java @@ -40,10 +40,14 @@ import java.net.ServerSocket; import java.net.Socket; import java.net.SocketTimeoutException; +import java.util.Arrays; import java.util.Map; import java.util.Objects; +import java.util.concurrent.BlockingQueue; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicLong; @@ -54,12 +58,17 @@ public class IoTDBAirGapReceiverAgent implements IService { private static final int UDP_PACKET_MAX_SIZE_IN_BYTES = 65_507; private static final int UDP_SESSION_CLEANUP_INTERVAL_MS = 60_000; private static final int UDP_MAX_CLIENT_SESSIONS = 1_024; - private static final int LISTENER_THREAD_COUNT = 2; + private static final int LISTENER_THREAD_COUNT = 3; + private static final int UDP_REQUEST_QUEUE_CAPACITY = 4_096; + private static final int UDP_REQUEST_QUEUE_BYTES = 32 * 1024 * 1024; private final ExecutorService listenExecutor = IoTDBThreadPoolFactory.newFixedThreadPool( LISTENER_THREAD_COUNT, ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName()); private final AtomicBoolean allowSubmitListen = new AtomicBoolean(false); + private final BlockingQueue udpRequests = + new LinkedBlockingQueue<>(UDP_REQUEST_QUEUE_CAPACITY); + private final Semaphore udpRequestBytes = new Semaphore(UDP_REQUEST_QUEUE_BYTES); private ServerSocket serverSocket; private DatagramSocket datagramSocket; @@ -88,38 +97,95 @@ public void listen() { } public void listenUdp() { - while (allowSubmitListen.get() && !datagramSocket.isClosed()) { - removeExpiredUdpClientSessions(System.nanoTime()); - final byte[] buffer = new byte[UDP_PACKET_MAX_SIZE_IN_BYTES]; - final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); - try { - datagramSocket.receive(packet); - - final long airGapReceiverId = receiverId.incrementAndGet(); - final IoTDBAirGapReceiver receiver = - new IoTDBAirGapReceiver(new Socket(), airGapReceiverId); - final String receiverKey = packet.getSocketAddress().toString(); - final boolean[] registeredSession = {false}; + final byte[] buffer = new byte[UDP_PACKET_MAX_SIZE_IN_BYTES]; + try { + while (allowSubmitListen.get() && !datagramSocket.isClosed()) { + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); try { - receiver.receiveUdp( - datagramSocket, - packet, - receiverKey, - buffer, - () -> registeredSession[0] = registerUdpSessionIfNecessary(packet)); - } finally { - if (registeredSession[0]) { - SessionManager.getInstance().removeCurrSession(); + datagramSocket.receive(packet); + // Drain the socket independently of database/authentication work. Keep one ordered + // consumer, and bound both queued packet count and retained payload memory. + udpRequestBytes.acquire(packet.getLength()); + boolean queued = false; + try { + udpRequests.put( + new DatagramPacket( + Arrays.copyOf(buffer, packet.getLength()), + packet.getLength(), + packet.getSocketAddress())); + queued = true; + } finally { + if (!queued) { + udpRequestBytes.release(packet.getLength()); + } } - } - } catch (final SocketTimeoutException e) { - removeExpiredUdpClientSessions(System.nanoTime()); - } catch (final IOException e) { - if (allowSubmitListen.get()) { + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } catch (final SocketTimeoutException e) { + // Session expiration is serialized with request handling by the consumer. + } catch (final IOException e) { + if (allowSubmitListen.get()) { + LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + } + } catch (final Exception e) { LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); } - } catch (final Exception e) { - LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + } + } finally { + clearUdpRequests(); + } + } + + void processUdpRequests() { + try { + while (allowSubmitListen.get()) { + final DatagramPacket packet = udpRequests.poll(1, TimeUnit.SECONDS); + if (packet == null) { + removeExpiredUdpClientSessions(System.nanoTime()); + continue; + } + try { + handleUdpPacket(packet); + } catch (final Exception e) { + if (allowSubmitListen.get()) { + LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER, e); + } + } finally { + udpRequestBytes.release(packet.getLength()); + } + } + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + } finally { + clearUdpRequests(); + udpClientSessions.forEach(this::removeUdpClientSession); + } + } + + private void clearUdpRequests() { + DatagramPacket packet; + while ((packet = udpRequests.poll()) != null) { + udpRequestBytes.release(packet.getLength()); + } + } + + void handleUdpPacket(final DatagramPacket packet) throws IOException { + removeExpiredUdpClientSessions(System.nanoTime()); + final IoTDBAirGapReceiver receiver = + new IoTDBAirGapReceiver(new Socket(), receiverId.incrementAndGet()); + final String receiverKey = packet.getSocketAddress().toString(); + final boolean[] registeredSession = {false}; + try { + receiver.receiveUdp( + datagramSocket, + packet, + receiverKey, + packet.getData(), + () -> registeredSession[0] = registerUdpSessionIfNecessary(packet)); + } finally { + if (registeredSession[0]) { + SessionManager.getInstance().removeCurrSession(); } } } @@ -196,6 +262,8 @@ public void start() throws StartupException { try { serverSocket = new ServerSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort()); datagramSocket = new DatagramSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort()); + datagramSocket.setReceiveBufferSize( + PipeConfig.getInstance().getPipeAirGapReceiverUdpSocketBufferSizeInBytes()); datagramSocket.setSoTimeout(UDP_SESSION_CLEANUP_INTERVAL_MS); } catch (final IOException e) { if (Objects.nonNull(serverSocket)) { @@ -211,6 +279,7 @@ public void start() throws StartupException { allowSubmitListen.set(true); listenExecutor.submit(this::listen); listenExecutor.submit(this::listenUdp); + listenExecutor.submit(this::processUdpRequests); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STARTED, serverSocket); } @@ -230,8 +299,7 @@ public void stop() { LOGGER.warn(DataNodePipeMessages.FAILED_TO_CLOSE_IOTDBAIRGAPRECEIVERAGENT_S_SERVER_SOCKET, e); } - udpClientSessions.forEach(this::removeUdpClientSession); - listenExecutor.shutdown(); + listenExecutor.shutdownNow(); LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STOPPED, serverSocket); } diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java index 0c5d2e09f9e24..e8470beca89dd 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java @@ -269,7 +269,9 @@ private TPipeTransferResp receiveInternal(final TPipeTransferReq req) { final short rawRequestType = req.getType(); if (PipeRequestType.isValidatedRequestType(rawRequestType)) { final PipeRequestType requestType = PipeRequestType.valueOf(rawRequestType); - if (requestType != PipeRequestType.TRANSFER_SLICE) { + if (requestType != PipeRequestType.TRANSFER_SLICE + && !(isStrictUnidirectionalRequest(req) + && requestType == PipeRequestType.TRANSFER_COMPRESSED)) { clearSliceReqHandler(); } final TPipeTransferResp authResp = checkPipeTransferAuthenticated(requestType); @@ -501,7 +503,9 @@ private TPipeTransferResp receiveInternal(final TPipeTransferReq req) { case TRANSFER_SLICE: { try { - return handleTransferSlice(PipeTransferSliceReq.fromTPipeTransferReq(req)); + return handleTransferSlice( + PipeTransferSliceReq.fromTPipeTransferReq(req), + isStrictUnidirectionalRequest(req)); } finally { PipeDataNodeReceiverMetrics.getInstance() .recordTransferSliceTimer(System.nanoTime() - startTime); @@ -515,7 +519,9 @@ private TPipeTransferResp receiveInternal(final TPipeTransferReq req) { PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(req); try (final PipeMemoryBlock ignored = tryAllocateReceiverMemory(requestedMemorySizeInBytes)) { - return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req)); + return receiveNestedRequest( + PipeTransferCompressedReq.fromTPipeTransferReq(req), + isStrictUnidirectionalRequest(req)); } } catch (final PipeRuntimeOutOfMemoryCriticalException e) { return new TPipeTransferResp( @@ -1135,7 +1141,25 @@ private String getConfigReceiverId() { return configReceiverId.get(); } - private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq pipeTransferSliceReq) { + private boolean isStrictUnidirectionalRequest(final TPipeTransferReq req) { + return req instanceof AirGapPseudoTPipeTransferRequest + && ((AirGapPseudoTPipeTransferRequest) req).isStrictUnidirectional(); + } + + private TPipeTransferResp receiveNestedRequest( + final TPipeTransferReq req, final boolean strictUnidirectional) { + return receive( + strictUnidirectional + ? new AirGapPseudoTPipeTransferRequest() + .setStrictUnidirectional(true) + .setVersion(req.version) + .setType(req.type) + .setBody(req.body) + : req); + } + + private TPipeTransferResp handleTransferSlice( + final PipeTransferSliceReq pipeTransferSliceReq, final boolean strictUnidirectional) { final long sliceBodySizeInBytes = getSliceBodySizeInBytes(pipeTransferSliceReq); long requestedMemorySizeInBytes = sliceBodySizeInBytes; String memoryAction = "buffering sliced pipe transfer request"; @@ -1174,7 +1198,7 @@ private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq pipeTra "Slice received, waiting for more slices to complete the request.")); } clearSliceReqHandler(); - return receive(req.get()); + return receiveNestedRequest(req.get(), strictUnidirectional); } } catch (final PipeRuntimeOutOfMemoryCriticalException e) { closeMemoryBlock(sliceMemoryBlock); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java index ffbebb5a66851..0335e208fc3af 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java @@ -39,12 +39,105 @@ import java.net.DatagramSocket; import java.net.InetAddress; import java.net.Socket; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; import java.util.Map; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Semaphore; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; public class IoTDBAirGapReceiverAgentTest { + @Test(timeout = 30_000) + public void testUdpSocketIsDrainedWhileRequestHandlingIsBlocked() throws Exception { + verifyUdpQueue(false); + } + + @Test(timeout = 30_000) + public void testStoppingUdpReceiverReleasesQueuedPackets() throws Exception { + verifyUdpQueue(true); + } + + private void verifyUdpQueue(final boolean stopWithPendingPackets) throws Exception { + final CountDownLatch handlingStarted = new CountDownLatch(1); + final CountDownLatch resumeHandling = new CountDownLatch(1); + final CountDownLatch handled = new CountDownLatch(17); + final List order = Collections.synchronizedList(new ArrayList<>()); + final IoTDBAirGapReceiverAgent agent = + new IoTDBAirGapReceiverAgent() { + @Override + void handleUdpPacket(final DatagramPacket packet) { + handlingStarted.countDown(); + try { + resumeHandling.await(); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + return; + } + order.add((int) packet.getData()[0]); + handled.countDown(); + } + }; + try (final DatagramSocket receiver = new DatagramSocket(0, InetAddress.getLoopbackAddress()); + final DatagramSocket sender = new DatagramSocket()) { + setAgentField(agent, "datagramSocket", receiver); + ((AtomicBoolean) getAgentField(agent, "allowSubmitListen")).set(true); + final ExecutorService executor = (ExecutorService) getAgentField(agent, "listenExecutor"); + @SuppressWarnings("unchecked") + final BlockingQueue queue = + (BlockingQueue) getAgentField(agent, "udpRequests"); + final Semaphore bytes = (Semaphore) getAgentField(agent, "udpRequestBytes"); + final int initialPermits = bytes.availablePermits(); + executor.submit(agent::listenUdp); + executor.submit(agent::processUdpRequests); + try { + final byte[] data = new byte[60 * 1024]; + sender.send(new DatagramPacket(data, data.length, receiver.getLocalSocketAddress())); + Assert.assertTrue(handlingStarted.await(5, TimeUnit.SECONDS)); + for (int i = 1; i <= 16; i++) { + data[0] = (byte) i; + sender.send(new DatagramPacket(data, data.length, receiver.getLocalSocketAddress())); + final long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (queue.size() < i && System.nanoTime() < deadline) { + Thread.sleep(1); + } + Assert.assertEquals(i, queue.size()); + } + Assert.assertEquals(initialPermits - 17 * data.length, bytes.availablePermits()); + if (!stopWithPendingPackets) { + resumeHandling.countDown(); + Assert.assertTrue(handled.await(5, TimeUnit.SECONDS)); + for (int i = 0; i < 17; i++) { + Assert.assertEquals(i, order.get(i).intValue()); + } + } + } finally { + agent.stop(); + resumeHandling.countDown(); + Assert.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + Assert.assertTrue(queue.isEmpty()); + Assert.assertEquals(initialPermits, bytes.availablePermits()); + } + } + + private static Object getAgentField(final Object target, final String name) throws Exception { + final Field field = IoTDBAirGapReceiverAgent.class.getDeclaredField(name); + field.setAccessible(true); + return field.get(target); + } + + private static void setAgentField(final Object target, final String name, final Object value) + throws Exception { + final Field field = IoTDBAirGapReceiverAgent.class.getDeclaredField(name); + field.setAccessible(true); + field.set(target, value); + } + @Test public void testExpiredUdpClientSessionIsRemoved() throws Exception { final IoTDBAirGapReceiverAgent agent = new IoTDBAirGapReceiverAgent(); diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java index dd2d9dce6898c..a450a311f57aa 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverUnidirectionalTest.java @@ -21,8 +21,11 @@ import org.apache.iotdb.common.rpc.thrift.TSStatus; import org.apache.iotdb.commons.conf.CommonDescriptor; +import org.apache.iotdb.commons.pipe.sink.compressor.PipeCompressor; +import org.apache.iotdb.commons.pipe.sink.compressor.PipeCompressorFactory; import org.apache.iotdb.commons.pipe.sink.payload.airgap.AirGapPseudoTPipeTransferRequest; import org.apache.iotdb.commons.pipe.sink.payload.thrift.common.PipeTransferHandshakeConstant; +import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferCompressedReq; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferHandshakeV2Req; import org.apache.iotdb.commons.pipe.sink.payload.thrift.request.PipeTransferSliceReq; import org.apache.iotdb.db.pipe.sink.payload.evolvable.request.PipeTransferDataNodeHandshakeV2Req; @@ -40,6 +43,7 @@ import java.nio.ByteBuffer; import java.nio.file.Files; import java.nio.file.Path; +import java.util.Collections; import java.util.HashMap; import java.util.Map; @@ -47,10 +51,20 @@ public class IoTDBDataNodeReceiverUnidirectionalTest { @Test public void testRepeatedHandshakePreservesFileAndPendingSlices() throws Exception { + verifyRepeatedHandshakePreservesFileAndPendingSlices(false); + } + + @Test + public void testCompressedHandshakePreservesFileAndPendingSlices() throws Exception { + verifyRepeatedHandshakePreservesFileAndPendingSlices(true); + } + + private void verifyRepeatedHandshakePreservesFileAndPendingSlices(final boolean compressed) + throws Exception { final Path baseDir = Files.createTempDirectory("unidirectional-receiver"); final TestReceiver receiver = new TestReceiver(baseDir); try { - assertSuccess(receiver.receive(handshake(true, true))); + assertSuccess(receiver.receive(handshakeForTransfer(compressed))); assertSuccess( receiver.receive( PipeTransferTsFilePieceReq.toTPipeTransferReq( @@ -61,7 +75,7 @@ public void testRepeatedHandshakePreservesFileAndPendingSlices() throws Exceptio final int split = piece.body.remaining() / 2; assertSuccess(receiver.receive(slice(piece, 0, 0, split))); - assertSuccess(receiver.receive(handshake(true, true))); + assertSuccess(receiver.receive(handshakeForTransfer(compressed))); Assert.assertEquals(receiverDir, receiver.getReceiverDir()); Assert.assertTrue(receiverDir.exists()); @@ -149,6 +163,25 @@ private static TPipeTransferReq handshake( .setBody(req.body); } + private static TPipeTransferReq handshakeForTransfer(final boolean compressed) + throws IOException { + final TPipeTransferReq handshake = handshake(true, true); + if (!compressed) { + return handshake; + } + final TPipeTransferReq req = + PipeTransferCompressedReq.toTPipeTransferReq( + handshake, + Collections.singletonList( + PipeCompressorFactory.getCompressor( + PipeCompressor.PipeCompressionType.GZIP.getIndex()))); + return new AirGapPseudoTPipeTransferRequest() + .setStrictUnidirectional(true) + .setVersion(req.version) + .setType(req.type) + .setBody(req.body); + } + private static TPipeTransferReq slice( final TPipeTransferReq req, final int index, final int start, final int end) throws IOException { 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 c162a20d90b62..e5b016942a126 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 @@ -61,6 +61,48 @@ import java.util.zip.CRC32; public class IoTDBDataRegionAirGapSinkTest { + @Test + public void testStrictUdpSendIntervalIsValidatedAndAppliedToEverySlice() throws Exception { + for (final String key : + Arrays.asList( + PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY, + PipeSinkConstant.SINK_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY)) { + try (final DatagramSocket receiverSocket = new DatagramSocket(0); + final UdpTestingIoTDBDataRegionAirGapSink sink = + new UdpTestingIoTDBDataRegionAirGapSink()) { + final PipeParameters parameters = buildParameters(false); + parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp"); + parameters + .getAttribute() + .put(PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY, "true"); + parameters + .getAttribute() + .put(PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, "1024"); + parameters.getAttribute().put(key, "-1"); + Assert.assertThrows( + PipeParameterNotValidException.class, + () -> sink.validate(new PipeParameterValidator(parameters))); + parameters.getAttribute().put(key, "20"); + sink.validate(new PipeParameterValidator(parameters)); + sink.customize( + parameters, + new PipeTaskRuntimeConfiguration(new PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1))); + final byte[] request = new byte[2500]; + request[0] = 1; + request[2] = (byte) 13; + final long start = System.nanoTime(); + Assert.assertTrue(sink.sendUdp(receiverSocket.getLocalPort(), request)); + Assert.assertTrue( + System.nanoTime() - start >= java.util.concurrent.TimeUnit.MILLISECONDS.toNanos(40)); + receiverSocket.setSoTimeout(1000); + for (int i = 0; i < 3; i++) { + final byte[] buffer = new byte[1024]; + final DatagramPacket packet = new DatagramPacket(buffer, buffer.length); + receiverSocket.receive(packet); + } + } + } + } @Test public void testStrictUnidirectionalRequiresUdpAndBooleanValue() throws Exception { diff --git a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java index 7cdf30ca6e696..fea25d97e676b 100644 --- a/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/en/org/apache/iotdb/commons/i18n/PipeMessages.java @@ -982,4 +982,6 @@ private PipeMessages() {} public static final String MESSAGE_DATAPARTITIONTABLE_GENERATION_FAILED_D85CD23A = "DataPartitionTable generation failed: "; public static final String MESSAGE_UNKNOWN_TASK_STATUS_E05D98F0 = "Unknown task status: "; + public static final String EXCEPTION_UDP_SEND_INTERVAL_MUST_BE_NON_NEGATIVE_BUT_GOT_ARG_36D41687 = + "UDP send interval must be non-negative, but got %d."; } diff --git a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java index 7b57a22fa3783..758eaae4d1c8e 100644 --- a/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java +++ b/iotdb-core/node-commons/src/main/i18n/zh/org/apache/iotdb/commons/i18n/PipeMessages.java @@ -948,4 +948,6 @@ private PipeMessages() {} public static final String MESSAGE_DATAPARTITIONTABLE_GENERATION_FAILED_D85CD23A = "DataPartitionTable 生成失败："; public static final String MESSAGE_UNKNOWN_TASK_STATUS_E05D98F0 = "未知任务状态："; + public static final String EXCEPTION_UDP_SEND_INTERVAL_MUST_BE_NON_NEGATIVE_BUT_GOT_ARG_36D41687 = + "UDP 发送间隔必须为非负数，实际为 %d。"; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java index d307e8bfa365f..e6a2acb8c96b9 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/conf/CommonConfig.java @@ -345,6 +345,7 @@ public class CommonConfig { Math.min( 64 * 1024 * 1024, (int) Math.min(Runtime.getRuntime().maxMemory() / 64, Integer.MAX_VALUE)); + private int pipeAirGapReceiverUdpSocketBufferSizeInBytes = 32 * 1024 * 1024; private boolean pipeReceiverLoadConversionEnabled = false; private volatile long loggerPeriodicalLogMinIntervalSeconds = 60; private volatile long loggerCacheMaxSizeInBytes = 64 * MB; @@ -2027,6 +2028,18 @@ public int getPipeAirGapReceiverMaxPayloadSizeInBytes() { return pipeAirGapReceiverMaxPayloadSizeInBytes; } + public int getPipeAirGapReceiverUdpSocketBufferSizeInBytes() { + return pipeAirGapReceiverUdpSocketBufferSizeInBytes; + } + + public void setPipeAirGapReceiverUdpSocketBufferSizeInBytes( + int pipeAirGapReceiverUdpSocketBufferSizeInBytes) { + if (pipeAirGapReceiverUdpSocketBufferSizeInBytes > 0) { + this.pipeAirGapReceiverUdpSocketBufferSizeInBytes = + pipeAirGapReceiverUdpSocketBufferSizeInBytes; + } + } + public double getPipeMetaReportMaxLogNumPerRound() { return pipeMetaReportMaxLogNumPerRound; } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java index e04237129e867..a48abdeea16ff 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeConfig.java @@ -379,6 +379,10 @@ public int getPipeAirGapReceiverMaxPayloadSizeInBytes() { return COMMON_CONFIG.getPipeAirGapReceiverMaxPayloadSizeInBytes(); } + public int getPipeAirGapReceiverUdpSocketBufferSizeInBytes() { + return COMMON_CONFIG.getPipeAirGapReceiverUdpSocketBufferSizeInBytes(); + } + public boolean isPipeReceiverLoadConversionEnabled() { return COMMON_CONFIG.isPipeReceiverLoadConversionEnabled(); } diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeDescriptor.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeDescriptor.java index 3a1da5f18ba46..b039b102d5120 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeDescriptor.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/PipeDescriptor.java @@ -487,6 +487,11 @@ public static void loadPipeInternalConfig(CommonConfig config, TrimProperties pr properties.getProperty( "pipe_air_gap_receiver_max_payload_size_in_bytes", String.valueOf(config.getPipeAirGapReceiverMaxPayloadSizeInBytes())))); + config.setPipeAirGapReceiverUdpSocketBufferSizeInBytes( + Integer.parseInt( + properties.getProperty( + "pipe_air_gap_receiver_udp_socket_buffer_size_in_bytes", + Integer.toString(config.getPipeAirGapReceiverUdpSocketBufferSizeInBytes())))); config.setPipeReceiverLoadConversionEnabled( Boolean.parseBoolean( properties.getProperty( diff --git a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java index c04f6b87c79a4..d761d4008300e 100644 --- a/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java +++ b/iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/pipe/config/constant/PipeSinkConstant.java @@ -211,6 +211,11 @@ private static String getDefaultConnectorOrSinkName(final PipeParameters paramet public static final int CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE = 60 * 1024; public static final int CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE = 1024; public static final int CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE = 65_507; + public static final String CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY = + "connector.air-gap.udp.send-interval-ms"; + public static final String SINK_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY = + "sink.air-gap.udp.send-interval-ms"; + public static final int CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_DEFAULT_VALUE = 5; public static final String CONNECTOR_IOTDB_SYNC_CONNECTOR_VERSION_KEY = "connector.version"; public static final String SINK_IOTDB_SYNC_CONNECTOR_VERSION_KEY = "sink.version"; 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 a7d59713d5109..0102cf4a093e5 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 @@ -48,6 +48,7 @@ import java.io.DataInputStream; import java.io.File; import java.io.IOException; +import java.io.InterruptedIOException; import java.io.RandomAccessFile; import java.net.DatagramPacket; import java.net.DatagramSocket; @@ -80,6 +81,8 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_DEFAULT_VALUE; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_PRIORITY_STRATEGY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_RANDOM_STRATEGY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_LOAD_BALANCE_ROUND_ROBIN_STRATEGY; @@ -88,6 +91,7 @@ import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_STRICT_UNIDIRECTIONAL_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_TRANSPORT_KEY; import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_UDP_PACKET_SIZE_KEY; +import static org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.SINK_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY; import static org.apache.iotdb.commons.utils.BasicStructureSerDeUtil.LONG_LEN; @TreeModel @@ -105,6 +109,8 @@ protected static class AirGapSocket extends Socket { private DatagramSocket datagramSocket; private InetAddress datagramAddress; private boolean isUdp; + private long lastStrictUdpSendNanos; + private boolean hasSentStrictUdpDatagram; public AirGapSocket(final String ip, final int port) { this.endPoint = new TEndPoint(ip, port); @@ -188,6 +194,7 @@ public String toString() { private boolean useUdpTransport; private boolean strictUnidirectional; private int udpPacketSizeInBytes; + private long strictUdpSendIntervalNanos; // The air gap connector does not use clientManager thus we put handshake type here protected boolean supportModsIfIsDataNodeReceiver = true; @@ -243,6 +250,17 @@ public void validate(final PipeParameterValidator validator) throws Exception { CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE, packetSize), packetSize); + final int sendIntervalMs = + parameters.getIntOrDefault( + Arrays.asList( + CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY, SINK_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY), + CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_DEFAULT_VALUE); + validator.validate( + arg -> sendIntervalMs >= 0, + String.format( + PipeMessages.EXCEPTION_UDP_SEND_INTERVAL_MUST_BE_NON_NEGATIVE_BUT_GOT_ARG_36D41687, + sendIntervalMs), + sendIntervalMs); } @Override @@ -308,6 +326,13 @@ public void customize( parameters.getIntOrDefault( Arrays.asList(CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY, SINK_AIR_GAP_UDP_PACKET_SIZE_KEY), CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE); + strictUdpSendIntervalNanos = + TimeUnit.MILLISECONDS.toNanos( + parameters.getIntOrDefault( + Arrays.asList( + CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY, + SINK_AIR_GAP_UDP_SEND_INTERVAL_MS_KEY), + CONNECTOR_AIR_GAP_UDP_SEND_INTERVAL_MS_DEFAULT_VALUE)); LOGGER.info( "Air gap transport is {}, udp packet size is {} bytes.", useUdpTransport ? "udp" : "tcp", @@ -584,6 +609,28 @@ private boolean sendOneDatagram(final AirGapSocket socket, final TPipeTransferRe final byte[] datagramBytes = datagramBuffer.array(); final DatagramSocket datagramSocket = socket.getDatagramSocket(); + if (strictUnidirectional && strictUdpSendIntervalNanos > 0) { + // No ACK supplies backpressure on a one-way link. Space all datagrams, including + // consecutive slices of a large request, so a single event cannot flood the socket. + if (socket.hasSentStrictUdpDatagram) { + long remainingNanos = + strictUdpSendIntervalNanos - (System.nanoTime() - socket.lastStrictUdpSendNanos); + while (remainingNanos > 0) { + try { + TimeUnit.NANOSECONDS.sleep(remainingNanos); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + final InterruptedIOException exception = new InterruptedIOException(e.getMessage()); + exception.initCause(e); + throw exception; + } + remainingNanos = + strictUdpSendIntervalNanos - (System.nanoTime() - socket.lastStrictUdpSendNanos); + } + } + socket.lastStrictUdpSendNanos = System.nanoTime(); + socket.hasSentStrictUdpDatagram = true; + } datagramSocket.send( new DatagramPacket( datagramBytes, diff --git a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/config/PipeDescriptorTest.java b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/config/PipeDescriptorTest.java index 00d982122719b..e8ae24f44b557 100644 --- a/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/config/PipeDescriptorTest.java +++ b/iotdb-core/node-commons/src/test/java/org/apache/iotdb/commons/pipe/config/PipeDescriptorTest.java @@ -33,15 +33,20 @@ public class PipeDescriptorTest { private final CommonConfig config = CommonDescriptor.getInstance().getConfig(); private int originalRequestSliceThresholdBytes; + private int originalAirGapUdpSocketBufferSizeInBytes; @Before public void setUp() { originalRequestSliceThresholdBytes = config.getPipeSinkRequestSliceThresholdBytes(); + originalAirGapUdpSocketBufferSizeInBytes = + config.getPipeAirGapReceiverUdpSocketBufferSizeInBytes(); } @After public void tearDown() { config.setPipeSinkRequestSliceThresholdBytes(originalRequestSliceThresholdBytes); + config.setPipeAirGapReceiverUdpSocketBufferSizeInBytes( + originalAirGapUdpSocketBufferSizeInBytes); } @Test @@ -62,4 +67,15 @@ public void testPipeRequestSliceThresholdSupportsSinkAndConnectorKeys() { PipeDescriptor.loadPipeInternalConfig(config, bothProperties); Assert.assertEquals(456, config.getPipeSinkRequestSliceThresholdBytes()); } + + @Test + public void testAirGapReceiverUdpSocketBufferIsLoadedFromGlobalProperties() { + final TrimProperties properties = new TrimProperties(); + properties.setProperty("pipe_air_gap_receiver_udp_socket_buffer_size_in_bytes", "1048576"); + + PipeDescriptor.loadPipeStaticConfig(config, properties); + PipeDescriptor.loadPipeInternalConfig(config, properties); + + Assert.assertEquals(1048576, config.getPipeAirGapReceiverUdpSocketBufferSizeInBytes()); + } } From 358020e0a4a476660d144ce6398ff3ecc989e90d Mon Sep 17 00:00:00 2001 From: Caideyipi <87789683+Caideyipi@users.noreply.github.com> Date: Mon, 2 Nov 2026 12:16:51 +0800 Subject: [PATCH 5/5] Stop air-gap receiver retries when interrupted --- .../protocol/airgap/IoTDBAirGapReceiver.java | 4 +++ .../airgap/IoTDBAirGapReceiverTest.java | 36 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java index 8c3bfe9d8b454..a8f030c7744f4 100644 --- a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java +++ b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java @@ -45,6 +45,7 @@ import java.io.FilterInputStream; import java.io.IOException; import java.io.InputStream; +import java.io.InterruptedIOException; import java.io.OutputStream; import java.net.DatagramPacket; import java.net.DatagramSocket; @@ -262,6 +263,9 @@ private void handleReq( Thread.sleep(PipeConfig.getInstance().getPipeAirGapRetryLocalIntervalMs()); } catch (final InterruptedException e) { Thread.currentThread().interrupt(); + final InterruptedIOException exception = new InterruptedIOException(e.getMessage()); + exception.initCause(e); + throw exception; } if (System.currentTimeMillis() - startTime < PipeConfig.getInstance().getPipeAirGapRetryMaxMs()) { diff --git a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java index f9933260fb984..bf12a2a1938b9 100644 --- a/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java +++ b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverTest.java @@ -48,6 +48,7 @@ import java.io.DataOutputStream; import java.io.IOException; import java.io.InputStream; +import java.io.InterruptedIOException; import java.io.OutputStream; import java.lang.reflect.Field; import java.lang.reflect.Method; @@ -260,6 +261,41 @@ public void send(final DatagramPacket packet) { } } + @Test + public void testInterruptedUdpRetryStopsWithoutReenteringReceiver() throws Exception { + final AtomicInteger receiveCount = new AtomicInteger(); + final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(new RecordingSocket(), 7L); + final StubIoTDBDataNodeReceiverAgent stubAgent = new StubIoTDBDataNodeReceiverAgent(); + stubAgent.setStubReceiver( + "udp-client", + new StubReceiver( + new TSStatus( + TSStatusCode.PIPE_RECEIVER_TEMPORARY_UNAVAILABLE_EXCEPTION.getStatusCode())) { + @Override + public TPipeTransferResp receive(final TPipeTransferReq req) { + Assert.assertEquals(1, receiveCount.incrementAndGet()); + return super.receive(req); + } + }); + setField(receiver, "agent", stubAgent); + try (final DatagramSocket receiverSocket = new NoResponseDatagramSocket()) { + final byte[] request = + toUdpDatagram(AirGapUdpProtocol.NO_ACK_REQUEST_ID_MASK | 75L, toRawRequestBytes()); + final DatagramPacket packet = + new DatagramPacket(request, request.length, InetAddress.getLoopbackAddress(), 12345); + Thread.currentThread().interrupt(); + try { + Assert.assertThrows( + InterruptedIOException.class, + () -> receiver.receiveUdp(receiverSocket, packet, "udp-client", request, () -> {})); + Assert.assertTrue(Thread.currentThread().isInterrupted()); + Assert.assertEquals(1, receiveCount.get()); + } finally { + Thread.interrupted(); + } + } + } + @Test public void testTemporaryUnavailableRetryUsesFreshRequestBody() throws Exception { assertTemporaryUnavailableRetryUsesFreshRequestBody(null);