NIFI-16084 - Load-balanced connection desynchronizes its socket after an in-flight transaction is unregistered

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();
     }
 

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 {

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<FlowFileRecord> 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();
+        }
+    }
+}

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<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);