This is an automated email from the ASF dual-hosted git repository.
exceptionfactory pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new c45c77cacb5 NIFI-16084 Fixed Load-balanced connection desynchronizing
Socket after in-flight transaction unregistered (#11402)
c45c77cacb5 is described below
commit c45c77cacb57041978938170ea4357392b939295
Author: Pierre Villard <[email protected]>
AuthorDate: Mon Jul 13 22:07:25 2026 +0200
NIFI-16084 Fixed Load-balanced connection desynchronizing Socket after
in-flight transaction unregistered (#11402)
Signed-off-by: David Handermann <[email protected]>
---
.../async/nio/NioAsyncLoadBalanceClient.java | 6 +-
.../client/async/nio/TestLoadBalanceSession.java | 43 ++++++++
.../async/nio/TestNioAsyncLoadBalanceClient.java | 110 +++++++++++++++++++++
.../server/TestStandardLoadBalanceProtocol.java | 42 ++++++++
4 files changed, 200 insertions(+), 1 deletion(-)
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/client/async/nio/NioAsyncLoadBalanceClient.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/client/async/nio/NioAsyncLoadBalanceClient.java
index 9248cc2d4c7..376dcac1e34 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/client/async/nio/NioAsyncLoadBalanceClient.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/main/java/org/apache/nifi/controller/queue/clustered/client/async/nio/NioAsyncLoadBalanceClient.java
@@ -138,6 +138,10 @@ public class NioAsyncLoadBalanceClient implements
AsyncLoadBalanceClient {
logger.debug("{} Triggering failure callback for {} FlowFiles
for Registered Partition {} because partition was unregistered", this,
flowFilesSent.size(), removedPartition);
removedPartition.getFailureCallback().onTransactionFailed(flowFilesSent,
TransactionFailureCallback.TransactionPhase.SENDING);
+
+ // The transaction was abandoned mid-stream. Close the
connection so the socket is not reused for the next
+ // transaction; a reused socket would leave the peer reading a
stray protocol byte and aborting the transfer.
+ close();
}
}
}
@@ -427,7 +431,7 @@ public class NioAsyncLoadBalanceClient implements
AsyncLoadBalanceClient {
return new SimpleLimitThreshold(1000, 10_000_000L);
}
- private synchronized boolean isConnectionEstablished() {
+ synchronized boolean isConnectionEstablished() {
return selector != null && channel != null && channel.isConnected();
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestLoadBalanceSession.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestLoadBalanceSession.java
index de5eea23023..267e0595cd3 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestLoadBalanceSession.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestLoadBalanceSession.java
@@ -189,6 +189,49 @@ public class TestLoadBalanceSession {
assertEquals(Arrays.asList(flowFile1, flowFile2),
transaction.getAndPurgeFlowFilesSent());
}
+ @Test
+ @Timeout(30)
+ public void
testSessionCancelLeavesChannelOpenSoReusedStreamRepeatsVersionByte() throws
InterruptedException, IOException {
+ final SocketChannel socketChannel = SocketChannel.open(new
InetSocketAddress("localhost", port));
+ socketChannel.configureBlocking(false);
+ final PeerChannel peerChannel = new PeerChannel(socketChannel, null,
"unit-test");
+
+ final RegisteredPartition partition1 = new
RegisteredPartition("unit-test-connection", () -> false,
+ () -> new MockFlowFileRecord(5), NOP_FAILURE_CALLBACK, (ff,
nodeId) -> { }, () -> LoadBalanceCompression.DO_NOT_COMPRESS, () -> true);
+ final FlowFileContentAccess contentAccess = ff -> new
ByteArrayInputStream("hello".getBytes());
+
+ final LoadBalanceSession session1 = new LoadBalanceSession(partition1,
contentAccess, new StandardLoadBalanceFlowFileCodec(), peerChannel, 30000,
+ new SimpleLimitThreshold(100, 10_000_000));
+
+ // Advance session1 until it has written its opening protocol-version
byte.
+ while (received.size() < 1) {
+ session1.communicate();
+ Thread.sleep(10L);
+ }
+
+ // A LoadBalanceSession cancel neither closes the channel nor sends
ABORT_TRANSACTION; closing a channel whose
+ // transaction was abandoned is the caller's responsibility
(NioAsyncLoadBalanceClient#unregister).
+ assertTrue(session1.cancel());
+ assertTrue(socketChannel.isOpen());
+
+ // If the same channel is reused for the next transaction, it again
writes protocol-version byte 1.
+ final LoadBalanceSession session2 = new LoadBalanceSession(partition1,
contentAccess, new StandardLoadBalanceFlowFileCodec(), peerChannel, 30000,
+ new SimpleLimitThreshold(100, 10_000_000));
+ while (session2.communicate()) {
+ }
+
+ // Two protocol-version bytes then sit back-to-back on the stream: the
abandoned transaction's, then the reused
+ // channel's. The second lands where the peer expects a continuation
of the first transaction (the desync).
+ while (received.size() < 2) {
+ Thread.sleep(10L);
+ }
+ final byte[] sent = received.toByteArray();
+ assertEquals(1, sent[0]);
+ assertEquals(1, sent[1]);
+
+ socketChannel.close();
+ }
+
@Test
@Timeout(10)
public void testLargeContent() throws InterruptedException, IOException {
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestNioAsyncLoadBalanceClient.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestNioAsyncLoadBalanceClient.java
new file mode 100644
index 00000000000..e676c448df8
--- /dev/null
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestNioAsyncLoadBalanceClient.java
@@ -0,0 +1,110 @@
+/*
+ * 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.nifi.controller.queue.clustered.client.async.nio;
+
+import org.apache.nifi.cluster.coordination.ClusterCoordinator;
+import org.apache.nifi.cluster.coordination.node.NodeConnectionState;
+import org.apache.nifi.cluster.coordination.node.NodeConnectionStatus;
+import org.apache.nifi.cluster.protocol.NodeIdentifier;
+import org.apache.nifi.controller.MockFlowFileRecord;
+import org.apache.nifi.controller.queue.LoadBalanceCompression;
+import org.apache.nifi.controller.queue.clustered.FlowFileContentAccess;
+import
org.apache.nifi.controller.queue.clustered.client.StandardLoadBalanceFlowFileCodec;
+import
org.apache.nifi.controller.queue.clustered.client.async.TransactionFailureCallback;
+import org.apache.nifi.controller.repository.FlowFileRecord;
+import org.apache.nifi.events.EventReporter;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+import java.io.ByteArrayInputStream;
+import java.io.InputStream;
+import java.net.ServerSocket;
+import java.net.Socket;
+import java.util.List;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+class TestNioAsyncLoadBalanceClient {
+
+ private static final TransactionFailureCallback NOP_FAILURE_CALLBACK = new
TransactionFailureCallback() {
+ @Override
+ public void onTransactionFailed(final List<FlowFileRecord> flowFiles,
final Exception cause, final TransactionPhase transactionPhase) {
+ }
+
+ @Override
+ public boolean isRebalanceOnFailure() {
+ return false;
+ }
+ };
+
+ @Test
+ @Timeout(30)
+ void testUnregisterOfInFlightTransactionClosesChannel() throws Exception {
+ final ServerSocket serverSocket = new ServerSocket(0);
+ final int port = serverSocket.getLocalPort();
+
+ final Thread server = new Thread(() -> {
+ try (final Socket socket = serverSocket.accept()) {
+ final InputStream in = socket.getInputStream();
+ while (in.read() != -1) {
+ }
+ } catch (final Exception ignored) {
+ }
+ });
+ server.setDaemon(true);
+ server.start();
+
+ final String localhost = "localhost";
+ final NodeIdentifier nodeId = new NodeIdentifier("node-1", localhost,
port, localhost, port,
+ localhost, port, localhost, port, port, false);
+
+ final ClusterCoordinator clusterCoordinator =
mock(ClusterCoordinator.class);
+ when(clusterCoordinator.getConnectionStatus(nodeId)).thenReturn(new
NodeConnectionStatus(nodeId, NodeConnectionState.CONNECTED));
+
+ final FlowFileContentAccess contentAccess = ff -> new
ByteArrayInputStream(new byte[0]);
+
+ final NioAsyncLoadBalanceClient client = new
NioAsyncLoadBalanceClient(nodeId, null, 30000, contentAccess,
+ new StandardLoadBalanceFlowFileCodec(), EventReporter.NO_OP,
clusterCoordinator);
+
+ final String connectionId = "unit-test-connection";
+ try {
+ client.start();
+ client.register(connectionId, () -> false, () -> new
MockFlowFileRecord(5), NOP_FAILURE_CALLBACK,
+ (flowFiles, node) -> { }, () ->
LoadBalanceCompression.DO_NOT_COMPRESS, () -> true);
+
+ final long deadline = System.currentTimeMillis() + 20_000L;
+ while (!client.isConnectionEstablished() &&
System.currentTimeMillis() < deadline) {
+ client.communicate();
+ Thread.sleep(10L);
+ }
+ assertTrue(client.isConnectionEstablished(), "Expected the client
to establish a connection and begin a transaction");
+
+ // Unregistering while a transaction is in flight must tear down
the (now desynchronized) socket so the
+ // next transaction reconnects on a clean stream rather than
reusing it.
+ client.unregister(connectionId);
+
+ assertFalse(client.isConnectionEstablished(), "Channel must be
closed after an in-flight transaction is unregistered");
+ } finally {
+ client.stop();
+ serverSocket.close();
+ }
+ }
+}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/server/TestStandardLoadBalanceProtocol.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/server/TestStandardLoadBalanceProtocol.java
index 04a57646648..e2df99f173f 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/server/TestStandardLoadBalanceProtocol.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/server/TestStandardLoadBalanceProtocol.java
@@ -213,6 +213,48 @@ public class TestStandardLoadBalanceProtocol {
assertTrue(flowFileRepoUpdateRecords.stream().allMatch(record ->
record.getType() == RepositoryRecordType.CREATE));
}
+ @Test
+ public void
testStrayVersionByteFromReusedSocketReadAsCompletionIndicator() throws
IOException, IllegalClusterStateException {
+ final StandardLoadBalanceProtocol protocol = new
StandardLoadBalanceProtocol(flowFileRepo, contentRepo, provenanceRepo,
flowController, ALWAYS_AUTHORIZED);
+
+ final PipedInputStream serverInput = new PipedInputStream();
+ final PipedOutputStream serverContentSource = new PipedOutputStream();
+ serverInput.connect(serverContentSource);
+
+ final ByteArrayOutputStream serverOutput = new ByteArrayOutputStream();
+
+ final Checksum checksum = new CRC32();
+ final OutputStream checkedOutput = new
CheckedOutputStream(serverContentSource, checksum);
+ final DataOutputStream dos = new DataOutputStream(checkedOutput);
+ dos.writeUTF("unit-test-connection-id");
+
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put("uuid", "unit-test-id");
+
+ dos.write(CHECK_SPACE);
+ dos.write(MORE_FLOWFILES);
+ writeAttributes(attributes, dos);
+ writeContent("hello".getBytes(), dos);
+ dos.write(NO_MORE_FLOWFILES);
+ dos.writeLong(checksum.getValue());
+
+ // A transaction abandoned on a reused socket leaves the next
transaction's leading protocol-version
+ // byte (1) where the server expects COMPLETE_TRANSACTION (0x23) or
ABORT_TRANSACTION (0x24).
+ dos.write(1);
+ dos.flush();
+
+ final IOException thrown = assertThrows(IOException.class,
+ () -> protocol.receiveFlowFiles(serverInput, serverOutput,
"Unit Test", 1));
+ assertTrue(thrown.getMessage().contains("received a value of 1"),
"Unexpected message: " + thrown.getMessage());
+
+ final byte[] serverResponse = serverOutput.toByteArray();
+ assertEquals(SPACE_AVAILABLE, serverResponse[0]);
+ assertEquals(CONFIRM_CHECKSUM, serverResponse[1]);
+ assertEquals(ABORT_TRANSACTION, serverResponse[serverResponse.length -
1]);
+
+ Mockito.verify(flowFileQueue,
times(0)).receiveFromPeer(anyCollection());
+ }
+
@Test
public void testMultipleFlowFiles() throws IOException {
final StandardLoadBalanceProtocol protocol = new
StandardLoadBalanceProtocol(flowFileRepo, contentRepo, provenanceRepo,
flowController, ALWAYS_AUTHORIZED);
