Implement UDP transport for air-gap sink

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,9 @@ public class IoTDBAirGapReceiver extends WrappedRunnable {
   private final IoTDBDataNodeReceiverAgent agent;
 
   private boolean isELanguagePayload;
+  private OutputStream currentOutputStream;
+  private DatagramSocket currentDatagramSocket;
+  private SocketAddress currentRemoteSocketAddress;
 
   public IoTDBAirGapReceiver(final Socket socket, final long receiverId) {
     this.socket = socket;
@@ -105,33 +112,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 +146,78 @@ private void receive() throws IOException {
     }
   }
 
-  private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long startTime)
+  boolean receive(final byte[] data, final String receiverKey) throws IOException {
+    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;
+      }
+
+      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;
+    }
+  }
+
+  void receiveUdp(
+      final DatagramSocket datagramSocket,
+      final DatagramPacket packet,
+      final String receiverKey,
+      final byte[] buffer)
+      throws IOException {
+    isELanguagePayload = false;
+    currentDatagramSocket = datagramSocket;
+    currentRemoteSocketAddress = packet.getSocketAddress();
+    boolean requestHandlingStarted = false;
+    try (final ByteArrayInputStream inputStream =
+        new ByteArrayInputStream(buffer, 0, packet.getLength())) {
+      final byte[] data = readData(inputStream);
+      requestHandlingStarted = true;
+      receive(data, receiverKey);
+    } catch (final Exception e) {
+      if (!requestHandlingStarted) {
+        fail();
+      }
+      throw e;
+    } finally {
+      currentOutputStream = null;
+      currentDatagramSocket = null;
+      currentRemoteSocketAddress = 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 = agent.receive(req);
+    final TPipeTransferResp resp =
+        receiverKey == null ? agent.receive(req) : agent.receive(receiverKey, req);
 
     final TSStatus status = resp.getStatus();
     if (status.getCode() == TSStatusCode.SUCCESS_STATUS.getStatusCode()) {
@@ -181,7 +235,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);
@@ -198,14 +252,23 @@ private void handleReq(final AirGapPseudoTPipeTransferRequest req, final long st
   }
 
   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) {
+      currentDatagramSocket.send(
+          new DatagramPacket(response, response.length, currentRemoteSocketAddress));
+      return;
+    }
+
+    final OutputStream outputStream =
+        currentOutputStream != null ? currentOutputStream : socket.getOutputStream();
+    outputStream.write(response);
     outputStream.flush();
   }
 

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

@@ -26,30 +26,42 @@
 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.util.Map;
 import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
 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 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);
+  private final Map<String, ClientSession> udpClientSessions = new ConcurrentHashMap<>();
 
   public void listen() {
     try {
@@ -61,39 +73,105 @@ public void listen() {
           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()) {
       listenExecutor.submit(this::listen);
     }
   }
 
+  public void listenUdp() {
+    while (allowSubmitListen.get() && !datagramSocket.isClosed()) {
+      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 = registerUdpSessionIfNecessary(packet);
+        try {
+          receiver.receiveUdp(datagramSocket, packet, receiverKey, buffer);
+        } finally {
+          if (registeredSession) {
+            SessionManager.getInstance().removeCurrSession();
+          }
+        }
+      } 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);
+      }
+    }
+  }
+
+  private boolean registerUdpSessionIfNecessary(final DatagramPacket packet) {
+    final String receiverKey = packet.getSocketAddress().toString();
+    final ClientSession session =
+        udpClientSessions.computeIfAbsent(
+            receiverKey,
+            key ->
+                new ClientSession(new DatagramClientSocket(packet.getAddress(), packet.getPort())));
+    return SessionManager.getInstance().registerSession(session);
+  }
+
   @Override
   public void start() throws StartupException {
     try {
       serverSocket = new ServerSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort());
+      datagramSocket = new DatagramSocket(PipeConfig.getInstance().getPipeAirGapReceiverPort());
     } 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(
+        (key, session) -> {
+          final boolean registeredSession = SessionManager.getInstance().registerSession(session);
+          try {
+            PipeDataNodeAgent.receiver().thrift().handleClientExit(key);
+          } finally {
+            if (registeredSession) {
+              SessionManager.getInstance().removeCurrSession();
+            }
+          }
+        });
+    udpClientSessions.clear();
     listenExecutor.shutdown();
 
     LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STOPPED, serverSocket);
@@ -103,4 +181,25 @@ public void stop() {
   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;
+    }
+  }
 }