From 649540d3a63363d158c6df40ebc3ef476ef439a2 Mon Sep 17 00:00:00 2001 From: Pierre Villard Date: Tue, 7 Jul 2026 14:01:07 +0200 Subject: [PATCH] NIFI-16084 - Load-balanced connection desynchronizes its socket after an in-flight transaction is unregistered --- .../async/nio/NioAsyncLoadBalanceClient.java | 6 +- .../async/nio/TestLoadBalanceSession.java | 43 +++++++ .../nio/TestNioAsyncLoadBalanceClient.java | 108 ++++++++++++++++++ .../TestStandardLoadBalanceProtocol.java | 42 +++++++ 4 files changed, 198 insertions(+), 1 deletion(-) create mode 100644 nifi-framework-bundle/nifi-framework/nifi-framework-core/src/test/java/org/apache/nifi/controller/queue/clustered/client/async/nio/TestNioAsyncLoadBalanceClient.java 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 9248cc2d4c72..376dcac1e346 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 synchronized void unregister(final String connectionId) { 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 @@ private TransactionThreshold newTransactionThreshold() { 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 de5eea230237..267e0595cd3d 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 void testSunnyCase() throws InterruptedException, IOException { 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 000000000000..5a65ffc79e84 --- /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,108 @@ +/* + * 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; + +public class TestNioAsyncLoadBalanceClient { + + private static final TransactionFailureCallback NOP_FAILURE_CALLBACK = new TransactionFailureCallback() { + @Override + public void onTransactionFailed(final List flowFiles, final Exception cause, final TransactionPhase transactionPhase) { + } + + @Override + public boolean isRebalanceOnFailure() { + return false; + } + }; + + @Test + @Timeout(30) + public 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 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); + + try { + client.start(); + client.register("unit-test-connection", () -> 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("unit-test-connection"); + + 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 04a57646648b..e2df99f173fe 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 void testSimpleFlowFileTransaction() throws IOException, IllegalClusterSt 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 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);