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 374e7cc68b9 NIFI-16284 Wait for Load Balancing to complete in
ClusteredConnectorBacklogIT (#11615)
374e7cc68b9 is described below
commit 374e7cc68b961c0941a517b52543f4720969fffe
Author: Mark Payne <[email protected]>
AuthorDate: Wed Sep 2 21:45:45 2026 -0400
NIFI-16284 Wait for Load Balancing to complete in
ClusteredConnectorBacklogIT (#11615)
- Complete Process Group drop requests only after every connection finishes
- Describe aggregate drop request lifecycle in StandardProcessGroup
Signed-off-by: David Handermann <[email protected]>
---
.../apache/nifi/groups/StandardProcessGroup.java | 17 +++++++
.../nifi/groups/StandardProcessGroupTest.java | 57 ++++++++++++++++++++++
.../connectors/ClusteredConnectorBacklogIT.java | 54 ++++++++++++++------
3 files changed, 112 insertions(+), 16 deletions(-)
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java
index 649aac012de..dec748330ef 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/groups/StandardProcessGroup.java
@@ -1660,6 +1660,23 @@ public final class StandardProcessGroup implements
ProcessGroup {
aggregateDropFlowFileStatus.setOriginalSize(aggregateOriginalSize);
aggregateDropFlowFileStatus.setDroppedSize(aggregateDroppedSize);
aggregateDropFlowFileStatus.setCurrentSize(aggregateCurrentSize);
+
+ // Merging never moves the aggregate request into a COMPLETE state;
only handleDropAllFlowFiles does that. The
+ // lifecycle is:
+ // 1. handleDropAllFlowFiles creates the aggregate request and
merges each connection's drop request into it
+ // by calling this method once per connection.
+ // 2. A drop request against an empty queue is already COMPLETE when
it is created. Copying that state onto
+ // the aggregate request would make DropFlowFileRequest.setState
complete the aggregate's Future, which
+ // callers treat as "every connection has been drained", even
though the connections merged later may
+ // still be dropping. The aggregate request is therefore held at
DROPPING_FLOWFILES while merging.
+ // 3. Once every connection has been merged, handleDropAllFlowFiles
sets the aggregate request to COMPLETE,
+ // or to FAILURE, when the Futures of all of those connections
finish. Those Futures are already finished
+ // when each connection's queue was empty, in which case the
aggregate request becomes COMPLETE without
+ // ever waiting.
+ if (aggregateState == DropFlowFileState.COMPLETE) {
+ aggregateState = DropFlowFileState.DROPPING_FLOWFILES;
+ }
+
aggregateDropFlowFileStatus.setState(aggregateState);
}
diff --git
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java
index 0e46d63531f..73d41805d43 100644
---
a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java
+++
b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/groups/StandardProcessGroupTest.java
@@ -21,11 +21,19 @@ import org.apache.nifi.components.state.Scope;
import org.apache.nifi.components.state.StateManager;
import org.apache.nifi.components.state.StateManagerProvider;
import org.apache.nifi.components.state.StateMap;
+import org.apache.nifi.connectable.Connectable;
+import org.apache.nifi.connectable.ConnectableType;
+import org.apache.nifi.connectable.Connection;
import org.apache.nifi.controller.ClusterTopologyProvider;
import org.apache.nifi.controller.NodeTypeProvider;
import org.apache.nifi.controller.ProcessScheduler;
import org.apache.nifi.controller.ReloadComponent;
import org.apache.nifi.controller.flow.FlowManager;
+import org.apache.nifi.controller.queue.DropFlowFileRequest;
+import org.apache.nifi.controller.queue.DropFlowFileState;
+import org.apache.nifi.controller.queue.DropFlowFileStatus;
+import org.apache.nifi.controller.queue.FlowFileQueue;
+import org.apache.nifi.controller.queue.QueueSize;
import org.apache.nifi.controller.service.ControllerServiceProvider;
import org.apache.nifi.encrypt.PropertyEncryptor;
import org.apache.nifi.flow.ExecutionEngine;
@@ -52,6 +60,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@@ -394,6 +403,33 @@ class StandardProcessGroupTest {
assertEquals(expected, leaf.getLoggingAttributes());
}
+ @Test
+ void testDropAllFlowFilesCompletionWaitsForEveryConnection() {
+ when(flowManager.getFlowAnalyzer()).thenReturn(Optional.empty());
+
+ final DropFlowFileRequest emptyDrop = new
DropFlowFileRequest("drop-all");
+ emptyDrop.setOriginalSize(new QueueSize(0, 0L));
+ emptyDrop.setCurrentSize(new QueueSize(0, 0L));
+ emptyDrop.setState(DropFlowFileState.COMPLETE);
+
+ final DropFlowFileRequest pendingDrop = new
DropFlowFileRequest("drop-all");
+ pendingDrop.setOriginalSize(new QueueSize(300, 300L));
+ pendingDrop.setCurrentSize(new QueueSize(300, 300L));
+ pendingDrop.setState(DropFlowFileState.DROPPING_FLOWFILES);
+
+ addConnectionWithDropStatus("empty-connection", emptyDrop);
+ addConnectionWithDropStatus("queued-connection", pendingDrop);
+
+ final DropFlowFileStatus dropStatus =
processGroup.dropAllFlowFiles("drop-all", "test-user");
+ final CompletableFuture<Void> completionFuture =
dropStatus.getCompletionFuture();
+
+ assertFalse(completionFuture.isDone());
+
+ pendingDrop.setState(DropFlowFileState.COMPLETE);
+
+ assertTrue(completionFuture.isDone());
+ }
+
private StandardProcessGroup createStandardProcessGroup(final String id) {
return new StandardProcessGroup(
id,
@@ -413,4 +449,25 @@ class StandardProcessGroupTest {
);
}
+ private void addConnectionWithDropStatus(final String connectionId, final
DropFlowFileStatus dropStatus) {
+ final Connectable source = mock(Connectable.class);
+ final Connectable destination = mock(Connectable.class);
+
when(source.getConnectableType()).thenReturn(ConnectableType.PROCESSOR);
+
when(destination.getConnectableType()).thenReturn(ConnectableType.PROCESSOR);
+ when(source.getProcessGroup()).thenReturn(processGroup);
+ when(destination.getProcessGroup()).thenReturn(processGroup);
+
+ final FlowFileQueue flowFileQueue = mock(FlowFileQueue.class);
+ when(flowFileQueue.dropFlowFiles(eq("drop-all"),
eq("test-user"))).thenReturn(dropStatus);
+
+ final Connection connection = mock(Connection.class);
+ when(connection.getIdentifier()).thenReturn(connectionId);
+ when(connection.getSource()).thenReturn(source);
+ when(connection.getDestination()).thenReturn(destination);
+ when(connection.getFlowFileQueue()).thenReturn(flowFileQueue);
+
when(connection.getVersionedComponentId()).thenReturn(Optional.empty());
+
+ processGroup.addConnection(connection);
+ }
+
}
diff --git
a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/connectors/ClusteredConnectorBacklogIT.java
b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/connectors/ClusteredConnectorBacklogIT.java
index 13b28c267da..89f19c96c94 100644
---
a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/connectors/ClusteredConnectorBacklogIT.java
+++
b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/connectors/ClusteredConnectorBacklogIT.java
@@ -21,7 +21,10 @@ import org.apache.nifi.tests.system.NiFiInstanceFactory;
import org.apache.nifi.tests.system.NiFiSystemIT;
import org.apache.nifi.toolkit.client.NiFiClientException;
import org.apache.nifi.web.api.dto.BacklogDTO;
+import org.apache.nifi.web.api.dto.ConnectionDTO;
+import org.apache.nifi.web.api.entity.ConnectionEntity;
import org.apache.nifi.web.api.entity.ConnectorEntity;
+import org.apache.nifi.web.api.entity.ProcessGroupFlowEntity;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -90,22 +93,41 @@ public class ClusteredConnectorBacklogIT extends
NiFiSystemIT {
"Record Backlog", "0",
"Load Balance Strategy", loadBalanceStrategy));
- getClientUtil().startConnector(connector.getId());
- waitForConnectorMinQueueCount(connector.getId(), 40);
- getClientUtil().stopConnector(connector.getId());
-
- final int totalQueued =
getConnectorQueuedFlowFileCount(connector.getId());
- logger.info("Load Balance Strategy {}: cluster has {} queued FlowFiles
before reading backlog", loadBalanceStrategy, totalQueued);
- assertTrue(totalQueued > 0);
-
- final BacklogDTO backlog =
getClientUtil().getConnectorBacklog(connector.getId());
- assertNotNull(backlog);
+ try {
+ getClientUtil().startConnector(connector.getId());
+ waitForConnectorMinQueueCount(connector.getId(), 40);
+ getClientUtil().stopConnector(connector.getId());
+
+ final int totalQueued =
getConnectorQueuedFlowFileCount(connector.getId());
+ logger.info("Load Balance Strategy {}: cluster has {} queued
FlowFiles before reading backlog", loadBalanceStrategy, totalQueued);
+ assertTrue(totalQueued > 0);
+
+ final BacklogDTO backlog =
getClientUtil().getConnectorBacklog(connector.getId());
+ assertNotNull(backlog);
+
+ final long flowFiles = backlog.getFlowFileCount();
+ assertTrue(flowFiles >= 0);
+ assertTrue(flowFiles <= totalQueued,
+ "Cluster aggregate FlowFile count from queue snapshot (" +
flowFiles + ") must not exceed total queued count (" + totalQueued + ")");
+ assertEquals(Long.valueOf(flowFiles), backlog.getRecordCount());
+ assertTrue(backlog.getByteCount() >= 0L);
+ } finally {
+ // Load Balancing keeps redistributing already-queued FlowFiles
between nodes even after the Connector is
+ // stopped. A FlowFile can be in flight - removed from the sending
node's queue but not yet enqueued on the
+ // receiving node - a window in which the teardown's Connector
purge won't see it. Wait for Load Balancing to
+ // settle here, regardless of whether start/stop or the assertions
above passed, so teardown does not fail
+ // trying to delete a Connector that still has FlowFiles queued.
+ waitForLoadBalancingComplete(connector.getId());
+ }
+ }
- final long flowFiles = backlog.getFlowFileCount();
- assertTrue(flowFiles >= 0);
- assertTrue(flowFiles <= totalQueued,
- "Cluster aggregate FlowFile count from queue snapshot (" +
flowFiles + ") must not exceed total queued count (" + totalQueued + ")");
- assertEquals(Long.valueOf(flowFiles), backlog.getRecordCount());
- assertTrue(backlog.getByteCount() >= 0L);
+ private void waitForLoadBalancingComplete(final String connectorId) throws
InterruptedException {
+ waitFor(() -> {
+ final ProcessGroupFlowEntity flowEntity =
getNifiClient().getConnectorClient().getFlow(connectorId);
+ return
flowEntity.getProcessGroupFlow().getFlow().getConnections().stream()
+ .map(ConnectionEntity::getComponent)
+ .map(ConnectionDTO::getLoadBalanceStatus)
+ .noneMatch(ConnectionDTO.LOAD_BALANCE_ACTIVE::equals);
+ }, 100L, "Load Balancing to complete for Connector " + connectorId);
}
}
