Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -44,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;
Expand Down Expand Up @@ -232,6 +234,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());
Expand All @@ -258,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()) {
Expand All @@ -281,6 +289,7 @@ private AirGapPseudoTPipeTransferRequest duplicateReq(
final AirGapPseudoTPipeTransferRequest req) {
return (AirGapPseudoTPipeTransferRequest)
new AirGapPseudoTPipeTransferRequest()
.setStrictUnidirectional(req.isStrictUnidirectional())
.setVersion(req.getVersion())
.setType(req.getType())
.setBody(req.body.duplicate());
Expand All @@ -296,7 +305,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 =
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -54,12 +58,17 @@
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;

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The new property only controls the kernel socket receive buffer. The copied application queue remains fixed at 4,096 packets and 32 MiB, and the sender UDP packet-size and request-slice settings do not configure these limits. Could we expose the queue bounds through receiver configuration, or explicitly confirm that they are intentionally fixed? As written, receive buffering cannot be tuned end to end.


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<DatagramPacket> udpRequests =
new LinkedBlockingQueue<>(UDP_REQUEST_QUEUE_CAPACITY);
private final Semaphore udpRequestBytes = new Semaphore(UDP_REQUEST_QUEUE_BYTES);

private ServerSocket serverSocket;
private DatagramSocket datagramSocket;
Expand Down Expand Up @@ -88,38 +97,95 @@
}

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 {

Check warning on line 148 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/airgap/IoTDBAirGapReceiverAgent.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Extract this nested try block into a separate method.

See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGafvJIjxf_orAgsFR-&open=AaGafvJIjxf_orAgsFR-&pullRequest=19063
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();
}
}
}
Expand Down Expand Up @@ -196,6 +262,8 @@
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)) {
Expand All @@ -211,6 +279,7 @@
allowSubmitListen.set(true);
listenExecutor.submit(this::listen);
listenExecutor.submit(this::listenUdp);
listenExecutor.submit(this::processUdpRequests);

LOGGER.info(DataNodePipeMessages.IOTDBAIRGAPRECEIVERAGENT_STARTED, serverSocket);
}
Expand All @@ -230,8 +299,7 @@
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);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -196,6 +196,9 @@

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();
Expand Down Expand Up @@ -227,12 +230,48 @@

@Override
public synchronized TPipeTransferResp receive(final TPipeTransferReq req) {
final boolean isUnidirectionalHandshake =
req instanceof AirGapPseudoTPipeTransferRequest

Check warning on line 234 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Replace this instanceof check and cast with 'instanceof AirGapPseudoTPipeTransferRequest airgappseudotpipetransferrequest'

See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGafvEpjxf_orAgsFR5&open=AaGafvEpjxf_orAgsFR5&pullRequest=19063
&& ((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) {

Check failure on line 266 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Refactor this method to reduce its Cognitive Complexity from 17 to the 15 allowed.

See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGafvEpjxf_orAgsFR6&open=AaGafvEpjxf_orAgsFR6&pullRequest=19063
try {
final long startTime = System.nanoTime();
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);
Expand Down Expand Up @@ -464,7 +503,9 @@
case TRANSFER_SLICE:
{
try {
return handleTransferSlice(PipeTransferSliceReq.fromTPipeTransferReq(req));
return handleTransferSlice(

Check warning on line 506 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

'try' child has incorrect indentation level 16, expected level should be 14.

See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGafvEpjxf_orAgsFR8&open=AaGafvEpjxf_orAgsFR8&pullRequest=19063
PipeTransferSliceReq.fromTPipeTransferReq(req),
isStrictUnidirectionalRequest(req));
} finally {
PipeDataNodeReceiverMetrics.getInstance()
.recordTransferSliceTimer(System.nanoTime() - startTime);
Expand All @@ -478,7 +519,9 @@
PipeTransferCompressedReq.getMaxAdditionalDecompressedLengthInBytes(req);
try (final PipeMemoryBlock ignored =
tryAllocateReceiverMemory(requestedMemorySizeInBytes)) {
return receive(PipeTransferCompressedReq.fromTPipeTransferReq(req));
return receiveNestedRequest(

Check warning on line 522 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

'try' child has incorrect indentation level 18, expected level should be 16.

See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGafvEpjxf_orAgsFR9&open=AaGafvEpjxf_orAgsFR9&pullRequest=19063
PipeTransferCompressedReq.fromTPipeTransferReq(req),
isStrictUnidirectionalRequest(req));
}
} catch (final PipeRuntimeOutOfMemoryCriticalException e) {
return new TPipeTransferResp(
Expand Down Expand Up @@ -1098,7 +1141,25 @@
return configReceiverId.get();
}

private TPipeTransferResp handleTransferSlice(final PipeTransferSliceReq pipeTransferSliceReq) {
private boolean isStrictUnidirectionalRequest(final TPipeTransferReq req) {
return req instanceof AirGapPseudoTPipeTransferRequest

Check warning on line 1145 in iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/receiver/protocol/thrift/IoTDBDataNodeReceiver.java

View check run for this annotation

SonarQubeCloud / SonarCloud Code Analysis

Replace this instanceof check and cast with 'instanceof AirGapPseudoTPipeTransferRequest airgappseudotpipetransferrequest'

See more on https://sonarcloud.io/project/issues?id=apache_iotdb&issues=AaGafvEpjxf_orAgsFR7&open=AaGafvEpjxf_orAgsFR7&pullRequest=19063
&& ((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";
Expand Down Expand Up @@ -1137,7 +1198,7 @@
"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);
Expand Down Expand Up @@ -1938,6 +1999,8 @@

@Override
public synchronized void handleExit() {
lastSuccessfulUnidirectionalHandshake = null;
lastSuccessfulUnidirectionalHandshakeResp = null;
PipeReceiverRuntimeRegistry.getInstance()
.deregister(configPipeReceiverRuntimeSessionKey.getAndSet(null));
clearSliceReqHandler();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
Loading
Loading