This is an automated email from the ASF dual-hosted git repository.
jt2594838 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iotdb.git
The following commit(s) were added to refs/heads/master by this push:
new 574d5eb6278 Implement UDP transport for air-gap sink (#17954)
574d5eb6278 is described below
commit 574d5eb6278e6f34aea9d92280e126bdc297bb38
Author: Caideyipi <[email protected]>
AuthorDate: Fri Oct 30 09:31:01 2026 +0800
Implement UDP transport for air-gap sink (#17954)
* Implement UDP transport for air-gap sink
* Update IoTDBDataRegionAirGapSinkTest.java
* Reuse thread pool
* Fix air-gap receiver Sonar issues
* Fix air gap receiver retry test
* Fix air-gap UDP ACK correlation and session cleanup
* Address air-gap UDP review comments
---
.../apache/iotdb/db/i18n/DataNodePipeMessages.java | 4 +
.../apache/iotdb/db/i18n/DataNodePipeMessages.java | 4 +
.../protocol/airgap/IoTDBAirGapReceiver.java | 163 +++++++++---
.../protocol/airgap/IoTDBAirGapReceiverAgent.java | 177 ++++++++++++-
.../thrift/IoTDBDataNodeReceiverAgent.java | 24 +-
.../airgap/IoTDBAirGapReceiverAgentTest.java | 120 +++++++++
.../protocol/airgap/IoTDBAirGapReceiverTest.java | 110 +++++++-
.../airgap/IoTDBDataRegionAirGapSinkTest.java | 238 ++++++++++++++++-
.../apache/iotdb/commons/i18n/PipeMessages.java | 8 +
.../apache/iotdb/commons/i18n/PipeMessages.java | 8 +
.../pipe/config/constant/PipeSinkConstant.java | 20 ++
.../pipe/sink/protocol/IoTDBAirGapSink.java | 290 ++++++++++++++++++++-
12 files changed, 1113 insertions(+), 53 deletions(-)
diff --git
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
index cf67838025e..2c33424cf99 100644
---
a/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
@@ -1323,6 +1323,10 @@ public final class DataNodePipeMessages {
"Pipe air gap receiver {} started. Socket: {}";
public static final String PIPE_AIR_GAP_RECEIVER_TEMPORARY_UNAVAILABLE_RETRY
=
"Pipe air gap receiver {}: Temporary unavailable retry timed out,
returning FAIL to sender.";
+ public static final String
EXCEPTION_AIR_GAP_UDP_PACKET_IS_MISSING_ITS_REQUEST_ID_C4228CC8 =
+ "Air-gap UDP packet is missing its request ID.";
+ public static final String
EXCEPTION_AIR_GAP_UDP_PACKET_CONTAINS_TRAILING_BYTES_2D34D418 =
+ "Air-gap UDP packet contains trailing bytes.";
public static final String PIPE_DATA_TRANSPORT_ERROR = "Pipe data transport
error, {}";
public static final String PIPE_INSERTING_ROW_CASTING_TYPE_FROM =
"Pipe: Inserting row. Casting type from {} to {}.";
diff --git
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
index 4e1e7a0f2fd..064fb8b4032 100644
---
a/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
+++
b/iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java
@@ -1242,6 +1242,10 @@ public final class DataNodePipeMessages {
"Pipe air gap receiver {}：监听数据时 socket {} 关闭。原因：{}";
public static final String PIPE_AIR_GAP_RECEIVER_STARTED_SOCKET =
"Pipe air gap receiver {} 已启动。Socket：{}";
+ public static final String
EXCEPTION_AIR_GAP_UDP_PACKET_IS_MISSING_ITS_REQUEST_ID_C4228CC8 =
+ "Air-gap UDP 数据包缺少 request ID。";
+ public static final String
EXCEPTION_AIR_GAP_UDP_PACKET_CONTAINS_TRAILING_BYTES_2D34D418 =
+ "Air-gap UDP 数据包包含多余字节。";
public static final String PIPE_AIR_GAP_RECEIVER_TEMPORARY_UNAVAILABLE_RETRY
=
"Pipe air gap receiver {}：临时不可用重试超时，向 sender 返回 FAIL。";
public static final String PIPE_DATA_TRANSPORT_ERROR = "Pipe 数据传输错误，{}";
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 48fb42d27c2..e3413b009b2 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
@@ -40,11 +40,15 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.BufferedInputStream;
+import java.io.ByteArrayInputStream;
import java.io.FilterInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
import java.net.Socket;
+import java.net.SocketAddress;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.util.Arrays;
@@ -63,6 +67,10 @@ public class IoTDBAirGapReceiver extends WrappedRunnable {
private final IoTDBDataNodeReceiverAgent agent;
private boolean isELanguagePayload;
+ private OutputStream currentOutputStream;
+ private DatagramSocket currentDatagramSocket;
+ private SocketAddress currentRemoteSocketAddress;
+ private Long currentDatagramRequestId;
public IoTDBAirGapReceiver(final Socket socket, final long receiverId) {
this.socket = socket;
@@ -105,33 +113,11 @@ public class IoTDBAirGapReceiver extends WrappedRunnable {
try {
final byte[] data = readData(inputStream);
+ currentOutputStream = socket.getOutputStream();
- // If check sum failed, it indicates that the length we read may not be
correct.
- // Namely, there may be remaining bytes in the socket stream, which will
fail any subsequent
- // attempts to read from that.
- // We directly close the socket here.
- if (!checkSum(data)) {
- LOGGER.warn(
- DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_CLOSED_BECAUSE_OF,
receiverId, socket);
- try {
- fail();
- } finally {
- socket.close();
- }
- return;
+ if (!receive(data, null)) {
+ socket.close();
}
-
- // Removed the used checksum
- final ByteBuffer byteBuffer = ByteBuffer.wrap(data, LONG_LEN,
data.length - LONG_LEN);
-
- // Pseudo request, to reuse logic in IoTDBThriftReceiverAgent
- final AirGapPseudoTPipeTransferRequest req =
- (AirGapPseudoTPipeTransferRequest)
- new AirGapPseudoTPipeTransferRequest()
- .setVersion(ReadWriteIOUtils.readByte(byteBuffer))
- .setType(ReadWriteIOUtils.readShort(byteBuffer))
- .setBody(byteBuffer.slice());
- handleReq(req, System.currentTimeMillis());
} catch (final SocketTimeoutException e) {
// It is normal for an air gap sender to remain idle. Only close the
connection when the
// timeout occurs after a request has started, because the stream can no
longer be decoded
@@ -161,9 +147,103 @@ public class IoTDBAirGapReceiver extends WrappedRunnable {
}
}
- private void handleReq(final AirGapPseudoTPipeTransferRequest req, final
long startTime)
+ boolean receive(final byte[] data, final String receiverKey) throws
IOException {
+ return receive(data, receiverKey, null);
+ }
+
+ private boolean receive(
+ final byte[] data, final String receiverKey, final Runnable
onValidRequest)
throws IOException {
- final TPipeTransferResp resp = agent.receive(duplicateReq(req));
+ try {
+ // If check sum failed, it indicates that the length we read may not be
correct.
+ // Namely, there may be remaining bytes in the socket stream, which will
fail any subsequent
+ // attempts to read from that.
+ if (!checkSum(data)) {
+ LOGGER.warn(
+ DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_CLOSED_BECAUSE_OF,
receiverId, socket);
+ fail();
+ return false;
+ }
+
+ if (onValidRequest != null) {
+ onValidRequest.run();
+ }
+ handleReq(toTPipeTransferReq(data), System.currentTimeMillis(),
receiverKey);
+ return true;
+ } catch (final IOException e) {
+ throw e;
+ } catch (final Exception e) {
+ if (currentDatagramSocket != null) {
+ fail();
+ }
+ throw e;
+ } finally {
+ currentOutputStream = null;
+ currentDatagramSocket = null;
+ currentRemoteSocketAddress = null;
+ currentDatagramRequestId = null;
+ }
+ }
+
+ void receiveUdp(
+ final DatagramSocket datagramSocket,
+ final DatagramPacket packet,
+ final String receiverKey,
+ final byte[] buffer,
+ final Runnable onValidRequest)
+ throws IOException {
+ isELanguagePayload = false;
+ currentDatagramSocket = datagramSocket;
+ currentRemoteSocketAddress = packet.getSocketAddress();
+ boolean requestHandlingStarted = false;
+ try {
+ if (packet.getLength() < Long.BYTES) {
+ throw new IOException(
+
DataNodePipeMessages.EXCEPTION_AIR_GAP_UDP_PACKET_IS_MISSING_ITS_REQUEST_ID_C4228CC8);
+ }
+ currentDatagramRequestId = ByteBuffer.wrap(buffer, 0,
Long.BYTES).getLong();
+ try (final ByteArrayInputStream inputStream =
+ new ByteArrayInputStream(buffer, Long.BYTES, packet.getLength() -
Long.BYTES)) {
+ final byte[] data = readData(inputStream);
+ if (inputStream.available() != 0) {
+ throw new IOException(
+
DataNodePipeMessages.EXCEPTION_AIR_GAP_UDP_PACKET_CONTAINS_TRAILING_BYTES_2D34D418);
+ }
+ requestHandlingStarted = true;
+ receive(data, receiverKey, onValidRequest);
+ }
+ } catch (final Exception e) {
+ if (!requestHandlingStarted) {
+ fail();
+ }
+ throw e;
+ } finally {
+ currentOutputStream = null;
+ currentDatagramSocket = null;
+ currentRemoteSocketAddress = null;
+ currentDatagramRequestId = null;
+ }
+ }
+
+ private AirGapPseudoTPipeTransferRequest toTPipeTransferReq(final byte[]
data) {
+ // Removed the used checksum
+ final ByteBuffer byteBuffer = ByteBuffer.wrap(data, LONG_LEN, data.length
- LONG_LEN);
+
+ // Pseudo request, to reuse logic in IoTDBThriftReceiverAgent
+ return (AirGapPseudoTPipeTransferRequest)
+ new AirGapPseudoTPipeTransferRequest()
+ .setVersion(ReadWriteIOUtils.readByte(byteBuffer))
+ .setType(ReadWriteIOUtils.readShort(byteBuffer))
+ .setBody(byteBuffer.slice());
+ }
+
+ private void handleReq(
+ final AirGapPseudoTPipeTransferRequest req, final long startTime, final
String receiverKey)
+ throws IOException {
+ final TPipeTransferResp resp =
+ receiverKey == null
+ ? agent.receive(duplicateReq(req))
+ : agent.receive(receiverKey, duplicateReq(req));
final TSStatus status = resp.getStatus();
if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
@@ -181,7 +261,7 @@ public class IoTDBAirGapReceiver extends WrappedRunnable {
}
if (System.currentTimeMillis() - startTime
< PipeConfig.getInstance().getPipeAirGapRetryMaxMs()) {
- handleReq(req, startTime);
+ handleReq(req, startTime, receiverKey);
} else {
LOGGER.warn(
DataNodePipeMessages.PIPE_AIR_GAP_RECEIVER_TEMPORARY_UNAVAILABLE_RETRY,
receiverId);
@@ -207,14 +287,31 @@ public class IoTDBAirGapReceiver extends WrappedRunnable {
}
private void ok() throws IOException {
- final OutputStream outputStream = socket.getOutputStream();
- outputStream.write(AirGapOneByteResponse.OK);
- outputStream.flush();
+ respond(AirGapOneByteResponse.OK);
}
private void fail() throws IOException {
- final OutputStream outputStream = socket.getOutputStream();
- outputStream.write(AirGapOneByteResponse.FAIL);
+ respond(AirGapOneByteResponse.FAIL);
+ }
+
+ private void respond(final byte[] response) throws IOException {
+ if (currentDatagramSocket != null && currentRemoteSocketAddress != null) {
+ if (currentDatagramRequestId == null) {
+ return;
+ }
+ final byte[] responseBytes =
+ ByteBuffer.allocate(Long.BYTES + response.length)
+ .putLong(currentDatagramRequestId)
+ .put(response)
+ .array();
+ currentDatagramSocket.send(
+ new DatagramPacket(responseBytes, responseBytes.length,
currentRemoteSocketAddress));
+ return;
+ }
+
+ final OutputStream outputStream =
+ currentOutputStream != null ? currentOutputStream :
socket.getOutputStream();
+ outputStream.write(response);
outputStream.flush();
}
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 c14ce2054b7..87683081d9b 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
@@ -26,30 +26,46 @@ import org.apache.iotdb.commons.pipe.config.PipeConfig;
import org.apache.iotdb.commons.service.IService;
import org.apache.iotdb.commons.service.ServiceType;
import org.apache.iotdb.db.i18n.DataNodePipeMessages;
+import org.apache.iotdb.db.pipe.agent.PipeDataNodeAgent;
+import org.apache.iotdb.db.protocol.session.ClientSession;
+import org.apache.iotdb.db.protocol.session.SessionManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.io.IOException;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
+import java.net.InetAddress;
import java.net.ServerSocket;
import java.net.Socket;
+import java.net.SocketTimeoutException;
+import java.util.Map;
import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
public class IoTDBAirGapReceiverAgent implements IService {
private static final Logger LOGGER =
LoggerFactory.getLogger(IoTDBAirGapReceiverAgent.class);
+ 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 final ExecutorService listenExecutor =
- IoTDBThreadPoolFactory.newSingleThreadExecutor(
- ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName());
+ IoTDBThreadPoolFactory.newFixedThreadPool(
+ LISTENER_THREAD_COUNT,
ThreadName.PIPE_RECEIVER_AIR_GAP_AGENT.getName());
private final AtomicBoolean allowSubmitListen = new AtomicBoolean(false);
private ServerSocket serverSocket;
+ private DatagramSocket datagramSocket;
private final AtomicLong receiverId = new AtomicLong(0);
+ final Map<String, UdpClientSession> udpClientSessions = new
ConcurrentHashMap<>();
public void listen() {
try {
@@ -61,7 +77,9 @@ public class IoTDBAirGapReceiverAgent implements IService {
ThreadName.PIPE_AIR_GAP_RECEIVER.getName() + "-" + airGapReceiverId);
airGapReceiverThread.start();
} catch (final IOException e) {
-
LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER,
e);
+ if (allowSubmitListen.get()) {
+
LOGGER.warn(DataNodePipeMessages.UNHANDLED_EXCEPTION_DURING_PIPE_AIR_GAP_RECEIVER,
e);
+ }
}
if (allowSubmitListen.get()) {
@@ -69,31 +87,150 @@ public class IoTDBAirGapReceiverAgent implements IService {
}
}
+ 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};
+ try {
+ receiver.receiveUdp(
+ datagramSocket,
+ packet,
+ receiverKey,
+ buffer,
+ () -> registeredSession[0] =
registerUdpSessionIfNecessary(packet));
+ } finally {
+ if (registeredSession[0]) {
+ SessionManager.getInstance().removeCurrSession();
+ }
+ }
+ } catch (final SocketTimeoutException e) {
+ removeExpiredUdpClientSessions(System.nanoTime());
+ } 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);
+ }
+ }
+ }
+
+ boolean registerUdpSessionIfNecessary(final DatagramPacket packet) {
+ final String receiverKey = packet.getSocketAddress().toString();
+ final long now = System.nanoTime();
+ removeExpiredUdpClientSessions(now);
+
+ UdpClientSession clientSession = udpClientSessions.get(receiverKey);
+ if (clientSession == null) {
+ ensureUdpClientSessionCapacity();
+ clientSession =
+ new UdpClientSession(
+ new ClientSession(new DatagramClientSocket(packet.getAddress(),
packet.getPort())),
+ now);
+ final UdpClientSession existing =
udpClientSessions.putIfAbsent(receiverKey, clientSession);
+ if (existing != null) {
+ clientSession = existing;
+ }
+ }
+ clientSession.lastAccessTimeNanos = now;
+ return SessionManager.getInstance().registerSession(clientSession.session);
+ }
+
+ private void ensureUdpClientSessionCapacity() {
+ if (udpClientSessions.size() < UDP_MAX_CLIENT_SESSIONS) {
+ return;
+ }
+
+ String leastRecentlyUsedKey = null;
+ UdpClientSession leastRecentlyUsedSession = null;
+ for (final Map.Entry<String, UdpClientSession> entry :
udpClientSessions.entrySet()) {
+ if (leastRecentlyUsedSession == null
+ || entry.getValue().lastAccessTimeNanos <
leastRecentlyUsedSession.lastAccessTimeNanos) {
+ leastRecentlyUsedKey = entry.getKey();
+ leastRecentlyUsedSession = entry.getValue();
+ }
+ }
+ if (leastRecentlyUsedKey != null) {
+ removeUdpClientSession(leastRecentlyUsedKey, leastRecentlyUsedSession);
+ }
+ }
+
+ void removeExpiredUdpClientSessions(final long nowNanos) {
+ final long idleTimeoutNanos =
+ TimeUnit.MILLISECONDS.toNanos(
+ Math.max(1,
PipeConfig.getInstance().getPipeSinkTransferTimeoutMs()));
+ udpClientSessions.forEach(
+ (key, session) -> {
+ if (nowNanos - session.lastAccessTimeNanos >= idleTimeoutNanos) {
+ removeUdpClientSession(key, session);
+ }
+ });
+ }
+
+ private void removeUdpClientSession(final String key, final UdpClientSession
session) {
+ if (!udpClientSessions.remove(key, session)) {
+ return;
+ }
+
+ final boolean registeredSession =
SessionManager.getInstance().registerSession(session.session);
+ try {
+ PipeDataNodeAgent.receiver().thrift().handleClientExit(key);
+ } finally {
+ if (registeredSession) {
+ SessionManager.getInstance().removeCurrSession();
+ }
+ }
+ }
+
@Override
public void start() throws StartupException {
try {
serverSocket = new
ServerSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort());
+ datagramSocket = new
DatagramSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort());
+ datagramSocket.setSoTimeout(UDP_SESSION_CLEANUP_INTERVAL_MS);
} catch (final IOException e) {
+ if (Objects.nonNull(serverSocket)) {
+ try {
+ serverSocket.close();
+ } catch (final IOException closeException) {
+ e.addSuppressed(closeException);
+ }
+ }
throw new StartupException(e);
}
allowSubmitListen.set(true);
listenExecutor.submit(this::listen);
+ listenExecutor.submit(this::listenUdp);
LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STARTED,
serverSocket);
}
@Override
public void stop() {
+ allowSubmitListen.set(false);
+
try {
if (Objects.nonNull(serverSocket)) {
serverSocket.close();
}
+ if (Objects.nonNull(datagramSocket)) {
+ datagramSocket.close();
+ }
} catch (final IOException e) {
LOGGER.warn(DataNodePipeMessages.FAILED_TO_CLOSE_IOTDBAIRGAPRECEIVERAGENT_S_SERVER_SOCKET,
e);
}
- allowSubmitListen.set(false);
+ udpClientSessions.forEach(this::removeUdpClientSession);
listenExecutor.shutdown();
LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STOPPED,
serverSocket);
@@ -103,4 +240,36 @@ public class IoTDBAirGapReceiverAgent implements IService {
public ServiceType getID() {
return ServiceType.AIR_GAP_SERVICE;
}
+
+ private static class DatagramClientSocket extends Socket {
+
+ private final InetAddress address;
+ private final int port;
+
+ private DatagramClientSocket(final InetAddress address, final int port) {
+ this.address = address;
+ this.port = port;
+ }
+
+ @Override
+ public InetAddress getInetAddress() {
+ return address;
+ }
+
+ @Override
+ public int getPort() {
+ return port;
+ }
+ }
+
+ static class UdpClientSession {
+
+ private final ClientSession session;
+ volatile long lastAccessTimeNanos;
+
+ UdpClientSession(final ClientSession session, final long
lastAccessTimeNanos) {
+ this.session = session;
+ this.lastAccessTimeNanos = lastAccessTimeNanos;
+ }
+ }
}
diff --git
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java
index 8a995f7656d..57ed64ec520 100644
---
a/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java
+++
b/iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiverAgent.java
@@ -24,9 +24,13 @@ import
org.apache.iotdb.commons.pipe.receiver.IoTDBReceiverAgent;
import
org.apache.iotdb.commons.pipe.sink.payload.thrift.request.IoTDBSinkRequestVersion;
import
org.apache.iotdb.db.pipe.processor.twostage.exchange.receiver.TwoStageAggregateReceiver;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
public class IoTDBDataNodeReceiverAgent extends IoTDBReceiverAgent {
private final ThreadLocal<IoTDBReceiver> receiverThreadLocal = new
ThreadLocal<>();
+ private final Map<String, IoTDBReceiver> receiverMap = new
ConcurrentHashMap<>();
@Override
protected void initConstructors() {
@@ -37,17 +41,25 @@ public class IoTDBDataNodeReceiverAgent extends
IoTDBReceiverAgent {
}
@Override
- protected IoTDBReceiver getReceiverWithSpecifiedClient(final String ignore) {
- return receiverThreadLocal.get();
+ protected IoTDBReceiver getReceiverWithSpecifiedClient(final String key) {
+ return key == null ? receiverThreadLocal.get() : receiverMap.get(key);
}
@Override
- protected void setReceiverWithSpecifiedClient(final String ignore, final
IoTDBReceiver receiver) {
- receiverThreadLocal.set(receiver);
+ protected void setReceiverWithSpecifiedClient(final String key, final
IoTDBReceiver receiver) {
+ if (key == null) {
+ receiverThreadLocal.set(receiver);
+ } else {
+ receiverMap.put(key, receiver);
+ }
}
@Override
- protected void removeReceiverWithSpecifiedClient(final String ignore) {
- receiverThreadLocal.remove();
+ protected void removeReceiverWithSpecifiedClient(final String key) {
+ if (key == null) {
+ receiverThreadLocal.remove();
+ } else {
+ receiverMap.remove(key);
+ }
}
}
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
new file mode 100644
index 00000000000..ffbebb5a668
--- /dev/null
+++
b/iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgentTest.java
@@ -0,0 +1,120 @@
+/*
+ * 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.airgap;
+
+import org.apache.iotdb.common.rpc.thrift.TSStatus;
+import org.apache.iotdb.commons.pipe.config.PipeConfig;
+import org.apache.iotdb.commons.pipe.receiver.IoTDBReceiver;
+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;
+import org.apache.iotdb.db.protocol.session.ClientSession;
+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.IOException;
+import java.lang.reflect.Field;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
+import java.net.InetAddress;
+import java.net.Socket;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+public class IoTDBAirGapReceiverAgentTest {
+
+ @Test
+ public void testExpiredUdpClientSessionIsRemoved() throws Exception {
+ final IoTDBAirGapReceiverAgent agent = new IoTDBAirGapReceiverAgent();
+ final String sessionKey = "expired-client";
+ final IoTDBAirGapReceiverAgent.UdpClientSession session =
+ new IoTDBAirGapReceiverAgent.UdpClientSession(new ClientSession(new
Socket()), 0L);
+ agent.udpClientSessions.put(sessionKey, session);
+
+ final AtomicBoolean receiverExited = new AtomicBoolean(false);
+ final IoTDBDataNodeReceiverAgent receiverAgent =
PipeDataNodeAgent.receiver().thrift();
+ final Field receiverMapField =
IoTDBDataNodeReceiverAgent.class.getDeclaredField("receiverMap");
+ receiverMapField.setAccessible(true);
+ @SuppressWarnings("unchecked")
+ final Map<String, IoTDBReceiver> receiverMap =
+ (Map<String, IoTDBReceiver>) receiverMapField.get(receiverAgent);
+ receiverMap.put(
+ sessionKey,
+ new IoTDBReceiver() {
+ @Override
+ public TPipeTransferResp receive(final TPipeTransferReq req) {
+ return new TPipeTransferResp(new
TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode()));
+ }
+
+ @Override
+ public void handleExit() {
+ receiverExited.set(true);
+ }
+
+ @Override
+ public IoTDBSinkRequestVersion getVersion() {
+ return IoTDBSinkRequestVersion.VERSION_1;
+ }
+ });
+
+ try {
+ final long timeoutNanos =
+
TimeUnit.MILLISECONDS.toNanos(PipeConfig.getInstance().getPipeSinkTransferTimeoutMs());
+ agent.removeExpiredUdpClientSessions(timeoutNanos + 1);
+
+ Assert.assertTrue(agent.udpClientSessions.isEmpty());
+ Assert.assertTrue(receiverExited.get());
+ } finally {
+ receiverMap.remove(sessionKey);
+ }
+ }
+
+ @Test
+ public void testMalformedUdpPacketDoesNotCreateClientSession() throws
Exception {
+ try (final DatagramSocket datagramSocket = new DatagramSocket()) {
+ final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(new
Socket(), 1L);
+ final byte[] malformedPacket = new byte[] {1};
+ final DatagramPacket packet =
+ new DatagramPacket(
+ malformedPacket,
+ malformedPacket.length,
+ InetAddress.getLoopbackAddress(),
+ datagramSocket.getLocalPort());
+ final AtomicBoolean sessionRegistrationAttempted = new
AtomicBoolean(false);
+
+ Assert.assertThrows(
+ IOException.class,
+ () ->
+ receiver.receiveUdp(
+ datagramSocket,
+ packet,
+ "invalid-client",
+ malformedPacket,
+ () -> sessionRegistrationAttempted.set(true)));
+
+ Assert.assertFalse(sessionRegistrationAttempted.get());
+ }
+ }
+}
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 d7e69b51cfd..57cacb38799 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
@@ -36,6 +36,7 @@ import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
import org.apache.iotdb.service.rpc.thrift.TPipeTransferResp;
import org.apache.tsfile.utils.BytesUtils;
+import org.apache.tsfile.utils.ReadWriteIOUtils;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
@@ -43,18 +44,23 @@ import org.junit.Test;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.lang.reflect.Field;
import java.lang.reflect.Method;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.Socket;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
+import java.util.Arrays;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.zip.CRC32;
public class IoTDBAirGapReceiverTest {
@@ -136,9 +142,9 @@ public class IoTDBAirGapReceiverTest {
final Method handleReq =
IoTDBAirGapReceiver.class.getDeclaredMethod(
- "handleReq", AirGapPseudoTPipeTransferRequest.class, long.class);
+ "handleReq", AirGapPseudoTPipeTransferRequest.class, long.class,
String.class);
handleReq.setAccessible(true);
- handleReq.invoke(receiver, req, System.currentTimeMillis() - 10_000L);
+ handleReq.invoke(receiver, req, System.currentTimeMillis() - 10_000L,
null);
Assert.assertArrayEquals(AirGapOneByteResponse.FAIL,
socket.getWrittenBytes());
} finally {
@@ -147,8 +153,68 @@ public class IoTDBAirGapReceiverTest {
}
}
+ @Test
+ public void testReceiveAirGapPayloadWithSpecifiedReceiverKey() throws
Exception {
+ final RecordingSocket socket = new RecordingSocket();
+ final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(socket, 4L);
+ final StubIoTDBDataNodeReceiverAgent stubAgent = new
StubIoTDBDataNodeReceiverAgent();
+ stubAgent.setStubReceiver(
+ "udp-client", new StubReceiver(new
TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())));
+ setField(receiver, "agent", stubAgent);
+
+ Assert.assertTrue(receiver.receive(toAirGapData(toRawRequestBytes()),
"udp-client"));
+ Assert.assertArrayEquals(AirGapOneByteResponse.OK,
socket.getWrittenBytes());
+ }
+
+ @Test
+ public void testUdpResponseContainsRequestId() throws Exception {
+ try (final DatagramSocket receiverSocket = new DatagramSocket(0);
+ final DatagramSocket senderSocket = new DatagramSocket()) {
+ final long requestId = 73L;
+ final byte[] request = toUdpDatagram(requestId, toRawRequestBytes());
+ senderSocket.send(
+ new DatagramPacket(
+ request,
+ request.length,
+ InetAddress.getLoopbackAddress(),
+ receiverSocket.getLocalPort()));
+
+ final byte[] buffer = new byte[65_507];
+ final DatagramPacket packet = new DatagramPacket(buffer, buffer.length);
+ receiverSocket.receive(packet);
+
+ final RecordingSocket socket = new RecordingSocket();
+ final IoTDBAirGapReceiver receiver = new IoTDBAirGapReceiver(socket, 5L);
+ final StubIoTDBDataNodeReceiverAgent stubAgent = new
StubIoTDBDataNodeReceiverAgent();
+ stubAgent.setStubReceiver(
+ "udp-client",
+ new StubReceiver(new
TSStatus(TSStatusCode.SUCCESS_STATUS.getStatusCode())));
+ setField(receiver, "agent", stubAgent);
+ receiver.receiveUdp(receiverSocket, packet, "udp-client", buffer, () ->
{});
+
+ final byte[] response = new byte[Long.BYTES + 2];
+ final DatagramPacket responsePacket = new DatagramPacket(response,
response.length);
+ senderSocket.setSoTimeout(1000);
+ senderSocket.receive(responsePacket);
+ Assert.assertEquals(Long.BYTES + 1, responsePacket.getLength());
+ Assert.assertEquals(requestId, ByteBuffer.wrap(response).getLong());
+ Assert.assertEquals(AirGapOneByteResponse.OK[0], response[Long.BYTES]);
+ }
+ }
+
@Test
public void testTemporaryUnavailableRetryUsesFreshRequestBody() throws
Exception {
+ assertTemporaryUnavailableRetryUsesFreshRequestBody(null);
+ }
+
+ @Test
+ public void
testTemporaryUnavailableRetryUsesFreshRequestBodyWithSpecifiedReceiverKey()
+ throws Exception {
+ assertTemporaryUnavailableRetryUsesFreshRequestBody("udp-client");
+ }
+
+ private void assertTemporaryUnavailableRetryUsesFreshRequestBody(final
String receiverKey)
+ throws Exception {
final CommonConfig commonConfig =
CommonDescriptor.getInstance().getConfig();
final long originalRetryLocalIntervalMs =
commonConfig.getPipeAirGapRetryLocalIntervalMs();
final long originalRetryMaxMs = commonConfig.getPipeAirGapRetryMaxMs();
@@ -163,6 +229,7 @@ public class IoTDBAirGapReceiverTest {
final byte[] expectedBody = new byte[] {1, 2, 3};
final AtomicInteger receiveCount = new AtomicInteger();
stubAgent.setStubReceiver(
+ receiverKey,
new IoTDBReceiver() {
@Override
public TPipeTransferResp receive(final TPipeTransferReq req) {
@@ -196,9 +263,9 @@ public class IoTDBAirGapReceiverTest {
final Method handleReq =
IoTDBAirGapReceiver.class.getDeclaredMethod(
- "handleReq", AirGapPseudoTPipeTransferRequest.class, long.class);
+ "handleReq", AirGapPseudoTPipeTransferRequest.class, long.class,
String.class);
handleReq.setAccessible(true);
- handleReq.invoke(receiver, req, System.currentTimeMillis());
+ handleReq.invoke(receiver, req, System.currentTimeMillis(), receiverKey);
Assert.assertEquals(2, receiveCount.get());
Assert.assertEquals(0, req.body.position());
@@ -239,7 +306,7 @@ public class IoTDBAirGapReceiverTest {
@Test
public void testIdleReadTimeoutDoesNotRespondOrCloseSocket() throws
Exception {
final RecordingSocket socket = new RecordingSocket(new
TimeoutInputStream(new byte[0]));
- invokeReceive(new IoTDBAirGapReceiver(socket, 4L));
+ invokeReceive(new IoTDBAirGapReceiver(socket, 5L));
Assert.assertArrayEquals(new byte[0], socket.getWrittenBytes());
Assert.assertFalse(socket.isClosed());
@@ -249,7 +316,7 @@ public class IoTDBAirGapReceiverTest {
public void testPartialRequestReadTimeoutClosesSocketWithoutResponse()
throws Exception {
final RecordingSocket socket =
new RecordingSocket(new TimeoutInputStream(new byte[] {(byte) 0xFF}));
- invokeReceive(new IoTDBAirGapReceiver(socket, 5L));
+ invokeReceive(new IoTDBAirGapReceiver(socket, 6L));
Assert.assertArrayEquals(new byte[0], socket.getWrittenBytes());
Assert.assertTrue(socket.isClosed());
@@ -387,6 +454,10 @@ public class IoTDBAirGapReceiverTest {
void setStubReceiver(final IoTDBReceiver receiver) {
setReceiverWithSpecifiedClient(null, receiver);
}
+
+ void setStubReceiver(final String key, final IoTDBReceiver receiver) {
+ setReceiverWithSpecifiedClient(key, receiver);
+ }
}
private static class StubReceiver implements IoTDBReceiver {
@@ -413,6 +484,33 @@ public class IoTDBAirGapReceiverTest {
}
}
+ private static byte[] toRawRequestBytes() throws IOException {
+ try (final ByteArrayOutputStream byteArrayOutputStream = new
ByteArrayOutputStream();
+ final DataOutputStream outputStream = new
DataOutputStream(byteArrayOutputStream)) {
+ ReadWriteIOUtils.write(IoTDBSinkRequestVersion.VERSION_1.getVersion(),
outputStream);
+ ReadWriteIOUtils.write((short) 0, outputStream);
+ return byteArrayOutputStream.toByteArray();
+ }
+ }
+
+ private static byte[] toAirGapData(final byte[] rawRequestBytes) {
+ final CRC32 crc32 = new CRC32();
+ crc32.update(rawRequestBytes, 0, rawRequestBytes.length);
+ return BytesUtils.concatByteArrayList(
+ Arrays.asList(BytesUtils.longToBytes(crc32.getValue()),
rawRequestBytes));
+ }
+
+ private static byte[] toUdpDatagram(final long requestId, final byte[]
rawRequestBytes) {
+ final byte[] data = toAirGapData(rawRequestBytes);
+ final ByteBuffer datagram =
+ ByteBuffer.allocate(Long.BYTES + 2 * Integer.BYTES + data.length)
+ .putLong(requestId)
+ .putInt(data.length)
+ .putInt(data.length)
+ .put(data);
+ return datagram.array();
+ }
+
private static class RuntimeCleanupReceiver implements IoTDBReceiver {
private final String sessionKey;
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 ca88f30fcc5..82000564cfe 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
@@ -23,6 +23,7 @@ import
org.apache.iotdb.commons.pipe.agent.plugin.builtin.BuiltinPipePlugin;
import org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant;
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.thrift.request.PipeRequestType;
import org.apache.iotdb.db.pipe.event.common.heartbeat.PipeHeartbeatEvent;
import
org.apache.iotdb.db.pipe.event.common.tablet.PipeRawTabletInsertionEvent;
@@ -42,13 +43,20 @@ import org.junit.Assert;
import org.junit.Test;
import java.io.IOException;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
+import java.net.Socket;
+import java.net.SocketTimeoutException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.zip.CRC32;
public class IoTDBDataRegionAirGapSinkTest {
@@ -140,7 +148,196 @@ public class IoTDBDataRegionAirGapSinkTest {
}
}
+ @Test
+ public void testUdpAirGapSlicing() throws Exception {
+ try (final DatagramSocket receiverSocket = new DatagramSocket(0);
+ final UdpTestingIoTDBDataRegionAirGapSink sink =
+ new UdpTestingIoTDBDataRegionAirGapSink()) {
+ final Map<String, String> attributes = buildParameterAttributes(false);
+ attributes.put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY, "udp");
+ attributes.put(PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY,
"1024");
+
+ final PipeParameters parameters = new PipeParameters(attributes);
+ sink.validate(new PipeParameterValidator(parameters));
+ sink.customize(
+ parameters,
+ new PipeTaskRuntimeConfiguration(new
PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1)));
+
+ final List<TPipeTransferReq> receivedRequests = new ArrayList<>();
+ final List<Integer> receivedPacketSizes = new ArrayList<>();
+ final AtomicBoolean senderFinished = new AtomicBoolean(false);
+ final AtomicReference<Throwable> receiverException = new
AtomicReference<>();
+ final Thread receiverThread =
+ new Thread(
+ () -> {
+ try {
+ receiverSocket.setSoTimeout(100);
+ while (!senderFinished.get()) {
+ final byte[] buffer = new byte[2048];
+ final DatagramPacket packet = new DatagramPacket(buffer,
buffer.length);
+ try {
+ receiverSocket.receive(packet);
+ receivedPacketSizes.add(packet.getLength());
+ receivedRequests.add(decodeAirGapDatagram(buffer,
packet.getLength()));
+ final ByteBuffer responseBuffer =
+ ByteBuffer.allocate(Long.BYTES +
AirGapOneByteResponse.OK.length)
+ .putLong(ByteBuffer.wrap(buffer, 0,
Long.BYTES).getLong())
+ .put(AirGapOneByteResponse.OK);
+ receiverSocket.send(
+ new DatagramPacket(
+ responseBuffer.array(),
+ responseBuffer.array().length,
+ packet.getAddress(),
+ packet.getPort()));
+ } catch (final SocketTimeoutException ignored) {
+ // Poll until the sender has finished all datagrams.
+ }
+ }
+ } catch (final Throwable ignored) {
+ receiverException.set(ignored);
+ }
+ });
+ receiverThread.start();
+
+ final byte[] request = new byte[2500];
+ request[0] = 1;
+ final byte[] type =
org.apache.tsfile.utils.BytesUtils.shortToBytes((short) 11);
+ request[1] = type[0];
+ request[2] = type[1];
+
+ try {
+ Assert.assertTrue(sink.sendUdp(receiverSocket.getLocalPort(),
request));
+ } finally {
+ senderFinished.set(true);
+ }
+
+ receiverThread.join(5000);
+ Assert.assertFalse(receiverThread.isAlive());
+ Assert.assertNull(receiverException.get());
+ Assert.assertTrue(receivedRequests.size() > 1);
+ for (final int packetSize : receivedPacketSizes) {
+ Assert.assertTrue(packetSize <= 1024);
+ }
+ for (final TPipeTransferReq receivedRequest : receivedRequests) {
+ Assert.assertEquals(PipeRequestType.TRANSFER_SLICE.getType(),
receivedRequest.type);
+ Assert.assertTrue(receivedRequest.body.remaining() < 1024);
+ }
+ }
+ }
+
+ @Test
+ public void testUdpAirGapRejectsOversizedAcknowledgement() throws Exception {
+ try (final DatagramSocket receiverSocket = new DatagramSocket(0);
+ final UdpTestingIoTDBDataRegionAirGapSink sink =
+ new UdpTestingIoTDBDataRegionAirGapSink()) {
+ final PipeParameters parameters = new
PipeParameters(buildParameterAttributes(false));
+
parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY,
"udp");
+ sink.validate(new PipeParameterValidator(parameters));
+ sink.customize(
+ parameters,
+ new PipeTaskRuntimeConfiguration(new
PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1)));
+
+ final AtomicBoolean oversizedAcknowledgementSent = new
AtomicBoolean(false);
+ final AtomicReference<Throwable> receiverException = new
AtomicReference<>();
+ final Thread receiverThread =
+ new Thread(
+ () -> {
+ try {
+ final byte[] buffer = new byte[2048];
+ final DatagramPacket packet = new DatagramPacket(buffer,
buffer.length);
+ receiverSocket.receive(packet);
+ final long requestId = ByteBuffer.wrap(buffer, 0,
Long.BYTES).getLong();
+ final byte[] oversizedAck =
+ ByteBuffer.allocate(Long.BYTES +
AirGapOneByteResponse.OK.length + 1)
+ .putLong(requestId)
+ .put(AirGapOneByteResponse.OK)
+ .put((byte) 1)
+ .array();
+ receiverSocket.send(
+ new DatagramPacket(
+ oversizedAck,
+ oversizedAck.length,
+ packet.getAddress(),
+ packet.getPort()));
+ oversizedAcknowledgementSent.set(true);
+ } catch (final Throwable t) {
+ receiverException.set(t);
+ }
+ });
+ receiverThread.start();
+
+ try {
+ Assert.assertThrows(
+ SocketTimeoutException.class,
+ () -> sink.sendUdp(receiverSocket.getLocalPort(), new byte[] {1,
0, 1}));
+ } finally {
+ receiverSocket.close();
+ }
+
+ receiverThread.join(5000);
+ Assert.assertFalse(receiverThread.isAlive());
+ Assert.assertNull(receiverException.get());
+ Assert.assertTrue(oversizedAcknowledgementSent.get());
+ }
+ }
+
+ @Test
+ public void testUdpAirGapIgnoresStaleAcknowledgement() throws Exception {
+ try (final DatagramSocket receiverSocket = new DatagramSocket(0);
+ final UdpTestingIoTDBDataRegionAirGapSink sink =
+ new UdpTestingIoTDBDataRegionAirGapSink()) {
+ final PipeParameters parameters = new
PipeParameters(buildParameterAttributes(false));
+
parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_KEY,
"udp");
+
parameters.getAttribute().put(PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_KEY,
"1024");
+ sink.validate(new PipeParameterValidator(parameters));
+ sink.customize(
+ parameters,
+ new PipeTaskRuntimeConfiguration(new
PipeTaskSinkRuntimeEnvironment("pipe", 1L, 1)));
+
+ final AtomicBoolean secondPacketReceived = new AtomicBoolean(false);
+ final AtomicReference<Throwable> receiverException = new
AtomicReference<>();
+ final Thread receiverThread =
+ new Thread(
+ () -> {
+ try {
+ final byte[] buffer = new byte[2048];
+ final DatagramPacket packet = new DatagramPacket(buffer,
buffer.length);
+ receiverSocket.receive(packet);
+ final long requestId = ByteBuffer.wrap(buffer, 0,
Long.BYTES).getLong();
+ sendUdpAck(receiverSocket, packet, requestId);
+ receiverSocket.receive(packet);
+ sendUdpAck(receiverSocket, packet, requestId);
+ secondPacketReceived.set(true);
+ } catch (final Throwable t) {
+ receiverException.set(t);
+ }
+ });
+ receiverThread.start();
+
+ try {
+ final byte[] request = new byte[2500];
+ request[0] = 1;
+ request[1] = 0;
+ request[2] = 1;
+ Assert.assertThrows(
+ SocketTimeoutException.class,
+ () -> sink.sendUdp(receiverSocket.getLocalPort(), request));
+ } finally {
+ receiverSocket.close();
+ }
+
+ receiverThread.join(5000);
+ Assert.assertFalse(receiverThread.isAlive());
+ Assert.assertNull(receiverException.get());
+ Assert.assertTrue(secondPacketReceived.get());
+ }
+ }
+
private PipeParameters buildParameters(final boolean useTsFileBatch) {
+ return new PipeParameters(buildParameterAttributes(useTsFileBatch));
+ }
+
+ private Map<String, String> buildParameterAttributes(final boolean
useTsFileBatch) {
final Map<String, String> attributes = new HashMap<>();
attributes.put(
PipeSinkConstant.CONNECTOR_KEY,
@@ -151,7 +348,7 @@ public class IoTDBDataRegionAirGapSinkTest {
if (useTsFileBatch) {
attributes.put(PipeSinkConstant.CONNECTOR_FORMAT_KEY, "tsfile");
}
- return new PipeParameters(attributes);
+ return attributes;
}
private PipeRawTabletInsertionEvent createPipeRawTabletInsertionEvent(
@@ -185,6 +382,34 @@ public class IoTDBDataRegionAirGapSinkTest {
return req;
}
+ private static TPipeTransferReq decodeAirGapDatagram(final byte[] datagram,
final int length) {
+ final ByteBuffer buffer = ByteBuffer.wrap(datagram, Long.BYTES, length -
Long.BYTES);
+ final int payloadLength = ReadWriteIOUtils.readInt(buffer);
+ Assert.assertEquals(payloadLength, ReadWriteIOUtils.readInt(buffer));
+ Assert.assertEquals(length - Long.BYTES - 2 * Integer.BYTES,
payloadLength);
+
+ final long expectedChecksum = ReadWriteIOUtils.readLong(buffer);
+ final byte[] payload = new byte[payloadLength - Long.BYTES];
+ buffer.get(payload);
+
+ final CRC32 crc32 = new CRC32();
+ crc32.update(payload, 0, payload.length);
+ Assert.assertEquals(expectedChecksum, crc32.getValue());
+ return toTPipeTransferReq(payload);
+ }
+
+ private static void sendUdpAck(
+ final DatagramSocket socket, final DatagramPacket requestPacket, final
long requestId)
+ throws IOException {
+ final byte[] ack =
+ ByteBuffer.allocate(Long.BYTES + AirGapOneByteResponse.OK.length)
+ .putLong(requestId)
+ .put(AirGapOneByteResponse.OK)
+ .array();
+ socket.send(
+ new DatagramPacket(ack, ack.length, requestPacket.getAddress(),
requestPacket.getPort()));
+ }
+
private static class RecordingIoTDBDataRegionAirGapSink extends
IoTDBDataRegionAirGapSink {
private final List<byte[]> sentRequests = new ArrayList<>();
@@ -227,4 +452,15 @@ public class IoTDBDataRegionAirGapSinkTest {
}
}
}
+
+ private static class UdpTestingIoTDBDataRegionAirGapSink extends
IoTDBDataRegionAirGapSink {
+
+ 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) {
+ socket.connectUdp(1000);
+ return sendBytes(socket, request);
+ }
+ }
+ }
}
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 0aafc358427..c09c42a188f 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
@@ -508,6 +508,14 @@ public final class PipeMessages {
"IoTDBAirGapConnector is customized with handshakeTimeoutMs: {}.";
public static final String AIR_GAP_CUSTOMIZED_E_LANGUAGE =
"IoTDBAirGapConnector is customized with eLanguageEnable: {}.";
+ public static final String
+
EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654
=
+ "Air gap UDP datagram size %d exceeds configured packet size %d.";
+ public static final String
EXCEPTION_TIMED_OUT_WAITING_FOR_THE_MATCHING_AIR_GAP_ACK_DD2F3706 =
+ "Timed out waiting for the matching air-gap ACK.";
+ public static final String
+
EXCEPTION_AIR_GAP_UDP_PACKET_SIZE_ARG_IS_TOO_SMALL_TO_TRANSFER_SLICED_REQUESTS_722B32AB
=
+ "Air gap UDP packet size %d is too small to transfer sliced
requests.";
public static final String CONNECTED_TO_TARGET_SERVER =
"Successfully connected to target server ip: {}, port: {}.";
public static final String FAILED_TO_CLOSE_SOCKET =
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 9746b71500d..b8fd1db6729 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
@@ -486,6 +486,14 @@ public final class PipeMessages {
public static final String AIR_GAP_CUSTOMIZED_HANDSHAKE_TIMEOUT =
"IoTDBAirGapConnector 自定义握手超时时间：{}。";
+ public static final String
+
EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654
=
+ "Air gap UDP 数据报大小 %d 超过配置的包大小 %d。";
+ public static final String
EXCEPTION_TIMED_OUT_WAITING_FOR_THE_MATCHING_AIR_GAP_ACK_DD2F3706 =
+ "等待匹配的 air-gap ACK 超时。";
+ public static final String
+
EXCEPTION_AIR_GAP_UDP_PACKET_SIZE_ARG_IS_TOO_SMALL_TO_TRANSFER_SLICED_REQUESTS_722B32AB
=
+ "Air gap UDP 包大小 %d 太小，无法传输切片请求。";
public static final String AIR_GAP_CUSTOMIZED_E_LANGUAGE =
"IoTDBAirGapConnector 自定义 eLanguageEnable：{}。";
public static final String CONNECTED_TO_TARGET_SERVER =
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 04a744611c0..e7a36577a06 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
@@ -186,6 +186,26 @@ public class PipeSinkConstant {
"sink.air-gap.handshake-timeout-ms";
public static final int CONNECTOR_AIR_GAP_HANDSHAKE_TIMEOUT_MS_DEFAULT_VALUE
= 5000;
+ public static final String CONNECTOR_AIR_GAP_TRANSPORT_KEY =
"connector.air-gap.transport";
+ public static final String SINK_AIR_GAP_TRANSPORT_KEY =
"sink.air-gap.transport";
+ public static final String CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE = "tcp";
+ public static final String CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE = "udp";
+ public static final String CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE =
+ CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE;
+ public static final Set<String> CONNECTOR_AIR_GAP_TRANSPORT_SET =
+ Collections.unmodifiableSet(
+ new HashSet<>(
+ Arrays.asList(
+ CONNECTOR_AIR_GAP_TRANSPORT_TCP_VALUE,
CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE)));
+
+ 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 =
+ "sink.air-gap.udp.packet-size-bytes";
+ 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_IOTDB_SYNC_CONNECTOR_VERSION_KEY =
"connector.version";
public static final String SINK_IOTDB_SYNC_CONNECTOR_VERSION_KEY =
"sink.version";
public static final String
CONNECTOR_IOTDB_SYNC_CONNECTOR_VERSION_DEFAULT_VALUE = "1.1";
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 6919b3b87ce..bdffcfe4734 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,16 +26,20 @@ import org.apache.iotdb.commons.pipe.config.PipeConfig;
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.thrift.common.PipeTransferSliceReqBuilder;
import org.apache.iotdb.commons.utils.ErrorHandlingCommonUtils;
import org.apache.iotdb.pipe.api.annotation.TableModel;
import org.apache.iotdb.pipe.api.annotation.TreeModel;
import
org.apache.iotdb.pipe.api.customizer.configuration.PipeConnectorRuntimeConfiguration;
+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.PipeException;
import org.apache.iotdb.rpc.TSStatusCode;
+import org.apache.iotdb.service.rpc.thrift.TPipeTransferReq;
import org.apache.tsfile.utils.BytesUtils;
+import org.apache.tsfile.utils.ReadWriteIOUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -44,34 +48,59 @@ import java.io.DataInputStream;
import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
+import java.net.DatagramPacket;
+import java.net.DatagramSocket;
+import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketException;
+import java.net.SocketTimeoutException;
+import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.zip.CRC32;
import static
org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_DEFAULT_VALUE;
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_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;
+import static
org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE;
+import static
org.apache.iotdb.commons.pipe.config.constant.PipeSinkConstant.CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_DEFAULT_VALUE;
+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_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;
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_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;
@TreeModel
@TableModel
public abstract class IoTDBAirGapSink extends IoTDBSink {
+ private static final int UDP_ENVELOPE_SIZE = 2 * Integer.BYTES + Long.BYTES;
+ private static final int UDP_REQUEST_ID_SIZE = Long.BYTES;
+ private static final int UDP_ACK_SIZE = UDP_REQUEST_ID_SIZE + 1;
+ private static final int UDP_SLICE_REQUEST_SERIALIZATION_RESERVED_SIZE = 64;
+
protected static class AirGapSocket extends Socket {
private final TEndPoint endPoint;
+ private DatagramSocket datagramSocket;
+ private InetAddress datagramAddress;
+ private boolean isUdp;
public AirGapSocket(final String ip, final int port) {
this.endPoint = new TEndPoint(ip, port);
@@ -81,9 +110,60 @@ public abstract class IoTDBAirGapSink extends IoTDBSink {
return endPoint;
}
+ public void connectUdp(final int timeoutMs) throws IOException {
+ datagramAddress = InetAddress.getByName(endPoint.getIp());
+ datagramSocket = new DatagramSocket();
+ datagramSocket.connect(datagramAddress, endPoint.getPort());
+ datagramSocket.setSoTimeout(timeoutMs);
+ isUdp = true;
+ }
+
+ public InetAddress getDatagramAddress() {
+ return datagramAddress;
+ }
+
+ public DatagramSocket getDatagramSocket() {
+ return datagramSocket;
+ }
+
+ public boolean isUdp() {
+ return isUdp;
+ }
+
+ @Override
+ public boolean isConnected() {
+ return isUdp
+ ? datagramSocket != null && datagramSocket.isConnected() &&
!datagramSocket.isClosed()
+ : super.isConnected();
+ }
+
+ @Override
+ public synchronized void setSoTimeout(final int timeout) throws
SocketException {
+ if (isUdp && datagramSocket != null) {
+ datagramSocket.setSoTimeout(timeout);
+ } else {
+ super.setSoTimeout(timeout);
+ }
+ }
+
+ @Override
+ public synchronized void close() throws IOException {
+ if (datagramSocket != null) {
+ datagramSocket.close();
+ }
+ super.close();
+ }
+
@Override
public String toString() {
- return "AirGapSocket{" + "endPoint=" + endPoint + "} (" +
super.toString() + ")";
+ return "AirGapSocket{"
+ + "endPoint="
+ + endPoint
+ + ", transport="
+ + (isUdp ? "udp" : "tcp")
+ + "} ("
+ + super.toString()
+ + ")";
}
}
@@ -101,11 +181,49 @@ public abstract class IoTDBAirGapSink extends IoTDBSink {
private int handshakeTimeoutMs;
private boolean eLanguageEnable;
+ private boolean useUdpTransport;
+ private int udpPacketSizeInBytes;
// The air gap connector does not use clientManager thus we put handshake
type here
protected boolean supportModsIfIsDataNodeReceiver = true;
private final Map<TEndPoint, Long> failLogTimes = new HashMap<>();
+ private final AtomicLong udpRequestIdGenerator = new AtomicLong();
+
+ @Override
+ public void validate(final PipeParameterValidator validator) throws
Exception {
+ super.validate(validator);
+
+ final PipeParameters parameters = validator.getParameters();
+ final String airGapTransport =
+ parameters
+ .getStringOrDefault(
+ Arrays.asList(CONNECTOR_AIR_GAP_TRANSPORT_KEY,
SINK_AIR_GAP_TRANSPORT_KEY),
+ CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE)
+ .trim()
+ .toLowerCase();
+ validator.validate(
+ arg -> CONNECTOR_AIR_GAP_TRANSPORT_SET.contains(airGapTransport),
+ String.format(
+ "Air gap transport should be one of %s, but got %s.",
+ CONNECTOR_AIR_GAP_TRANSPORT_SET, airGapTransport),
+ airGapTransport);
+
+ final int packetSize =
+ 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);
+ validator.validate(
+ arg ->
+ packetSize >= CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE
+ && packetSize <= CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE,
+ String.format(
+ "UDP packet size should be in the range [%d, %d], but got %d.",
+ CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MIN_VALUE,
+ CONNECTOR_AIR_GAP_UDP_PACKET_SIZE_MAX_VALUE,
+ packetSize),
+ packetSize);
+ }
@Override
public void customize(
@@ -146,6 +264,23 @@ public abstract class IoTDBAirGapSink extends IoTDBSink {
CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_KEY,
SINK_AIR_GAP_E_LANGUAGE_ENABLE_KEY),
CONNECTOR_AIR_GAP_E_LANGUAGE_ENABLE_DEFAULT_VALUE);
LOGGER.info(PipeMessages.AIR_GAP_CUSTOMIZED_E_LANGUAGE, eLanguageEnable);
+
+ useUdpTransport =
+ CONNECTOR_AIR_GAP_TRANSPORT_UDP_VALUE.equals(
+ parameters
+ .getStringOrDefault(
+ Arrays.asList(CONNECTOR_AIR_GAP_TRANSPORT_KEY,
SINK_AIR_GAP_TRANSPORT_KEY),
+ CONNECTOR_AIR_GAP_TRANSPORT_DEFAULT_VALUE)
+ .trim()
+ .toLowerCase());
+ udpPacketSizeInBytes =
+ 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);
+ LOGGER.info(
+ "Air gap transport is {}, udp packet size is {} bytes.",
+ useUdpTransport ? "udp" : "tcp",
+ udpPacketSizeInBytes);
}
@Override
@@ -180,8 +315,12 @@ public abstract class IoTDBAirGapSink extends IoTDBSink {
final AirGapSocket socket = new AirGapSocket(ip, port);
try {
- socket.connect(new InetSocketAddress(ip, port), handshakeTimeoutMs);
- socket.setKeepAlive(true);
+ if (useUdpTransport) {
+ socket.connectUdp(handshakeTimeoutMs);
+ } else {
+ socket.connect(new InetSocketAddress(ip, port), handshakeTimeoutMs);
+ socket.setKeepAlive(true);
+ }
sockets.set(i, socket);
LOGGER.debug(PipeMessages.CONNECTED_TO_TARGET_SERVER, ip, port);
failLogTimes.remove(nodeUrls.get(i));
@@ -324,6 +463,10 @@ public abstract class IoTDBAirGapSink extends IoTDBSink {
PipeMessages.EXCEPTION_SOCKET_ARG_CLOSED_WILL_TRY_HANDSHAKE_02562BF1, socket));
}
+ if (socket.isUdp()) {
+ return sendBytesByUdp(socket, bytes);
+ }
+
final BufferedOutputStream outputStream = new
BufferedOutputStream(socket.getOutputStream());
bytes = enrichWithLengthAndChecksum(bytes);
outputStream.write(eLanguageEnable ? enrichWithELanguage(bytes) : bytes);
@@ -334,6 +477,147 @@ public abstract class IoTDBAirGapSink extends IoTDBSink {
return Arrays.equals(AirGapOneByteResponse.OK, response);
}
+ private boolean sendBytesByUdp(final AirGapSocket socket, final byte[]
bytes) throws IOException {
+ for (final TPipeTransferReq request : sliceIfNeededForUdp(bytes)) {
+ if (!sendOneDatagram(socket, request)) {
+ return false;
+ }
+ }
+ return true;
+ }
+
+ private boolean sendOneDatagram(final AirGapSocket socket, final
TPipeTransferReq request)
+ throws IOException {
+ final long requestId = udpRequestIdGenerator.incrementAndGet();
+ final ByteBuffer body = request.body.duplicate();
+ final int requestSize = Byte.BYTES + Short.BYTES + body.remaining();
+ final int eLanguageSize =
+ eLanguageEnable
+ ? AirGapELanguageConstant.E_LANGUAGE_PREFIX.length
+ + AirGapELanguageConstant.E_LANGUAGE_SUFFIX.length
+ : 0;
+ final int datagramSize = UDP_REQUEST_ID_SIZE + UDP_ENVELOPE_SIZE +
requestSize + eLanguageSize;
+ if (datagramSize > udpPacketSizeInBytes) {
+ throw new IOException(
+ String.format(
+ PipeMessages
+
.EXCEPTION_AIR_GAP_UDP_DATAGRAM_SIZE_ARG_EXCEEDS_CONFIGURED_PACKET_SIZE_ARG_559D0654,
+ datagramSize,
+ udpPacketSizeInBytes));
+ }
+
+ final CRC32 crc32 = new CRC32();
+ crc32.update(request.version);
+ crc32.update((request.type >>> Byte.SIZE) & 0xFF);
+ crc32.update(request.type & 0xFF);
+ crc32.update(body.duplicate());
+
+ final ByteBuffer datagramBuffer =
ByteBuffer.allocate(datagramSize).putLong(requestId);
+ if (eLanguageEnable) {
+ datagramBuffer.put(AirGapELanguageConstant.E_LANGUAGE_PREFIX);
+ }
+ datagramBuffer
+ .putInt(requestSize + Long.BYTES)
+ .putInt(requestSize + Long.BYTES)
+ .putLong(crc32.getValue())
+ .put(request.version)
+ .putShort(request.type)
+ .put(body);
+ if (eLanguageEnable) {
+ datagramBuffer.put(AirGapELanguageConstant.E_LANGUAGE_SUFFIX);
+ }
+ final byte[] datagramBytes = datagramBuffer.array();
+
+ final DatagramSocket datagramSocket = socket.getDatagramSocket();
+ datagramSocket.send(
+ new DatagramPacket(
+ datagramBytes,
+ datagramBytes.length,
+ socket.getDatagramAddress(),
+ socket.getEndPoint().getPort()));
+
+ final int originalTimeoutMs = datagramSocket.getSoTimeout();
+ final long deadlineNanos =
+ originalTimeoutMs == 0
+ ? Long.MAX_VALUE
+ : System.nanoTime() +
TimeUnit.MILLISECONDS.toNanos(originalTimeoutMs);
+ // One extra byte lets us detect datagrams that receive() would otherwise
silently truncate.
+ final byte[] response = new byte[UDP_ACK_SIZE + 1];
+ final DatagramPacket responsePacket = new DatagramPacket(response,
response.length);
+ try {
+ while (true) {
+ if (deadlineNanos != Long.MAX_VALUE) {
+ final long remainingNanos = deadlineNanos - System.nanoTime();
+ if (remainingNanos <= 0) {
+ throw new SocketTimeoutException(
+
PipeMessages.EXCEPTION_TIMED_OUT_WAITING_FOR_THE_MATCHING_AIR_GAP_ACK_DD2F3706);
+ }
+ datagramSocket.setSoTimeout(
+ (int) Math.max(1,
TimeUnit.NANOSECONDS.toMillis(remainingNanos)));
+ }
+
+ responsePacket.setLength(response.length);
+ datagramSocket.receive(responsePacket);
+ if (responsePacket.getLength() != UDP_ACK_SIZE) {
+ continue;
+ }
+
+ final ByteBuffer responseBuffer = ByteBuffer.wrap(response, 0,
responsePacket.getLength());
+ if (responseBuffer.getLong() != requestId) {
+ continue;
+ }
+ return responseBuffer.get() == AirGapOneByteResponse.OK[0];
+ }
+ } finally {
+ datagramSocket.setSoTimeout(originalTimeoutMs);
+ }
+ }
+
+ private List<TPipeTransferReq> sliceIfNeededForUdp(final byte[]
requestBytes) throws IOException {
+ final TPipeTransferReq request = toTPipeTransferReq(requestBytes);
+ final int rawPayloadSizeLimit =
+ udpPacketSizeInBytes
+ - UDP_ENVELOPE_SIZE
+ - UDP_REQUEST_ID_SIZE
+ - (eLanguageEnable
+ ? AirGapELanguageConstant.E_LANGUAGE_PREFIX.length
+ + AirGapELanguageConstant.E_LANGUAGE_SUFFIX.length
+ : 0);
+ if (Byte.BYTES + Short.BYTES + request.body.remaining() <=
rawPayloadSizeLimit) {
+ return Arrays.asList(request);
+ }
+
+ final int sliceBodySizeLimit =
+ rawPayloadSizeLimit - UDP_SLICE_REQUEST_SERIALIZATION_RESERVED_SIZE;
+ if (sliceBodySizeLimit <= 0) {
+ throw new IOException(
+ String.format(
+ PipeMessages
+
.EXCEPTION_AIR_GAP_UDP_PACKET_SIZE_ARG_IS_TOO_SMALL_TO_TRANSFER_SLICED_REQUESTS_722B32AB,
+ udpPacketSizeInBytes));
+ }
+
+ final int sliceOrderId = PipeTransferSliceReqBuilder.nextSliceOrderId();
+ final int sliceCount = PipeTransferSliceReqBuilder.getSliceCount(request,
sliceBodySizeLimit);
+
+ final List<TPipeTransferReq> slicedRequests = new ArrayList<>(sliceCount);
+ for (int i = 0; i < sliceCount; i++) {
+ slicedRequests.add(
+ PipeTransferSliceReqBuilder.buildSliceReq(
+ request, sliceOrderId, i, sliceCount, sliceBodySizeLimit));
+ }
+ return slicedRequests;
+ }
+
+ private TPipeTransferReq toTPipeTransferReq(final byte[] requestBytes) {
+ final ByteBuffer byteBuffer = ByteBuffer.wrap(requestBytes);
+ final TPipeTransferReq request = new TPipeTransferReq();
+ request.version = ReadWriteIOUtils.readByte(byteBuffer);
+ request.type = ReadWriteIOUtils.readShort(byteBuffer);
+ request.body = byteBuffer.slice();
+ return request;
+ }
+
protected boolean send(final AirGapSocket socket, final byte[] bytes) throws
IOException {
return send(null, 0, socket, bytes);
}
