Implement UDP transport for air-gap sink

iotdb-core/datanode/src/main/i18n/en/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -1367,6 +1367,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 {}.";

iotdb-core/datanode/src/main/i18n/zh/org/apache/iotdb/db/i18n/DataNodePipeMessages.java

@@ -1282,6 +1282,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 数据传输错误，{}";

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiver.java

@@ -40,11 +40,15 @@
 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 @@ private void receive() throws IOException {
 
     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 @@ private void receive() throws IOException {
     }
   }
 
-  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 @@ private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long st
       }
       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 @@ private AirGapPseudoTPipeTransferRequest duplicateReq(
   }
 
   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();
   }