diff --git a/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java b/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java index 1b8c49e49f40..00520e72383b 100644 --- a/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java +++ b/nifi-connector-mock-bundle/nifi-connector-mock-test-bundle/nifi-connector-mock-integration-tests/src/test/java/org/apache/nifi/mock/connectors/tests/CreateConnectorIT.java @@ -32,15 +32,15 @@ import java.util.List; import java.util.Optional; import java.util.Set; +import java.util.concurrent.TimeoutException; -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertTrue; public class CreateConnectorIT { @Test - public void testCreateStartAndStopGenerateAndUpdateConnector() throws IOException { + public void testCreateStartAndStopGenerateAndUpdateConnector() throws IOException, TimeoutException { try (final ConnectorTestRunner testRunner = new StandardConnectorTestRunner.Builder() .connectorClassName("org.apache.nifi.mock.connectors.GenerateAndLog") .narLibraryDirectory(new File("target/libDir")) @@ -63,12 +63,12 @@ public void testCreateStartAndStopGenerateAndUpdateConnector() throws IOExceptio assertEquals("org.apache.nifi.lookup.SimpleKeyValueLookupService", controllerServices.iterator().next().getType()); testRunner.startConnector(); - testRunner.stopConnector(); + testRunner.stopConnector(Duration.ofSeconds(120)); } } @Test - public void testStopConnectorWithTimeoutStopsRunningConnector() throws IOException { + public void testStopConnectorWithTimeoutStopsRunningConnector() throws IOException, TimeoutException { try (final ConnectorTestRunner testRunner = new StandardConnectorTestRunner.Builder() .connectorClassName("org.apache.nifi.mock.connectors.GenerateAndLog") .narLibraryDirectory(new File("target/libDir")) @@ -82,7 +82,7 @@ public void testStopConnectorWithTimeoutStopsRunningConnector() throws IOExcepti // ride through the node's internal stop retries (every 10 seconds) before the state settles, so the // budget is generous enough to stay deterministic on a slow CI runner; the poll still returns the // instant the Connector reports STOPPED. - assertDoesNotThrow(() -> testRunner.stopConnector(Duration.ofSeconds(120))); + testRunner.stopConnector(Duration.ofSeconds(120)); } } diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/FlowUpdateImpact.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/FlowUpdateImpact.java new file mode 100644 index 000000000000..8247d2d7fe62 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/FlowUpdateImpact.java @@ -0,0 +1,85 @@ +/* + * 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.web; + +import org.apache.nifi.web.api.entity.AffectedComponentEntity; + +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.Objects; +import java.util.Set; + +public final class FlowUpdateImpact { + private final Set affectedComponents; + private final Set removedConnections; + private final Set removedProcessGroupIds; + private final Set removedEndpointIds; + + public FlowUpdateImpact(final Set affectedComponents, + final Set removedConnections, + final Set removedProcessGroupIds, + final Set removedEndpointIds) { + this.affectedComponents = unmodifiableCopy(affectedComponents); + this.removedConnections = unmodifiableCopy(removedConnections); + this.removedProcessGroupIds = unmodifiableCopy(removedProcessGroupIds); + this.removedEndpointIds = unmodifiableCopy(removedEndpointIds); + } + + public Set getAffectedComponents() { + return affectedComponents; + } + + public Set getRemovedConnections() { + return removedConnections; + } + + public Set getRemovedProcessGroupIds() { + return removedProcessGroupIds; + } + + public Set getRemovedEndpointIds() { + return removedEndpointIds; + } + + private static Set unmodifiableCopy(final Set values) { + if (values == null || values.isEmpty()) { + return Collections.emptySet(); + } + + return Collections.unmodifiableSet(new LinkedHashSet<>(values)); + } + + @Override + public boolean equals(final Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof final FlowUpdateImpact other)) { + return false; + } + + return Objects.equals(affectedComponents, other.affectedComponents) + && Objects.equals(removedConnections, other.removedConnections) + && Objects.equals(removedProcessGroupIds, other.removedProcessGroupIds) + && Objects.equals(removedEndpointIds, other.removedEndpointIds); + } + + @Override + public int hashCode() { + return Objects.hash(affectedComponents, removedConnections, removedProcessGroupIds, removedEndpointIds); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/NiFiServiceFacade.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/NiFiServiceFacade.java index 2313b34f977c..e3ced6bc0320 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/NiFiServiceFacade.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/NiFiServiceFacade.java @@ -2049,6 +2049,18 @@ VersionControlInformationEntity setVersionControlInformation(Revision processGro */ String getFlowRegistryName(String flowRegistryId); + /** + * Determines which components currently exist in the Process Group with the given identifier and calculates which of those components + * would be impacted by updating the Process Group to the provided snapshot + * + * @param processGroupId the ID of the Process Group to update + * @param updatedSnapshot the snapshot to update the Process Group to + * @return the impact of updating the Process Group + */ + FlowUpdateImpact getFlowUpdateImpact(String processGroupId, RegisteredFlowSnapshot updatedSnapshot); + + RemovedConnectionDrainClassifier.Context getRemovedConnectionDrainContext(); + /** * Determines which components currently exist in the Process Group with the given identifier and calculates which of those components * would be impacted by updating the Process Group to the provided snapshot diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovalReason.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovalReason.java new file mode 100644 index 000000000000..3926d6849b8c --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovalReason.java @@ -0,0 +1,22 @@ +/* + * 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.web; + +public enum RemovalReason { + COMPONENT_REMOVED, + SOURCE_CHANGED +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDescriptor.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDescriptor.java new file mode 100644 index 000000000000..1694711aec09 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDescriptor.java @@ -0,0 +1,136 @@ +/* + * 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.web; + +import org.apache.nifi.connectable.ConnectableType; + +import java.util.Objects; + +public final class RemovedConnectionDescriptor { + private final String connectionInstanceId; + private final String connectionVersionedId; + private final String containingProcessGroupId; + private final String sourceInstanceId; + private final String sourceVersionedId; + private final String sourceProcessGroupId; + private final ConnectableType sourceType; + private final String destinationInstanceId; + private final String destinationVersionedId; + private final String destinationProcessGroupId; + private final ConnectableType destinationType; + private final RemovalReason removalReason; + + public RemovedConnectionDescriptor(final String connectionInstanceId, final String connectionVersionedId, + final String containingProcessGroupId, + final String sourceInstanceId, final String sourceVersionedId, final String sourceProcessGroupId, + final ConnectableType sourceType, + final String destinationInstanceId, final String destinationVersionedId, + final String destinationProcessGroupId, final ConnectableType destinationType, + final RemovalReason removalReason) { + this.connectionInstanceId = connectionInstanceId; + this.connectionVersionedId = connectionVersionedId; + this.containingProcessGroupId = containingProcessGroupId; + this.sourceInstanceId = sourceInstanceId; + this.sourceVersionedId = sourceVersionedId; + this.sourceProcessGroupId = sourceProcessGroupId; + this.sourceType = sourceType; + this.destinationInstanceId = destinationInstanceId; + this.destinationVersionedId = destinationVersionedId; + this.destinationProcessGroupId = destinationProcessGroupId; + this.destinationType = destinationType; + this.removalReason = removalReason; + } + + public String getConnectionInstanceId() { + return connectionInstanceId; + } + + public String getConnectionVersionedId() { + return connectionVersionedId; + } + + public String getContainingProcessGroupId() { + return containingProcessGroupId; + } + + public String getSourceInstanceId() { + return sourceInstanceId; + } + + public String getSourceVersionedId() { + return sourceVersionedId; + } + + public String getSourceProcessGroupId() { + return sourceProcessGroupId; + } + + public ConnectableType getSourceType() { + return sourceType; + } + + public String getDestinationInstanceId() { + return destinationInstanceId; + } + + public String getDestinationVersionedId() { + return destinationVersionedId; + } + + public String getDestinationProcessGroupId() { + return destinationProcessGroupId; + } + + public ConnectableType getDestinationType() { + return destinationType; + } + + public RemovalReason getRemovalReason() { + return removalReason; + } + + @Override + public boolean equals(final Object obj) { + if (this == obj) { + return true; + } + if (!(obj instanceof final RemovedConnectionDescriptor other)) { + return false; + } + + return Objects.equals(connectionInstanceId, other.connectionInstanceId) + && Objects.equals(connectionVersionedId, other.connectionVersionedId) + && Objects.equals(containingProcessGroupId, other.containingProcessGroupId) + && Objects.equals(sourceInstanceId, other.sourceInstanceId) + && Objects.equals(sourceVersionedId, other.sourceVersionedId) + && Objects.equals(sourceProcessGroupId, other.sourceProcessGroupId) + && sourceType == other.sourceType + && Objects.equals(destinationInstanceId, other.destinationInstanceId) + && Objects.equals(destinationVersionedId, other.destinationVersionedId) + && Objects.equals(destinationProcessGroupId, other.destinationProcessGroupId) + && destinationType == other.destinationType + && removalReason == other.removalReason; + } + + @Override + public int hashCode() { + return Objects.hash(connectionInstanceId, connectionVersionedId, containingProcessGroupId, + sourceInstanceId, sourceVersionedId, sourceProcessGroupId, sourceType, + destinationInstanceId, destinationVersionedId, destinationProcessGroupId, destinationType, + removalReason); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDrainClassifier.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDrainClassifier.java new file mode 100644 index 000000000000..9fcbd581dcc0 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDrainClassifier.java @@ -0,0 +1,527 @@ +/* + * 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.web; + +import org.apache.nifi.components.validation.ValidationStatus; +import org.apache.nifi.connectable.Connectable; +import org.apache.nifi.connectable.ConnectableType; +import org.apache.nifi.connectable.Connection; +import org.apache.nifi.connectable.Port; +import org.apache.nifi.controller.ProcessorNode; +import org.apache.nifi.controller.ScheduledState; +import org.apache.nifi.controller.flow.FlowManager; +import org.apache.nifi.groups.ProcessGroup; + +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Comparator; +import java.util.Deque; +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.stream.Collector; +import java.util.stream.Collectors; + +public final class RemovedConnectionDrainClassifier { + private static final Set SUPPORTED_SOURCE_TYPES = Set.of( + ConnectableType.PROCESSOR, + ConnectableType.INPUT_PORT, + ConnectableType.OUTPUT_PORT, + ConnectableType.FUNNEL); + private static final Set SUPPORTED_PRODUCER_BARRIER_TYPES = Set.of( + ConnectableType.PROCESSOR, + ConnectableType.INPUT_PORT, + ConnectableType.OUTPUT_PORT); + private static final Set SUPPORTED_DESTINATION_TYPES = Set.of( + ConnectableType.PROCESSOR, + ConnectableType.INPUT_PORT, + ConnectableType.OUTPUT_PORT); + private static final Comparator CONNECTION_ORDER = + Comparator.comparing(RemovedConnectionDescriptor::getConnectionInstanceId, Comparator.nullsLast(String::compareTo)) + .thenComparing(RemovedConnectionDescriptor::getConnectionVersionedId, Comparator.nullsLast(String::compareTo)); + + BatchResult classify(final FlowUpdateImpact flowUpdateImpact, final Context context) { + Objects.requireNonNull(flowUpdateImpact, "Flow Update Impact required"); + Objects.requireNonNull(context, "Removed Connection Drain Context required"); + + final List orderedConnections = flowUpdateImpact.getRemovedConnections().stream() + .sorted(CONNECTION_ORDER) + .toList(); + + final Set removedConnectionIds = orderedConnections.stream() + .map(RemovedConnectionDescriptor::getConnectionInstanceId) + .filter(Objects::nonNull) + .collect(toOrderedSet()); + + final Set nonEmptyRemovedDestinationIds = new LinkedHashSet<>(); + final Map initialResults = new LinkedHashMap<>(); + + for (final RemovedConnectionDescriptor descriptor : orderedConnections) { + final LiveConnection liveConnection = context.getConnection(descriptor.getConnectionInstanceId()); + if (liveConnection == null) { + initialResults.put(descriptor.getConnectionInstanceId(), ConnectionResult.unsupported(descriptor, UnsupportedReason.CONNECTION_NOT_FOUND)); + continue; + } + + if (liveConnection.knownQueueEmpty()) { + initialResults.put(descriptor.getConnectionInstanceId(), ConnectionResult.noDrain(descriptor)); + continue; + } + + if (descriptor.getDestinationInstanceId() != null) { + nonEmptyRemovedDestinationIds.add(descriptor.getDestinationInstanceId()); + } + } + + final Map classifiedResults = new LinkedHashMap<>(initialResults); + for (final RemovedConnectionDescriptor descriptor : orderedConnections) { + if (classifiedResults.containsKey(descriptor.getConnectionInstanceId())) { + continue; + } + + classifiedResults.put(descriptor.getConnectionInstanceId(), classifyNonEmpty(descriptor, flowUpdateImpact, context, + nonEmptyRemovedDestinationIds, removedConnectionIds)); + } + + final Set candidateProducerBarrierIds = getCandidateProducerBarrierIds(orderedConnections, classifiedResults); + + for (final RemovedConnectionDescriptor descriptor : orderedConnections) { + final ConnectionResult connectionResult = classifiedResults.get(descriptor.getConnectionInstanceId()); + if (connectionResult.classification() != Classification.CANDIDATE) { + continue; + } + + if (hasRetainedFeedbackPath(descriptor.getDestinationInstanceId(), candidateProducerBarrierIds, removedConnectionIds, context)) { + classifiedResults.put(descriptor.getConnectionInstanceId(), ConnectionResult.unsupported(descriptor, UnsupportedReason.RETAINED_FEEDBACK_PATH)); + } + } + + final List connectionResults = new ArrayList<>(orderedConnections.size()); + for (final RemovedConnectionDescriptor descriptor : orderedConnections) { + connectionResults.add(classifiedResults.get(descriptor.getConnectionInstanceId())); + } + + return new BatchResult(connectionResults, getCandidateProducerBarrierIds(orderedConnections, classifiedResults)); + } + + private ConnectionResult classifyNonEmpty(final RemovedConnectionDescriptor descriptor, final FlowUpdateImpact flowUpdateImpact, + final Context context, final Set nonEmptyRemovedDestinationIds, + final Set removedConnectionIds) { + if (descriptor.getRemovalReason() == RemovalReason.SOURCE_CHANGED) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.SOURCE_CHANGED_REMOVAL); + } + + if (!isRetainedGroupHierarchy(descriptor.getContainingProcessGroupId(), flowUpdateImpact.getRemovedProcessGroupIds(), context)) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.CONNECTION_IN_REMOVED_GROUP); + } + + if (isRemovedEndpoint(flowUpdateImpact, descriptor.getSourceInstanceId())) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.SOURCE_COMPONENT_REMOVED); + } + + if (isRemovedEndpoint(flowUpdateImpact, descriptor.getDestinationInstanceId())) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.DESTINATION_COMPONENT_REMOVED); + } + + if (!SUPPORTED_SOURCE_TYPES.contains(descriptor.getSourceType())) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.UNSUPPORTED_SOURCE_TYPE); + } + + if (!SUPPORTED_DESTINATION_TYPES.contains(descriptor.getDestinationType())) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.UNSUPPORTED_DESTINATION_TYPE); + } + + final LiveConnectable destination = context.getConnectable(descriptor.getDestinationInstanceId()); + if (destination == null) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.DESTINATION_COMPONENT_NOT_FOUND); + } + + final Optional destinationValidationFailure = validateDestination(destination); + if (destinationValidationFailure.isPresent()) { + return ConnectionResult.unsupported(descriptor, destinationValidationFailure.get()); + } + + final Set producerBarrierIds = resolveProducerBarrierIds(descriptor, flowUpdateImpact, context, nonEmptyRemovedDestinationIds); + if (producerBarrierIds.isEmpty()) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.NO_SUPPORTED_PRODUCER_FOUND); + } + + for (final String producerBarrierId : producerBarrierIds) { + if (Objects.equals(producerBarrierId, descriptor.getDestinationInstanceId())) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.SELF_LOOP); + } + + if (nonEmptyRemovedDestinationIds.contains(producerBarrierId)) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.PRODUCER_BARRIER_IS_REMOVED_DESTINATION); + } + + final LiveConnectable producerBarrier = context.getConnectable(producerBarrierId); + if (producerBarrier == null) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.SOURCE_COMPONENT_NOT_FOUND); + } + + if (!SUPPORTED_PRODUCER_BARRIER_TYPES.contains(producerBarrier.type())) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.UNSUPPORTED_SOURCE_TYPE); + } + + if (!isRetainedGroupHierarchy(producerBarrier.processGroupId(), flowUpdateImpact.getRemovedProcessGroupIds(), context)) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.SOURCE_COMPONENT_REMOVED); + } + + if (isRemovedEndpoint(flowUpdateImpact, producerBarrierId)) { + return ConnectionResult.unsupported(descriptor, UnsupportedReason.SOURCE_COMPONENT_REMOVED); + } + } + + return ConnectionResult.candidate(descriptor, producerBarrierIds); + } + + private Set resolveProducerBarrierIds(final RemovedConnectionDescriptor descriptor, final FlowUpdateImpact flowUpdateImpact, + final Context context, final Set nonEmptyRemovedDestinationIds) { + if (descriptor.getSourceType() != ConnectableType.FUNNEL) { + if (descriptor.getSourceInstanceId() == null) { + return Collections.emptySet(); + } + + final LiveConnectable source = context.getConnectable(descriptor.getSourceInstanceId()); + if (source == null) { + return Collections.emptySet(); + } + + if (!isRetainedGroupHierarchy(source.processGroupId(), flowUpdateImpact.getRemovedProcessGroupIds(), context)) { + return Collections.emptySet(); + } + + return Set.of(source.id()); + } + + final LiveConnectable funnel = context.getConnectable(descriptor.getSourceInstanceId()); + if (funnel == null) { + return Collections.emptySet(); + } + + final Set producerBarrierIds = new LinkedHashSet<>(); + collectUpstreamProducerBarrierIds(funnel, context, producerBarrierIds, new HashSet<>()); + return producerBarrierIds; + } + + private void collectUpstreamProducerBarrierIds(final LiveConnectable connectable, final Context context, + final Set producerBarrierIds, final Set visitedConnectionIds) { + if (connectable == null) { + return; + } + + if (connectable.type() != ConnectableType.FUNNEL) { + producerBarrierIds.add(connectable.id()); + return; + } + + for (final String incomingConnectionId : connectable.incomingConnectionIds()) { + if (!visitedConnectionIds.add(incomingConnectionId)) { + continue; + } + + final LiveConnection incomingConnection = context.getConnection(incomingConnectionId); + if (incomingConnection == null) { + continue; + } + + collectUpstreamProducerBarrierIds(context.getConnectable(incomingConnection.sourceId()), context, producerBarrierIds, visitedConnectionIds); + } + } + + private Optional validateDestination(final LiveConnectable destination) { + if (destination.type() == ConnectableType.PROCESSOR) { + if (destination.physicalScheduledState() != ScheduledState.RUNNING) { + return Optional.of(UnsupportedReason.DESTINATION_NOT_RUNNING); + } + + if (destination.validationStatus() != ValidationStatus.VALID) { + return Optional.of(UnsupportedReason.DESTINATION_NOT_VALID); + } + } else if ((destination.type() == ConnectableType.INPUT_PORT || destination.type() == ConnectableType.OUTPUT_PORT) && !destination.running()) { + return Optional.of(UnsupportedReason.DESTINATION_NOT_RUNNING); + } + + return Optional.empty(); + } + + private boolean hasRetainedFeedbackPath(final String destinationId, final Set producerBarrierIds, + final Set removedConnectionIds, final Context context) { + if (destinationId == null || producerBarrierIds.isEmpty()) { + return false; + } + + final Deque pendingConnectables = new ArrayDeque<>(); + final Set visitedConnectables = new HashSet<>(); + pendingConnectables.add(destinationId); + + while (!pendingConnectables.isEmpty()) { + final String connectableId = pendingConnectables.removeFirst(); + if (!visitedConnectables.add(connectableId)) { + continue; + } + + final LiveConnectable connectable = context.getConnectable(connectableId); + if (connectable == null) { + continue; + } + + for (final String outgoingConnectionId : connectable.outgoingConnectionIds()) { + if (removedConnectionIds.contains(outgoingConnectionId)) { + continue; + } + + final LiveConnection outgoingConnection = context.getConnection(outgoingConnectionId); + if (outgoingConnection == null) { + continue; + } + + final String downstreamConnectableId = outgoingConnection.destinationId(); + if (producerBarrierIds.contains(downstreamConnectableId)) { + return true; + } + + pendingConnectables.addLast(downstreamConnectableId); + } + } + + return false; + } + + private boolean isRemovedEndpoint(final FlowUpdateImpact flowUpdateImpact, final String componentInstanceId) { + return componentInstanceId != null && flowUpdateImpact.getRemovedEndpointIds().contains(componentInstanceId); + } + + private boolean isRetainedGroupHierarchy(final String processGroupId, final Set removedProcessGroupIds, final Context context) { + if (processGroupId == null) { + return false; + } + + String currentGroupId = processGroupId; + final Set visitedGroupIds = new HashSet<>(); + while (currentGroupId != null) { + if (!visitedGroupIds.add(currentGroupId)) { + return false; + } + + if (removedProcessGroupIds.contains(currentGroupId)) { + return false; + } + + final LiveProcessGroup processGroup = context.getProcessGroup(currentGroupId); + if (processGroup == null) { + return false; + } + + currentGroupId = processGroup.parentProcessGroupId(); + } + + return true; + } + + private Set getCandidateProducerBarrierIds(final List orderedConnections, + final Map classifiedResults) { + final Set candidateProducerBarrierIds = new LinkedHashSet<>(); + for (final RemovedConnectionDescriptor descriptor : orderedConnections) { + final ConnectionResult connectionResult = classifiedResults.get(descriptor.getConnectionInstanceId()); + if (connectionResult != null && connectionResult.classification() == Classification.CANDIDATE) { + candidateProducerBarrierIds.addAll(connectionResult.producerBarrierComponentIds()); + } + } + + return candidateProducerBarrierIds; + } + + private static Collector> toOrderedSet() { + return Collectors.toCollection(LinkedHashSet::new); + } + + public interface Context { + LiveConnection getConnection(String connectionId); + + LiveConnectable getConnectable(String connectableId); + + LiveProcessGroup getProcessGroup(String processGroupId); + } + + enum Classification { + NO_DRAIN, + CANDIDATE, + UNSUPPORTED + } + + enum UnsupportedReason { + CONNECTION_NOT_FOUND, + SOURCE_CHANGED_REMOVAL, + CONNECTION_IN_REMOVED_GROUP, + SOURCE_COMPONENT_REMOVED, + DESTINATION_COMPONENT_REMOVED, + SOURCE_COMPONENT_NOT_FOUND, + DESTINATION_COMPONENT_NOT_FOUND, + SELF_LOOP, + UNSUPPORTED_SOURCE_TYPE, + UNSUPPORTED_DESTINATION_TYPE, + DESTINATION_NOT_RUNNING, + DESTINATION_NOT_VALID, + PRODUCER_BARRIER_IS_REMOVED_DESTINATION, + RETAINED_FEEDBACK_PATH, + NO_SUPPORTED_PRODUCER_FOUND + } + + record LiveConnection(String id, String sourceId, String destinationId, boolean knownQueueEmpty) { + } + + record LiveConnectable(String id, ConnectableType type, String processGroupId, + ScheduledState physicalScheduledState, ValidationStatus validationStatus, boolean running, + Set incomingConnectionIds, Set outgoingConnectionIds) { + LiveConnectable { + incomingConnectionIds = copyOrderedSet(incomingConnectionIds); + outgoingConnectionIds = copyOrderedSet(outgoingConnectionIds); + } + } + + record LiveProcessGroup(String id, String parentProcessGroupId) { + } + + record ConnectionResult(RemovedConnectionDescriptor connection, Classification classification, + UnsupportedReason unsupportedReason, Set producerBarrierComponentIds) { + ConnectionResult { + producerBarrierComponentIds = copyOrderedSet(producerBarrierComponentIds); + } + + static ConnectionResult noDrain(final RemovedConnectionDescriptor connection) { + return new ConnectionResult(connection, Classification.NO_DRAIN, null, Collections.emptySet()); + } + + static ConnectionResult candidate(final RemovedConnectionDescriptor connection, final Set producerBarrierComponentIds) { + return new ConnectionResult(connection, Classification.CANDIDATE, null, producerBarrierComponentIds); + } + + static ConnectionResult unsupported(final RemovedConnectionDescriptor connection, final UnsupportedReason unsupportedReason) { + return new ConnectionResult(connection, Classification.UNSUPPORTED, unsupportedReason, Collections.emptySet()); + } + } + + record BatchResult(List connectionResults, Set producerBarrierComponentIds) { + BatchResult { + connectionResults = Collections.unmodifiableList(new ArrayList<>(connectionResults)); + producerBarrierComponentIds = copyOrderedSet(producerBarrierComponentIds); + } + + boolean isSupported() { + return connectionResults.stream().noneMatch(result -> result.classification() == Classification.UNSUPPORTED); + } + + Optional getFirstUnsupportedConnectionResult() { + return connectionResults.stream().filter(result -> result.classification() == Classification.UNSUPPORTED).findFirst(); + } + } + + static final class FlowManagerContext implements Context { + private final FlowManager flowManager; + + FlowManagerContext(final FlowManager flowManager) { + this.flowManager = Objects.requireNonNull(flowManager, "Flow Manager required"); + } + + @Override + public LiveConnection getConnection(final String connectionId) { + final Connection connection = flowManager.getConnection(connectionId); + if (connection == null) { + return null; + } + + return new LiveConnection( + connection.getIdentifier(), + getConnectableIdentifier(connection.getSource()), + getConnectableIdentifier(connection.getDestination()), + connection.getFlowFileQueue().isEmpty()); + } + + @Override + public LiveConnectable getConnectable(final String connectableId) { + final Connectable connectable = flowManager.findConnectable(connectableId); + if (connectable == null) { + return null; + } + + ScheduledState physicalScheduledState = null; + ValidationStatus validationStatus = null; + boolean running = false; + if (connectable instanceof final ProcessorNode processorNode) { + physicalScheduledState = processorNode.getPhysicalScheduledState(); + validationStatus = processorNode.getValidationStatus(); + } else if (connectable instanceof final Port port) { + running = port.isRunning(); + } + + return new LiveConnectable( + connectable.getIdentifier(), + connectable.getConnectableType(), + connectable.getProcessGroupIdentifier(), + physicalScheduledState, + validationStatus, + running, + getConnectionIdentifiers(connectable.getIncomingConnections()), + getConnectionIdentifiers(connectable.getConnections())); + } + + @Override + public LiveProcessGroup getProcessGroup(final String processGroupId) { + final ProcessGroup processGroup = flowManager.getGroup(processGroupId); + if (processGroup == null) { + return null; + } + + final ProcessGroup parent = processGroup.getParent(); + return new LiveProcessGroup(processGroup.getIdentifier(), parent == null ? null : parent.getIdentifier()); + } + + private String getConnectableIdentifier(final Connectable connectable) { + return connectable == null ? null : connectable.getIdentifier(); + } + + private Set getConnectionIdentifiers(final Collection connections) { + if (connections == null || connections.isEmpty()) { + return Collections.emptySet(); + } + + final Set connectionIdentifiers = new LinkedHashSet<>(connections.size()); + for (final Connection connection : connections) { + connectionIdentifiers.add(connection.getIdentifier()); + } + + return connectionIdentifiers; + } + } + + private static Set copyOrderedSet(final Collection values) { + if (values == null || values.isEmpty()) { + return Collections.emptySet(); + } + + return Collections.unmodifiableSet(new LinkedHashSet<>(values)); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDrainCoordinator.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDrainCoordinator.java new file mode 100644 index 000000000000..4543a18326b5 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/RemovedConnectionDrainCoordinator.java @@ -0,0 +1,567 @@ +/* + * 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.web; + +import org.apache.nifi.connectable.ConnectableType; +import org.apache.nifi.controller.ScheduledState; +import org.apache.nifi.web.api.dto.AffectedComponentDTO; +import org.apache.nifi.web.api.entity.AffectedComponentEntity; +import org.apache.nifi.web.util.CancellableTimedPause; +import org.apache.nifi.web.util.ComponentLifecycle; +import org.apache.nifi.web.util.InvalidComponentAction; +import org.apache.nifi.web.util.LifecycleManagementException; +import org.apache.nifi.web.util.Pause; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.net.URI; +import java.time.Duration; +import java.util.Collection; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.function.LongSupplier; +import java.util.stream.Collectors; + +public final class RemovedConnectionDrainCoordinator { + private static final Logger logger = LoggerFactory.getLogger(RemovedConnectionDrainCoordinator.class); + static final Duration DEFAULT_DRAIN_TIMEOUT = Duration.ofSeconds(30); + private static final Duration DEFAULT_POLL_INTERVAL = Duration.ofMillis(250); + + private final RemovedConnectionDrainClassifier classifier; + private final PauseFactory pauseFactory; + private final Duration drainTimeout; + + public RemovedConnectionDrainCoordinator() { + this(new RemovedConnectionDrainClassifier(), new MonotonicPauseFactory(DEFAULT_POLL_INTERVAL, System::nanoTime), DEFAULT_DRAIN_TIMEOUT); + } + + RemovedConnectionDrainCoordinator(final RemovedConnectionDrainClassifier classifier, final PauseFactory pauseFactory, final Duration drainTimeout) { + this.classifier = Objects.requireNonNull(classifier, "Removed Connection Drain Classifier required"); + this.pauseFactory = Objects.requireNonNull(pauseFactory, "Pause Factory required"); + this.drainTimeout = Objects.requireNonNull(drainTimeout, "Drain Timeout required"); + } + + public DrainResult coordinateDrain(final FlowUpdateImpact flowUpdateImpact, final RemovedConnectionDrainClassifier.Context context, + final ComponentLifecycle componentLifecycle, final URI requestUri, final String groupId, + final CancellationHandle cancellationHandle) throws LifecycleManagementException { + Objects.requireNonNull(flowUpdateImpact, "Flow Update Impact required"); + Objects.requireNonNull(context, "Removed Connection Drain Context required"); + Objects.requireNonNull(componentLifecycle, "Component Lifecycle required"); + Objects.requireNonNull(requestUri, "Request URI required"); + Objects.requireNonNull(groupId, "Group ID required"); + Objects.requireNonNull(cancellationHandle, "Cancellation Handle required"); + + final RemovedConnectionDrainClassifier.Context queueAwareContext = createQueueAwareContext(flowUpdateImpact, context, componentLifecycle, requestUri); + final RemovedConnectionDrainClassifier.BatchResult batchResult = classifier.classify(flowUpdateImpact, queueAwareContext); + if (!batchResult.isSupported()) { + throw new LifecycleManagementException(buildClassificationFailureMessage(batchResult)); + } + + final Set candidateConnectionIds = batchResult.connectionResults().stream() + .filter(result -> result.classification() == RemovedConnectionDrainClassifier.Classification.CANDIDATE) + .map(result -> result.connection().getConnectionInstanceId()) + .collect(Collectors.toCollection(LinkedHashSet::new)); + + if (candidateConnectionIds.isEmpty()) { + return DrainResult.success(Collections.emptySet(), Collections.emptySet()); + } + + final Map affectedComponentsById = flowUpdateImpact.getAffectedComponents().stream() + .collect(Collectors.toMap(AffectedComponentEntity::getId, entity -> entity, (left, right) -> left, LinkedHashMap::new)); + final Set componentsToStop = new LinkedHashSet<>(); + for (final String producerBarrierComponentId : batchResult.producerBarrierComponentIds()) { + final AffectedComponentEntity entity = getProducerBarrierEntity(affectedComponentsById, queueAwareContext, producerBarrierComponentId); + if (entity == null || entity.getComponent() == null) { + continue; + } + + if (isActive(entity.getComponent())) { + componentsToStop.add(entity); + } + } + + final List orderedCandidateConnectionIds = candidateConnectionIds.stream().sorted().toList(); + final List orderedProducerBarrierIds = componentsToStop.stream().map(AffectedComponentEntity::getId).sorted().toList(); + logger.info("Starting drain of removed connections {} with producer barriers {}", orderedCandidateConnectionIds, orderedProducerBarrierIds); + + final DeadlinePause drainPause = pauseFactory.createDrainPause(drainTimeout); + cancellationHandle.setCancelCallback(drainPause::cancel); + + final Set drainStoppedComponents = new LinkedHashSet<>(); + try { + if (!componentsToStop.isEmpty()) { + final Set updatedStoppedComponents = componentLifecycle.scheduleComponents( + requestUri, groupId, componentsToStop, ScheduledState.STOPPED, drainPause, InvalidComponentAction.SKIP); + drainStoppedComponents.addAll(getStoppedComponents(componentsToStop, updatedStoppedComponents)); + + if (!allComponentsStopped(componentsToStop, updatedStoppedComponents)) { + if (cancellationHandle.isCancelled()) { + return restoreAfterCancellation(componentLifecycle, requestUri, groupId, candidateConnectionIds, drainStoppedComponents); + } + + final Set producerBarrierIds = componentsToStop.stream() + .map(AffectedComponentEntity::getId) + .collect(Collectors.toCollection(LinkedHashSet::new)); + throw new LifecycleManagementException(buildStopTimeoutMessage(producerBarrierIds)); + } + } + + if (cancellationHandle.isCancelled()) { + return restoreAfterCancellation(componentLifecycle, requestUri, groupId, candidateConnectionIds, drainStoppedComponents); + } + + final boolean queuesDrained = componentLifecycle.waitForConnectionQueuesEmpty(requestUri, candidateConnectionIds, drainPause); + if (queuesDrained) { + if (cancellationHandle.isCancelled()) { + return restoreAfterCancellation(componentLifecycle, requestUri, groupId, candidateConnectionIds, drainStoppedComponents); + } + + logger.info("Completed draining removed connections {}", orderedCandidateConnectionIds); + return DrainResult.success(candidateConnectionIds, drainStoppedComponents); + } + + if (cancellationHandle.isCancelled()) { + return restoreAfterCancellation(componentLifecycle, requestUri, groupId, candidateConnectionIds, drainStoppedComponents); + } + + throw new LifecycleManagementException(buildQueueTimeoutMessage(candidateConnectionIds)); + } catch (final LifecycleManagementException e) { + final Set stoppedComponents = getStoppedComponentsToRestore(queueAwareContext, componentsToStop, drainStoppedComponents); + if (cancellationHandle.isCancelled()) { + return restoreAfterCancellation(componentLifecycle, requestUri, groupId, candidateConnectionIds, stoppedComponents); + } + + final LifecycleManagementException failure = decorateFailure(e, componentsToStop, candidateConnectionIds); + logger.warn("Removed connection drain failed for connections {}", orderedCandidateConnectionIds, failure); + restoreOrSuppress(componentLifecycle, requestUri, groupId, stoppedComponents, failure); + throw failure; + } catch (final RuntimeException e) { + final Set stoppedComponents = getStoppedComponentsToRestore(queueAwareContext, componentsToStop, drainStoppedComponents); + if (cancellationHandle.isCancelled()) { + return restoreAfterCancellation(componentLifecycle, requestUri, groupId, candidateConnectionIds, stoppedComponents); + } + + final LifecycleManagementException failure = new LifecycleManagementException( + "Removed connection drain failed for connections " + candidateConnectionIds.stream().sorted().toList(), e); + logger.warn("Removed connection drain failed for connections {}", orderedCandidateConnectionIds, failure); + restoreOrSuppress(componentLifecycle, requestUri, groupId, stoppedComponents, failure); + throw failure; + } finally { + cancellationHandle.setCancelCallback(null); + } + } + + private DrainResult restoreAfterCancellation(final ComponentLifecycle componentLifecycle, final URI requestUri, final String groupId, + final Set candidateConnectionIds, + final Set drainStoppedComponents) { + LifecycleManagementException restorationFailure = null; + try { + restoreStoppedComponents(componentLifecycle, requestUri, groupId, drainStoppedComponents); + } catch (final LifecycleManagementException e) { + restorationFailure = e; + logger.warn("Failed to restore producer barriers {} after removed connection drain cancellation", + drainStoppedComponents.stream().map(AffectedComponentEntity::getId).sorted().toList(), e); + } + + return DrainResult.cancelled(candidateConnectionIds, drainStoppedComponents, restorationFailure); + } + + private String buildClassificationFailureMessage(final RemovedConnectionDrainClassifier.BatchResult batchResult) { + final List unsupportedConnections = batchResult.connectionResults().stream() + .filter(result -> result.classification() == RemovedConnectionDrainClassifier.Classification.UNSUPPORTED) + .map(result -> result.connection().getConnectionInstanceId() + "[reason=" + result.unsupportedReason().name() + "]") + .toList(); + return "Removed connection drain preflight failed: " + unsupportedConnections; + } + + private String buildStopTimeoutMessage(final Set producerBarrierIds) { + return "Removed connection drain timed out [timedOut=true, cancelled=false, phase=stopping-producer-barriers, timeout=" + + drainTimeout.toSeconds() + "s, componentIds=" + producerBarrierIds.stream().sorted().toList() + "]"; + } + + private String buildQueueTimeoutMessage(final Set connectionIds) { + return "Removed connection drain timed out [timedOut=true, cancelled=false, phase=waiting-for-queues, timeout=" + + drainTimeout.toSeconds() + "s, connectionIds=" + connectionIds.stream().sorted().toList() + "]"; + } + + private AffectedComponentEntity getProducerBarrierEntity(final Map affectedComponentsById, + final RemovedConnectionDrainClassifier.Context context, + final String producerBarrierComponentId) { + final AffectedComponentEntity affectedComponentEntity = affectedComponentsById.get(producerBarrierComponentId); + if (affectedComponentEntity != null && affectedComponentEntity.getComponent() != null) { + return affectedComponentEntity; + } + + final RemovedConnectionDrainClassifier.LiveConnectable liveConnectable = context.getConnectable(producerBarrierComponentId); + if (liveConnectable == null) { + return null; + } + + final String referenceType = getReferenceType(liveConnectable.type()); + if (referenceType == null) { + return null; + } + + final AffectedComponentDTO componentDto = new AffectedComponentDTO(); + componentDto.setId(liveConnectable.id()); + componentDto.setName(liveConnectable.id()); + componentDto.setProcessGroupId(liveConnectable.processGroupId()); + componentDto.setReferenceType(referenceType); + componentDto.setState(getState(liveConnectable)); + + final AffectedComponentEntity componentEntity = new AffectedComponentEntity(); + componentEntity.setId(liveConnectable.id()); + componentEntity.setReferenceType(referenceType); + componentEntity.setComponent(componentDto); + return componentEntity; + } + + private String getReferenceType(final ConnectableType connectableType) { + if (connectableType == ConnectableType.PROCESSOR) { + return AffectedComponentDTO.COMPONENT_TYPE_PROCESSOR; + } + if (connectableType == ConnectableType.INPUT_PORT) { + return AffectedComponentDTO.COMPONENT_TYPE_INPUT_PORT; + } + if (connectableType == ConnectableType.OUTPUT_PORT) { + return AffectedComponentDTO.COMPONENT_TYPE_OUTPUT_PORT; + } + + return null; + } + + private String getState(final RemovedConnectionDrainClassifier.LiveConnectable liveConnectable) { + if (liveConnectable.type() == ConnectableType.PROCESSOR) { + return liveConnectable.physicalScheduledState() == null ? null : liveConnectable.physicalScheduledState().name(); + } + + return liveConnectable.running() ? "RUNNING" : "STOPPED"; + } + + private LifecycleManagementException decorateFailure(final LifecycleManagementException failure, + final Set componentsToStop, + final Set candidateConnectionIds) { + final String message = failure.getMessage(); + if (message != null && message.startsWith("Removed connection drain")) { + return failure; + } + + final String decoratedMessage; + if (message != null && message.contains("waiting for connection queues")) { + decoratedMessage = "Removed connection drain failed while waiting for connections " + + candidateConnectionIds.stream().sorted().toList() + ": " + message; + } else { + final List componentIds = componentsToStop.stream().map(AffectedComponentEntity::getId).sorted().toList(); + decoratedMessage = "Removed connection drain failed while stopping producer barriers " + componentIds + ": " + message; + } + + return new LifecycleManagementException(decoratedMessage, failure); + } + + private void restoreOrSuppress(final ComponentLifecycle componentLifecycle, final URI requestUri, final String groupId, + final Set stoppedComponents, final LifecycleManagementException failure) { + try { + restoreStoppedComponents(componentLifecycle, requestUri, groupId, stoppedComponents); + } catch (final LifecycleManagementException restorationFailure) { + logger.warn("Failed to restore producer barriers {} after removed connection drain failure", + stoppedComponents.stream().map(AffectedComponentEntity::getId).sorted().toList(), restorationFailure); + failure.addSuppressed(restorationFailure); + } + } + + private void restoreStoppedComponents(final ComponentLifecycle componentLifecycle, final URI requestUri, final String groupId, + final Set stoppedComponents) throws LifecycleManagementException { + if (stoppedComponents.isEmpty()) { + return; + } + + componentLifecycle.scheduleComponents(requestUri, groupId, stoppedComponents, ScheduledState.RUNNING, + pauseFactory.createRestorationPause(), InvalidComponentAction.SKIP); + } + + private RemovedConnectionDrainClassifier.Context createQueueAwareContext(final FlowUpdateImpact flowUpdateImpact, + final RemovedConnectionDrainClassifier.Context context, + final ComponentLifecycle componentLifecycle, + final URI requestUri) throws LifecycleManagementException { + final Map knownQueueEmptyByConnectionId = new LinkedHashMap<>(); + final Pause noWaitPause = NoWaitPause.INSTANCE; + for (final RemovedConnectionDescriptor removedConnection : flowUpdateImpact.getRemovedConnections()) { + final String connectionId = removedConnection.getConnectionInstanceId(); + if (connectionId == null) { + continue; + } + + final boolean knownQueueEmpty = componentLifecycle.waitForConnectionQueuesEmpty(requestUri, Set.of(connectionId), noWaitPause); + knownQueueEmptyByConnectionId.put(connectionId, knownQueueEmpty); + } + + return new QueueAwareContext(context, knownQueueEmptyByConnectionId); + } + + private boolean allComponentsStopped(final Set componentsToStop, final Set updatedStoppedComponents) { + if (componentsToStop.isEmpty()) { + return true; + } + + final Map updatedComponentsById = toOrderedMap(updatedStoppedComponents); + for (final AffectedComponentEntity componentToStop : componentsToStop) { + final AffectedComponentEntity updatedComponent = updatedComponentsById.getOrDefault(componentToStop.getId(), componentToStop); + if (updatedComponent.getComponent() == null || isActive(updatedComponent.getComponent())) { + return false; + } + } + + return true; + } + + private Set getStoppedComponents(final Collection componentsToStop, + final Collection updatedComponents) { + if (componentsToStop == null || componentsToStop.isEmpty()) { + return Collections.emptySet(); + } + + final Map updatedComponentsById = toOrderedMap(updatedComponents); + final Set stoppedComponents = new LinkedHashSet<>(); + for (final AffectedComponentEntity componentToStop : componentsToStop) { + final AffectedComponentEntity updatedComponent = updatedComponentsById.getOrDefault(componentToStop.getId(), componentToStop); + if (updatedComponent.getComponent() != null && !isActive(updatedComponent.getComponent())) { + stoppedComponents.add(updatedComponent); + } + } + + return stoppedComponents; + } + + private Set getStoppedComponentsToRestore(final RemovedConnectionDrainClassifier.Context context, + final Collection componentsToStop, + final Collection updatedComponents) { + if (componentsToStop == null || componentsToStop.isEmpty()) { + return Collections.emptySet(); + } + + final Map updatedComponentsById = toOrderedMap(updatedComponents); + final Set stoppedComponents = new LinkedHashSet<>(); + for (final AffectedComponentEntity componentToStop : componentsToStop) { + final AffectedComponentEntity updatedComponent = updatedComponentsById.get(componentToStop.getId()); + if (updatedComponent != null) { + if (updatedComponent.getComponent() != null && !isActive(updatedComponent.getComponent())) { + stoppedComponents.add(updatedComponent); + } + continue; + } + + final RemovedConnectionDrainClassifier.LiveConnectable liveConnectable = context.getConnectable(componentToStop.getId()); + if (liveConnectable == null || isActive(liveConnectable)) { + continue; + } + + final AffectedComponentEntity liveComponent = getProducerBarrierEntity(Collections.emptyMap(), context, componentToStop.getId()); + if (liveComponent != null && liveComponent.getComponent() != null) { + stoppedComponents.add(liveComponent); + } + } + + return stoppedComponents; + } + + private Map toOrderedMap(final Collection components) { + if (components == null || components.isEmpty()) { + return Collections.emptyMap(); + } + + final Map orderedMap = new LinkedHashMap<>(); + for (final AffectedComponentEntity component : components) { + orderedMap.put(component.getId(), component); + } + return orderedMap; + } + + private boolean isActive(final AffectedComponentDTO affectedComponentDto) { + final String state = affectedComponentDto.getState(); + if ("Running".equalsIgnoreCase(state) || "Starting".equalsIgnoreCase(state)) { + return true; + } + + final Integer activeThreadCount = affectedComponentDto.getActiveThreadCount(); + return activeThreadCount != null && activeThreadCount > 0; + } + + private boolean isActive(final RemovedConnectionDrainClassifier.LiveConnectable liveConnectable) { + if (liveConnectable == null) { + return false; + } + + if (liveConnectable.type() == ConnectableType.PROCESSOR) { + return liveConnectable.physicalScheduledState() == ScheduledState.RUNNING + || liveConnectable.physicalScheduledState() == ScheduledState.STARTING; + } + + return liveConnectable.running(); + } + + public interface CancellationHandle { + boolean isCancelled(); + + void setCancelCallback(Runnable runnable); + } + + interface PauseFactory { + DeadlinePause createDrainPause(Duration timeout); + + Pause createRestorationPause(); + } + + interface DeadlinePause extends Pause { + void cancel(); + } + + public record DrainResult(Set candidateConnectionIds, Set drainStoppedComponents, boolean cancelled, + LifecycleManagementException restorationFailure) { + public DrainResult { + candidateConnectionIds = copyOrderedSet(candidateConnectionIds); + drainStoppedComponents = copyOrderedSet(drainStoppedComponents); + } + + static DrainResult success(final Set candidateConnectionIds, final Set drainStoppedComponents) { + return new DrainResult(candidateConnectionIds, drainStoppedComponents, false, null); + } + + static DrainResult cancelled(final Set candidateConnectionIds, final Set drainStoppedComponents) { + return cancelled(candidateConnectionIds, drainStoppedComponents, null); + } + + static DrainResult cancelled(final Set candidateConnectionIds, final Set drainStoppedComponents, + final LifecycleManagementException restorationFailure) { + return new DrainResult(candidateConnectionIds, drainStoppedComponents, true, restorationFailure); + } + } + + private static final class MonotonicPauseFactory implements PauseFactory { + private final Duration pollInterval; + private final LongSupplier nanoTimeSupplier; + + private MonotonicPauseFactory(final Duration pollInterval, final LongSupplier nanoTimeSupplier) { + this.pollInterval = pollInterval; + this.nanoTimeSupplier = nanoTimeSupplier; + } + + @Override + public DeadlinePause createDrainPause(final Duration timeout) { + return new TimedDeadlinePause(pollInterval, timeout, nanoTimeSupplier); + } + + @Override + public Pause createRestorationPause() { + return new CancellableTimedPause(pollInterval.toMillis(), Long.MAX_VALUE, TimeUnit.MILLISECONDS); + } + } + + private static final class TimedDeadlinePause implements DeadlinePause { + private final long pauseNanos; + private final long deadlineNanos; + private final LongSupplier nanoTimeSupplier; + private volatile boolean cancelled; + + private TimedDeadlinePause(final Duration pollInterval, final Duration timeout, final LongSupplier nanoTimeSupplier) { + this.pauseNanos = Math.max(1L, pollInterval.toNanos()); + this.nanoTimeSupplier = nanoTimeSupplier; + this.deadlineNanos = nanoTimeSupplier.getAsLong() + timeout.toNanos(); + } + + @Override + public void cancel() { + cancelled = true; + } + + @Override + public boolean pause() { + if (cancelled) { + return false; + } + + final long now = nanoTimeSupplier.getAsLong(); + if (now >= deadlineNanos) { + return false; + } + + try { + TimeUnit.NANOSECONDS.sleep(Math.min(pauseNanos, Math.max(1L, deadlineNanos - now))); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + return false; + } + + return !cancelled && nanoTimeSupplier.getAsLong() < deadlineNanos; + } + } + + private enum NoWaitPause implements Pause { + INSTANCE; + + @Override + public boolean pause() { + return false; + } + } + + private record QueueAwareContext(RemovedConnectionDrainClassifier.Context delegate, + Map knownQueueEmptyByConnectionId) implements RemovedConnectionDrainClassifier.Context { + private QueueAwareContext { + delegate = Objects.requireNonNull(delegate, "Removed Connection Drain Context required"); + knownQueueEmptyByConnectionId = Collections.unmodifiableMap(new LinkedHashMap<>(knownQueueEmptyByConnectionId)); + } + + @Override + public RemovedConnectionDrainClassifier.LiveConnection getConnection(final String connectionId) { + final RemovedConnectionDrainClassifier.LiveConnection connection = delegate.getConnection(connectionId); + if (connection == null) { + return null; + } + + final Boolean knownQueueEmpty = knownQueueEmptyByConnectionId.get(connectionId); + if (knownQueueEmpty == null) { + return connection; + } + + return new RemovedConnectionDrainClassifier.LiveConnection(connection.id(), connection.sourceId(), connection.destinationId(), knownQueueEmpty); + } + + @Override + public RemovedConnectionDrainClassifier.LiveConnectable getConnectable(final String connectableId) { + return delegate.getConnectable(connectableId); + } + + @Override + public RemovedConnectionDrainClassifier.LiveProcessGroup getProcessGroup(final String processGroupId) { + return delegate.getProcessGroup(processGroupId); + } + } + + private static Set copyOrderedSet(final Collection values) { + if (values == null || values.isEmpty()) { + return Collections.emptySet(); + } + + return Collections.unmodifiableSet(new LinkedHashSet<>(values)); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/StandardNiFiServiceFacade.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/StandardNiFiServiceFacade.java index 6688dd48152a..65c135745613 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/StandardNiFiServiceFacade.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/StandardNiFiServiceFacade.java @@ -95,6 +95,7 @@ import org.apache.nifi.components.validation.ValidationState; import org.apache.nifi.components.validation.ValidationStatus; import org.apache.nifi.connectable.Connectable; +import org.apache.nifi.connectable.ConnectableType; import org.apache.nifi.connectable.Connection; import org.apache.nifi.connectable.Funnel; import org.apache.nifi.connectable.Port; @@ -134,6 +135,8 @@ import org.apache.nifi.diagnostics.SystemDiagnostics; import org.apache.nifi.events.BulletinFactory; import org.apache.nifi.expression.ExpressionLanguageScope; +import org.apache.nifi.flow.ConnectableComponent; +import org.apache.nifi.flow.ConnectableComponentType; import org.apache.nifi.flow.ExecutionEngine; import org.apache.nifi.flow.ExternalControllerServiceReference; import org.apache.nifi.flow.ParameterProviderReference; @@ -7153,24 +7156,21 @@ public void verifyCanRevertLocalModifications(final String groupId, final Regist @Override public Set getComponentsAffectedByFlowUpdate(final String processGroupId, final RegisteredFlowSnapshot updatedSnapshot) { - final ProcessGroup group = processGroupDAO.getProcessGroup(processGroupId); - - final VersionedComponentFlowMapper mapper = makeNiFiRegistryFlowMapper(controllerFacade.getExtensionManager()); - final VersionedProcessGroup localContents = mapper.mapProcessGroup(group, controllerFacade.getControllerServiceProvider(), controllerFacade.getFlowManager(), true); - - final ComparableDataFlow localFlow = new StandardComparableDataFlow("Current Flow", localContents); - final ComparableDataFlow proposedFlow = new StandardComparableDataFlow("New Flow", updatedSnapshot.getFlowContents()); + return getFlowUpdateImpact(processGroupId, updatedSnapshot).getAffectedComponents(); + } - final FlowComparator flowComparator = new StandardFlowComparator(localFlow, proposedFlow, new StaticDifferenceDescriptor(), - Function.identity(), VersionedComponent::getIdentifier, FlowComparatorVersionedStrategy.DEEP); - final FlowComparison comparison = flowComparator.compare(); + @Override + public FlowUpdateImpact getFlowUpdateImpact(final String processGroupId, final RegisteredFlowSnapshot updatedSnapshot) { + final ProcessGroup group = processGroupDAO.getProcessGroup(processGroupId); + final FlowComparison comparison = compareFlowUpdate(group, updatedSnapshot); + final VersionedProcessGroup proposedContents = updatedSnapshot.getFlowContents(); final FlowManager flowManager = controllerFacade.getFlowManager(); final Set affectedComponents = comparison.getDifferences().stream() .filter(difference -> difference.getDifferenceType() != DifferenceType.COMPONENT_ADDED) // components that are added are not components that will be affected in the local flow. .filter(FlowDifferenceFilters.FILTER_ADDED_REMOVED_REMOTE_PORTS) .filter(difference -> difference.getComponentA() != null) // a difference that would not affect a local component - .filter(diff -> FlowDifferenceFilters.isComponentUpdateRequired(diff, proposedFlow.getContents(), flowManager)) + .filter(diff -> FlowDifferenceFilters.isComponentUpdateRequired(diff, proposedContents, flowManager)) // A local rename of a public port is preserved during a version-control update (it is not overwritten with the // registry name), so the port must not be reported as affected/stopped for that name change. Applied unconditionally here because // this affected-components calculation serves only the version-control update path. @@ -7242,6 +7242,10 @@ public Set getComponentsAffectedByFlowUpdate(final Stri }) .collect(Collectors.toCollection(HashSet::new)); + final Set removedConnections = new LinkedHashSet<>(); + final Set removedProcessGroupIds = new LinkedHashSet<>(); + final Set removedEndpointIds = new LinkedHashSet<>(); + for (final FlowDifference difference : comparison.getDifferences()) { // Ignore differences for adding remote ports if (FlowDifferenceFilters.isAddedOrRemovedRemotePort(difference)) { @@ -7266,6 +7270,8 @@ public Set getComponentsAffectedByFlowUpdate(final Stri continue; } + addRemovedFlowUpdateImpact(difference, removedConnections, removedProcessGroupIds, removedEndpointIds); + // If any Process Group is removed, consider all components below that Process Group as an affected component if (difference.getDifferenceType() == DifferenceType.COMPONENT_REMOVED && localComponent.getComponentType() == org.apache.nifi.flow.ComponentType.PROCESS_GROUP) { final String localGroupId = localComponent.getInstanceIdentifier(); @@ -7353,7 +7359,24 @@ public Set getComponentsAffectedByFlowUpdate(final Stri } } - return affectedComponents; + return new FlowUpdateImpact(affectedComponents, removedConnections, removedProcessGroupIds, removedEndpointIds); + } + + @Override + public RemovedConnectionDrainClassifier.Context getRemovedConnectionDrainContext() { + return new RemovedConnectionDrainClassifier.FlowManagerContext(controllerFacade.getFlowManager()); + } + + FlowComparison compareFlowUpdate(final ProcessGroup group, final RegisteredFlowSnapshot updatedSnapshot) { + final VersionedComponentFlowMapper mapper = makeNiFiRegistryFlowMapper(controllerFacade.getExtensionManager()); + final VersionedProcessGroup localContents = mapper.mapProcessGroup(group, controllerFacade.getControllerServiceProvider(), controllerFacade.getFlowManager(), true); + + final ComparableDataFlow localFlow = new StandardComparableDataFlow("Current Flow", localContents); + final ComparableDataFlow proposedFlow = new StandardComparableDataFlow("New Flow", updatedSnapshot.getFlowContents()); + + final FlowComparator flowComparator = new StandardFlowComparator(localFlow, proposedFlow, new StaticDifferenceDescriptor(), + Function.identity(), VersionedComponent::getIdentifier, FlowComparatorVersionedStrategy.DEEP); + return flowComparator.compare(); } private Port getInputPort(final InstantiatedVersionedPort port) { @@ -7374,6 +7397,76 @@ private Port getOutputPort(final InstantiatedVersionedPort port) { return processGroup.getOutputPort(port.getInstanceIdentifier()); } + private void addRemovedFlowUpdateImpact(final FlowDifference difference, final Set removedConnections, + final Set removedProcessGroupIds, final Set removedEndpointIds) { + final DifferenceType differenceType = difference.getDifferenceType(); + if (differenceType == DifferenceType.COMPONENT_REMOVED) { + final VersionedComponent localComponent = difference.getComponentA(); + if (localComponent instanceof final VersionedConnection removedConnection) { + removedConnections.add(createRemovedConnectionDescriptor(removedConnection, RemovalReason.COMPONENT_REMOVED)); + } else if (localComponent.getComponentType() == org.apache.nifi.flow.ComponentType.PROCESS_GROUP) { + removedProcessGroupIds.add(localComponent.getInstanceIdentifier()); + } else if (isRemovedEndpoint(localComponent)) { + removedEndpointIds.add(localComponent.getInstanceIdentifier()); + } + } else if (differenceType == DifferenceType.SOURCE_CHANGED && difference.getComponentA() instanceof final VersionedConnection removedConnection) { + removedConnections.add(createRemovedConnectionDescriptor(removedConnection, RemovalReason.SOURCE_CHANGED)); + } + } + + private boolean isRemovedEndpoint(final VersionedComponent component) { + return component.getComponentType() == org.apache.nifi.flow.ComponentType.INPUT_PORT + || component.getComponentType() == org.apache.nifi.flow.ComponentType.OUTPUT_PORT + || component.getComponentType() == org.apache.nifi.flow.ComponentType.PROCESSOR + || component.getComponentType() == org.apache.nifi.flow.ComponentType.FUNNEL + || component.getComponentType() == org.apache.nifi.flow.ComponentType.REMOTE_INPUT_PORT + || component.getComponentType() == org.apache.nifi.flow.ComponentType.REMOTE_OUTPUT_PORT; + } + + private RemovedConnectionDescriptor createRemovedConnectionDescriptor(final VersionedConnection connection, final RemovalReason removalReason) { + final ConnectableComponent source = connection.getSource(); + final ConnectableComponent destination = connection.getDestination(); + + return new RemovedConnectionDescriptor( + connection.getInstanceIdentifier(), + connection.getIdentifier(), + getComponentGroupRuntimeId(connection), + source == null ? null : source.getInstanceIdentifier(), + source == null ? null : source.getId(), + getConnectableGroupRuntimeId(source), + getConnectableType(source), + destination == null ? null : destination.getInstanceIdentifier(), + destination == null ? null : destination.getId(), + getConnectableGroupRuntimeId(destination), + getConnectableType(destination), + removalReason); + } + + private String getComponentGroupRuntimeId(final VersionedComponent component) { + if (component instanceof final InstantiatedVersionedComponent instantiatedComponent) { + return instantiatedComponent.getInstanceGroupId(); + } + + return component.getGroupIdentifier(); + } + + private String getConnectableGroupRuntimeId(final ConnectableComponent connectableComponent) { + if (connectableComponent instanceof final InstantiatedVersionedComponent instantiatedComponent) { + return instantiatedComponent.getInstanceGroupId(); + } + + return connectableComponent == null ? null : connectableComponent.getGroupId(); + } + + private ConnectableType getConnectableType(final ConnectableComponent connectableComponent) { + if (connectableComponent == null || connectableComponent.getType() == null) { + return null; + } + + final ConnectableComponentType type = connectableComponent.getType(); + return ConnectableType.valueOf(type.name()); + } + private void mapToConnectableId(final Collection connectables, final Map> destination) { for (final Connectable connectable : connectables) { final Optional versionedIdOption = connectable.getVersionedComponentId(); diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/FlowUpdateResource.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/FlowUpdateResource.java index 143b784605d8..fed169aa5de2 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/FlowUpdateResource.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/FlowUpdateResource.java @@ -35,7 +35,9 @@ import org.apache.nifi.flow.VersionedParameterContext; import org.apache.nifi.registry.flow.FlowSnapshotContainer; import org.apache.nifi.registry.flow.RegisteredFlowSnapshot; +import org.apache.nifi.web.FlowUpdateImpact; import org.apache.nifi.web.NiFiServiceFacade; +import org.apache.nifi.web.RemovedConnectionDrainCoordinator; import org.apache.nifi.web.ResourceNotFoundException; import org.apache.nifi.web.ResumeFlowException; import org.apache.nifi.web.Revision; @@ -90,6 +92,7 @@ */ public abstract class FlowUpdateResource extends ApplicationResource { private static final String DISABLED_COMPONENT_STATE = "DISABLED"; + protected static final String UPDATE_REQUEST_TYPE = "update-requests"; private static final Logger logger = LoggerFactory.getLogger(FlowUpdateResource.class); protected NiFiServiceFacade serviceFacade; @@ -180,17 +183,18 @@ protected Response initiateFlowUpdate(final String groupId, final T requestEntit // 2. Verify READ and WRITE permissions for user, for every component. // 3. Verify that all components in the snapshot exist on all nodes (i.e., the NAR exists)? // 4: Verify that Process Group can be updated. Only versioned flows care about the verifyNotDirty flag. - // 5. Stop all Processors, Funnels, Ports that are affected. - // 6. Wait for all of the components to finish stopping. - // 7. Disable all Controller Services that are affected. - // 8. Wait for all Controller Services to finish disabling. - // 9. Ensure that if any connection was deleted, that it has no data in it. Ensure that no Input Port + // 5. Drain non-empty removed Connections after stopping their producer barriers. + // 6. Stop all Processors, Funnels, Ports that are affected. + // 7. Wait for all of the components to finish stopping. + // 8. Disable all Controller Services that are affected. + // 9. Wait for all Controller Services to finish disabling. + // 10. Ensure that if any connection was deleted, that it has no data in it. Ensure that no Input Port // was removed, unless it currently has no incoming connections. Ensure that no Output Port was removed, // unless it currently has no outgoing connections. Checking ports & connections could be done before // stopping everything, but removal of Connections cannot. - // 10.-11. Update components in the Process Group; update Version Control Information (registry version change only). - // 12. Re-Enable all affected Controller Services that were not removed. - // 13. Re-Start all Processors, Funnels, Ports that are affected and not removed. + // 11.-12. Update components in the Process Group; update Version Control Information (registry version change only). + // 13. Re-Enable all affected Controller Services that were not removed. + // 14. Re-Start all Processors, Funnels, Ports that are affected and not removed. // Step 0: Obtain the versioned flow snapshot to use for the update final FlowSnapshotContainer flowSnapshotContainer = flowSnapshotContainerSupplier.get(); @@ -208,12 +212,12 @@ protected Response initiateFlowUpdate(final String groupId, final T requestEntit final Set unresolvedParameterProviders = serviceFacade.resolveParameterProviders(flowSnapshot, user); // Step 1: Determine which components will be affected by updating the flow - final Set affectedComponents = serviceFacade.getComponentsAffectedByFlowUpdate(groupId, flowSnapshot); + final FlowUpdateImpact flowUpdateImpact = serviceFacade.getFlowUpdateImpact(groupId, flowSnapshot); // build a request wrapper final InitiateUpdateFlowRequestWrapper requestWrapper = new InitiateUpdateFlowRequestWrapper(requestEntity, componentLifecycle, requestType, getAbsolutePath(), replicateUriPath, - affectedComponents, replicateRequest, flowSnapshot); + flowUpdateImpact, replicateRequest, flowSnapshot); final Revision requestRevision = getRevision(revisionDto, groupId); return withWriteLock( @@ -284,14 +288,14 @@ protected Response submitFlowUpdateRequest(final NiFiUser user, final String gro // result in stopping components, which can take an indeterminate amount of time. final String requestId = UUID.randomUUID().toString(); final AsynchronousWebRequest request = - new StandardAsynchronousWebRequest<>(requestId, wrapper.getRequestEntity(), groupId, user, getUpdateFlowSteps()); + new StandardAsynchronousWebRequest<>(requestId, wrapper.getRequestEntity(), groupId, user, getUpdateFlowSteps(requestType)); // Submit the request to be performed in the background final Consumer> updateTask = vcur -> { try { updateFlow(groupId, wrapper.getComponentLifecycle(), wrapper.getRequestUri(), - wrapper.getAffectedComponents(), wrapper.isReplicateRequest(), wrapper.getReplicateUriPath(), + wrapper.getFlowUpdateImpact(), wrapper.isReplicateRequest(), wrapper.getReplicateUriPath(), revision, wrapper.getRequestEntity(), wrapper.getFlowSnapshot(), request, idGenerationSeed, allowDirtyFlowUpdate, requestType); @@ -331,13 +335,15 @@ private boolean isActive(final AffectedComponentDTO affectedComponentDto) { * Perform the specified flow update */ private void updateFlow(final String groupId, final ComponentLifecycle componentLifecycle, final URI requestUri, - final Set affectedComponents, final boolean replicateRequest, + final FlowUpdateImpact flowUpdateImpact, final boolean replicateRequest, final String replicateUriPath, final Revision revision, final T requestEntity, final RegisteredFlowSnapshot flowSnapshot, final AsynchronousWebRequest asyncRequest, final String idGenerationSeed, final boolean allowDirtyFlowUpdate, final String requestType) throws LifecycleManagementException, ResumeFlowException { - // Steps 5-6: Determine which components must be stopped and stop them. + final Set affectedComponents = flowUpdateImpact.getAffectedComponents(); + + // Steps 5-7: Drain removed connections, determine which components must be stopped, and stop them. final Set stoppableReferenceTypes = new HashSet<>(); stoppableReferenceTypes.add(AffectedComponentDTO.COMPONENT_TYPE_PROCESSOR); stoppableReferenceTypes.add(AffectedComponentDTO.COMPONENT_TYPE_REMOTE_INPUT_PORT); @@ -349,7 +355,21 @@ private void updateFlow(final String groupId, final ComponentLifecycle component final Set runningComponents = affectedComponents.stream() .filter(entity -> stoppableReferenceTypes.contains(entity.getComponent().getReferenceType())) .filter(entity -> isActive(entity.getComponent())) - .collect(Collectors.toSet()); + .collect(Collectors.toCollection(LinkedHashSet::new)); + + if (UPDATE_REQUEST_TYPE.equals(requestType)) { + final RemovedConnectionDrainCoordinator.DrainResult drainResult = preDrainRemovedConnections( + flowUpdateImpact, componentLifecycle, requestUri, groupId, asyncRequest); + if (drainResult.cancelled()) { + if (drainResult.restorationFailure() != null) { + logger.warn("Failed to restore components after removed connection drain cancellation", drainResult.restorationFailure()); + asyncRequest.appendFailureDetail("restoration failed: " + drainResult.restorationFailure().getMessage()); + } + return; + } + runningComponents.addAll(drainResult.drainStoppedComponents()); + asyncRequest.markStepComplete(); + } logger.info("Stopping {} Processors", runningComponents.size()); final CancellableTimedPause stopComponentsPause = new CancellableTimedPause(250, Long.MAX_VALUE, TimeUnit.MILLISECONDS); @@ -361,7 +381,7 @@ private void updateFlow(final String groupId, final ComponentLifecycle component } asyncRequest.markStepComplete(); - // Steps 7-8. Disable enabled controller services that are affected. + // Steps 8-9. Disable enabled controller services that are affected. // We don't want to disable services that are already disabling. But we need to wait for their state to transition from Disabling to Disabled. final Set servicesToWaitFor = affectedComponents.stream() .filter(dto -> AffectedComponentDTO.COMPONENT_TYPE_CONTROLLER_SERVICE.equals(dto.getComponent().getReferenceType())) @@ -621,8 +641,11 @@ private NodeResponse replicateFlowUpdateRequest(final URI replicateUri, final Ni /** * Get a list of steps to perform for upload flow */ - private static List getUpdateFlowSteps() { + static List getUpdateFlowSteps(final String requestType) { final List updateSteps = new ArrayList<>(); + if (UPDATE_REQUEST_TYPE.equals(requestType)) { + updateSteps.add(new StandardUpdateStep("Draining Removed Connections")); + } updateSteps.add(new StandardUpdateStep("Stopping Affected Processors")); updateSteps.add(new StandardUpdateStep("Disabling Affected Controller Services")); updateSteps.add(new StandardUpdateStep("Updating Flow")); @@ -631,6 +654,24 @@ private static List getUpdateFlowSteps() { return updateSteps; } + protected RemovedConnectionDrainCoordinator.DrainResult preDrainRemovedConnections( + final FlowUpdateImpact flowUpdateImpact, final ComponentLifecycle componentLifecycle, final URI requestUri, + final String groupId, final AsynchronousWebRequest asyncRequest) throws LifecycleManagementException { + return new RemovedConnectionDrainCoordinator().coordinateDrain( + flowUpdateImpact, serviceFacade.getRemovedConnectionDrainContext(), componentLifecycle, requestUri, groupId, + new RemovedConnectionDrainCoordinator.CancellationHandle() { + @Override + public boolean isCancelled() { + return asyncRequest.isCancelled(); + } + + @Override + public void setCancelCallback(final Runnable runnable) { + asyncRequest.setCancelCallback(runnable); + } + }); + } + /** * Extracts the response entity from the specified node response. * @@ -773,20 +814,20 @@ protected class InitiateUpdateFlowRequestWrapper extends Entity { private final String requestType; private final URI requestUri; private final String replicateUriPath; - private final Set affectedComponents; + private final FlowUpdateImpact flowUpdateImpact; private final boolean replicateRequest; private final RegisteredFlowSnapshot flowSnapshot; public InitiateUpdateFlowRequestWrapper(final T requestEntity, final ComponentLifecycle componentLifecycle, final String requestType, final URI requestUri, final String replicateUriPath, - final Set affectedComponents, + final FlowUpdateImpact flowUpdateImpact, final boolean replicateRequest, final RegisteredFlowSnapshot flowSnapshot) { this.requestEntity = requestEntity; this.componentLifecycle = componentLifecycle; this.requestType = requestType; this.requestUri = requestUri; this.replicateUriPath = replicateUriPath; - this.affectedComponents = affectedComponents; + this.flowUpdateImpact = flowUpdateImpact; this.replicateRequest = replicateRequest; this.flowSnapshot = flowSnapshot; } @@ -812,7 +853,11 @@ public String getReplicateUriPath() { } public Set getAffectedComponents() { - return affectedComponents; + return flowUpdateImpact.getAffectedComponents(); + } + + public FlowUpdateImpact getFlowUpdateImpact() { + return flowUpdateImpact; } public boolean isReplicateRequest() { diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/VersionsResource.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/VersionsResource.java index 6bbaf98b676a..8032142f6bc6 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/VersionsResource.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/VersionsResource.java @@ -57,12 +57,12 @@ import org.apache.nifi.registry.flow.RegisteredFlowSnapshot; import org.apache.nifi.registry.flow.RegisteredFlowSnapshotMetadata; import org.apache.nifi.registry.flow.VersionedFlowState; +import org.apache.nifi.web.FlowUpdateImpact; import org.apache.nifi.web.Revision; import org.apache.nifi.web.api.dto.RevisionDTO; import org.apache.nifi.web.api.dto.VersionControlInformationDTO; import org.apache.nifi.web.api.dto.VersionedFlowDTO; import org.apache.nifi.web.api.dto.VersionedFlowUpdateRequestDTO; -import org.apache.nifi.web.api.entity.AffectedComponentEntity; import org.apache.nifi.web.api.entity.CreateActiveRequestEntity; import org.apache.nifi.web.api.entity.CreateFlowBranchRequestEntity; import org.apache.nifi.web.api.entity.Entity; @@ -1085,7 +1085,7 @@ public Response applyRebasedFlowVersion( public Response getUpdateRequest( @Parameter(description = "The ID of the Update Request") @PathParam("id") final String updateRequestId) { - return retrieveFlowUpdateRequest("update-requests", updateRequestId); + return retrieveFlowUpdateRequest(UPDATE_REQUEST_TYPE, updateRequestId); } @GET @@ -1145,7 +1145,7 @@ public Response deleteUpdateRequest( @QueryParam(DISCONNECTED_NODE_ACKNOWLEDGED) @DefaultValue("false") final Boolean disconnectedNodeAcknowledged, @Parameter(description = "The ID of the Update Request") @PathParam("id") final String updateRequestId) { - return deleteFlowUpdateRequest("update-requests", updateRequestId, disconnectedNodeAcknowledged); + return deleteFlowUpdateRequest(UPDATE_REQUEST_TYPE, updateRequestId, disconnectedNodeAcknowledged); } @DELETE @@ -1423,7 +1423,7 @@ public Response initiateVersionControlUpdate( } // supplier retrieves Versioned Flow Snapshot from the Flow Registry - return initiateFlowUpdate(groupId, requestEntity, false, "update-requests", + return initiateFlowUpdate(groupId, requestEntity, false, UPDATE_REQUEST_TYPE, "/nifi-api/versions/process-groups/" + groupId, () -> serviceFacade.getVersionedFlowSnapshot(requestVersionControlInfoDto, true) ); @@ -1523,12 +1523,12 @@ public Response initiateRevertFlowVersion( final Set unresolvedParameterProviders = serviceFacade.resolveParameterProviders(flowSnapshot, NiFiUserUtils.getNiFiUser()); // Step 1: Determine which components will be affected by updating the version - final Set affectedComponents = serviceFacade.getComponentsAffectedByFlowUpdate(groupId, flowSnapshot); + final FlowUpdateImpact flowUpdateImpact = serviceFacade.getFlowUpdateImpact(groupId, flowSnapshot); // build a request wrapper final InitiateUpdateFlowRequestWrapper requestWrapper = new InitiateUpdateFlowRequestWrapper(requestEntity, componentLifecycle, "revert-requests", getAbsolutePath(), - "/nifi-api/versions/process-groups/" + groupId, affectedComponents, replicateRequest, flowSnapshot); + "/nifi-api/versions/process-groups/" + groupId, flowUpdateImpact, replicateRequest, flowSnapshot); final Revision requestRevision = getRevision(requestEntity.getProcessGroupRevision(), groupId); return withWriteLock( diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/AsynchronousWebRequest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/AsynchronousWebRequest.java index 9e21e4ccefce..52b03087ed49 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/AsynchronousWebRequest.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/AsynchronousWebRequest.java @@ -68,6 +68,13 @@ public interface AsynchronousWebRequest { */ void fail(String explanation); + /** + * Appends detail to the existing failure reason, or establishes the detail as the failure reason when none exists. + * + * @param detail additional failure detail + */ + void appendFailureDetail(String detail); + /** * Indicates the reason that the request failed, or null if the request has not failed * diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequest.java index 02a77552c5b3..5d4f9db943e1 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequest.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequest.java @@ -156,6 +156,12 @@ public synchronized void fail(final String explanation) { applyFailure(explanation); } + @Override + public synchronized void appendFailureDetail(final String detail) { + final String failureDetail = Objects.requireNonNull(detail); + applyFailure(failureReason == null ? failureDetail : failureReason + "; " + failureDetail); + } + private void applyFailure(final String explanation) { this.failureReason = Objects.requireNonNull(explanation); this.complete = true; diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycle.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycle.java index c96a7f5aa954..13bbed9ba527 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycle.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycle.java @@ -26,7 +26,9 @@ import org.apache.nifi.authorization.user.NiFiUser; import org.apache.nifi.authorization.user.NiFiUserUtils; import org.apache.nifi.cluster.coordination.ClusterCoordinator; +import org.apache.nifi.cluster.coordination.http.replication.AsyncClusterResponse; import org.apache.nifi.cluster.coordination.http.replication.RequestReplicator; +import org.apache.nifi.cluster.coordination.node.NodeConnectionState; import org.apache.nifi.cluster.exception.NoClusterCoordinatorException; import org.apache.nifi.cluster.manager.NodeResponse; import org.apache.nifi.cluster.protocol.NodeIdentifier; @@ -39,6 +41,7 @@ import org.apache.nifi.web.api.dto.AffectedComponentDTO; import org.apache.nifi.web.api.dto.ControllerServiceDTO; import org.apache.nifi.web.api.dto.DtoFactory; +import org.apache.nifi.web.api.dto.ListingRequestDTO; import org.apache.nifi.web.api.dto.ProcessorRunStatusDetailsDTO; import org.apache.nifi.web.api.dto.RevisionDTO; import org.apache.nifi.web.api.dto.status.ProcessGroupStatusSnapshotDTO; @@ -47,6 +50,7 @@ import org.apache.nifi.web.api.entity.ComponentEntity; import org.apache.nifi.web.api.entity.ControllerServiceEntity; import org.apache.nifi.web.api.entity.ControllerServicesEntity; +import org.apache.nifi.web.api.entity.ListingRequestEntity; import org.apache.nifi.web.api.entity.ProcessGroupEntity; import org.apache.nifi.web.api.entity.ProcessGroupStatusEntity; import org.apache.nifi.web.api.entity.ProcessorRunStatusDetailsEntity; @@ -60,9 +64,13 @@ import java.net.URISyntaxException; import java.util.Collection; import java.util.Collections; +import java.util.Comparator; import java.util.HashMap; +import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; @@ -631,6 +639,57 @@ public Set activateControllerServices(final URI origina .collect(Collectors.toSet()); } + @Override + public boolean waitForConnectionQueuesEmpty(final URI originalUri, final Set connectionIds, final Pause pause) throws LifecycleManagementException { + if (connectionIds.isEmpty()) { + return true; + } + + final Set expectedNodes = new HashSet<>(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)); + final List orderedConnectionIds = connectionIds.stream() + .sorted(Comparator.naturalOrder()) + .toList(); + final NiFiUser user = NiFiUserUtils.getNiFiUser(); + final Map queuePollResults = new LinkedHashMap<>(); + + boolean continuePolling = true; + while (continuePolling) { + boolean allQueuesEmpty = true; + for (final String connectionId : orderedConnectionIds) { + final ListingRequestResult listingRequest = createFlowFileListingRequest(user, originalUri, connectionId, expectedNodes); + final Integer queuedFlowFiles = getQueuedFlowFiles(listingRequest.entity()); + queuePollResults.put(connectionId, new QueuePollResult(queuedFlowFiles, listingRequest.hasExpectedCoverage(), + listingRequest.involvedNodeIds(), listingRequest.completedNodeIds(), listingRequest.successfulNodeIds())); + logger.debug("Removed connection drain cluster queue poll [connectionId={}, queuedFlowFiles={}, expectedNodeIds={}, involvedNodeIds={}, " + + "completedNodeIds={}, successfulNodeIds={}, expectedCoverage={}]", connectionId, queuedFlowFiles, + getExpectedNodeIds(expectedNodes), listingRequest.involvedNodeIds(), listingRequest.completedNodeIds(), + listingRequest.successfulNodeIds(), listingRequest.hasExpectedCoverage()); + if (!listingRequest.hasExpectedCoverage() || queuedFlowFiles == null || queuedFlowFiles != 0) { + allQueuesEmpty = false; + } + + final String listingRequestId = getListingRequestId(listingRequest.entity()); + final boolean listingDeleted = deleteFlowFileListingRequest(user, originalUri, connectionId, listingRequestId, expectedNodes); + if (!listingDeleted) { + logger.debug("Failed to delete replicated flow file listing request {} for connection {}", listingRequestId, connectionId); + } + + if (!allQueuesEmpty) { + break; + } + } + + if (allQueuesEmpty) { + return true; + } + + continuePolling = pause.pause(); + } + + logger.warn("Removed connection drain cluster queue wait ended with remaining queues {}", queuePollResults); + return false; + } + private boolean waitForControllerServiceValidation(final NiFiUser user, final URI originalUri, final String groupId, final Set serviceIds, final Pause pause) throws InterruptedException { @@ -680,6 +739,183 @@ private boolean waitForControllerServiceValidation(final NiFiUser user, final UR return false; } + private ListingRequestResult createFlowFileListingRequest(final NiFiUser user, final URI originalUri, final String connectionId, + final Set expectedNodes) throws LifecycleManagementException { + final URI createListingRequestUri; + try { + createListingRequestUri = new URI(originalUri.getScheme(), originalUri.getUserInfo(), originalUri.getHost(), originalUri.getPort(), + "/nifi-api/flowfile-queues/" + connectionId + "/listing-requests", null, originalUri.getFragment()); + } catch (final URISyntaxException e) { + throw new RuntimeException(e); + } + + try { + final AsyncClusterResponse clusterResponse = replicateFlowFileListingRequest(expectedNodes, user, HttpMethod.POST, createListingRequestUri); + + final NodeResponse mergedResponse = clusterResponse.awaitMergedResponse(); + final ListingRequestEntity listingRequestEntity = mergedResponse != null && mergedResponse.is2xx() + ? getResponseEntity(mergedResponse, ListingRequestEntity.class) + : null; + final Set involvedNodeIds = getNodeIds(clusterResponse.getNodesInvolved()); + final Set completedNodeIds = getNodeIds(clusterResponse.getCompletedNodeIdentifiers()); + final Set successfulNodeIds = getSuccessfulNodeIds(clusterResponse.getCompletedNodeResponses()); + + if (!hasExpectedSuccessfulNodeCoverage(clusterResponse, expectedNodes)) { + return new ListingRequestResult(listingRequestEntity, false, involvedNodeIds, completedNodeIds, successfulNodeIds); + } + + if (mergedResponse == null || !mergedResponse.is2xx()) { + return new ListingRequestResult(listingRequestEntity, false, involvedNodeIds, completedNodeIds, successfulNodeIds); + } + + return new ListingRequestResult(listingRequestEntity, true, involvedNodeIds, completedNodeIds, successfulNodeIds); + } catch (final InterruptedException e) { + Thread.currentThread().interrupt(); + throw new LifecycleManagementException("Interrupted while waiting for connection queues to empty"); + } + } + + private boolean deleteFlowFileListingRequest(final NiFiUser user, final URI originalUri, final String connectionId, final String requestId, + final Set expectedNodes) { + if (requestId == null) { + return false; + } + + final URI listingRequestUri; + try { + listingRequestUri = new URI(originalUri.getScheme(), originalUri.getUserInfo(), originalUri.getHost(), originalUri.getPort(), + "/nifi-api/flowfile-queues/" + connectionId + "/listing-requests/" + requestId, null, originalUri.getFragment()); + } catch (final URISyntaxException e) { + throw new RuntimeException(e); + } + + try { + final AsyncClusterResponse clusterResponse = replicateFlowFileListingRequest(expectedNodes, user, HttpMethod.DELETE, listingRequestUri); + + final NodeResponse mergedResponse = clusterResponse.awaitMergedResponse(); + return mergedResponse != null && mergedResponse.is2xx() && hasExpectedSuccessfulNodeCoverage(clusterResponse, expectedNodes); + } catch (final Exception e) { + logger.debug("Failed to delete replicated flow file listing request {} for connection {}", requestId, connectionId, e); + return false; + } + } + + private AsyncClusterResponse replicateFlowFileListingRequest(final Set expectedNodes, final NiFiUser user, final String method, final URI requestUri) { + return getRequestReplicator().replicate(expectedNodes, user, method, requestUri, Collections.emptyMap(), Collections.emptyMap(), true, true); + } + + private String getListingRequestId(final ListingRequestEntity entity) { + if (entity == null || entity.getListingRequest() == null) { + return null; + } + + return entity.getListingRequest().getId(); + } + + private record ListingRequestResult(ListingRequestEntity entity, boolean hasExpectedCoverage, Set involvedNodeIds, + Set completedNodeIds, Set successfulNodeIds) { + } + + private record QueuePollResult(Integer queuedFlowFiles, boolean expectedCoverage, Set involvedNodeIds, + Set completedNodeIds, Set successfulNodeIds) { + } + + private boolean hasExpectedSuccessfulNodeCoverage(final AsyncClusterResponse clusterResponse, final Set expectedNodes) { + if (clusterResponse == null || !clusterResponse.isComplete()) { + return false; + } + + final Set expectedNodeIds = getExpectedNodeIds(expectedNodes); + if (!hasExpectedNodeIds(clusterResponse.getNodesInvolved(), expectedNodeIds)) { + return false; + } + + if (!hasExpectedNodeIds(clusterResponse.getCompletedNodeIdentifiers(), expectedNodeIds)) { + return false; + } + + final Set completedNodeResponses = clusterResponse.getCompletedNodeResponses(); + if (completedNodeResponses == null || completedNodeResponses.size() != expectedNodes.size()) { + return false; + } + + final Set completedNodeIds = new HashSet<>(); + for (final NodeResponse nodeResponse : completedNodeResponses) { + if (nodeResponse == null || nodeResponse.getNodeId() == null || !nodeResponse.is2xx()) { + return false; + } + + final String nodeId = nodeResponse.getNodeId().getId(); + if (nodeId == null || !expectedNodeIds.contains(nodeId) || !completedNodeIds.add(nodeId)) { + return false; + } + } + + return completedNodeIds.equals(expectedNodeIds); + } + + private Set getExpectedNodeIds(final Set expectedNodes) { + final Set expectedNodeIds = new HashSet<>(); + for (final NodeIdentifier nodeIdentifier : expectedNodes) { + if (nodeIdentifier == null || nodeIdentifier.getId() == null || !expectedNodeIds.add(nodeIdentifier.getId())) { + throw new IllegalArgumentException("Expected connected nodes must contain unique node identifiers"); + } + } + + return expectedNodeIds; + } + + private boolean hasExpectedNodeIds(final Set actualNodes, final Set expectedNodeIds) { + if (actualNodes == null || actualNodes.size() != expectedNodeIds.size()) { + return false; + } + + final Set actualNodeIds = new HashSet<>(); + for (final NodeIdentifier actualNode : actualNodes) { + if (actualNode == null || actualNode.getId() == null || !actualNodeIds.add(actualNode.getId())) { + return false; + } + } + + return actualNodeIds.equals(expectedNodeIds); + } + + private Integer getQueuedFlowFiles(final ListingRequestEntity entity) { + if (entity == null || entity.getListingRequest() == null) { + return null; + } + + final ListingRequestDTO listingRequest = entity.getListingRequest(); + return listingRequest.getQueueSize() == null ? null : listingRequest.getQueueSize().getObjectCount(); + } + + private Set getNodeIds(final Set nodeIdentifiers) { + if (nodeIdentifiers == null) { + return Collections.emptySet(); + } + + return nodeIdentifiers.stream() + .filter(Objects::nonNull) + .map(NodeIdentifier::getId) + .filter(Objects::nonNull) + .collect(Collectors.toSet()); + } + + private Set getSuccessfulNodeIds(final Set nodeResponses) { + if (nodeResponses == null) { + return Collections.emptySet(); + } + + return nodeResponses.stream() + .filter(Objects::nonNull) + .filter(NodeResponse::is2xx) + .map(NodeResponse::getNodeId) + .filter(Objects::nonNull) + .map(NodeIdentifier::getId) + .filter(Objects::nonNull) + .collect(Collectors.toSet()); + } + private boolean isControllerServiceValidationComplete(final Set controllerServiceEntities, final Map affectedComponents) { updateAffectedControllerServices(controllerServiceEntities, affectedComponents); for (final ControllerServiceEntity entity : controllerServiceEntities) { diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ComponentLifecycle.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ComponentLifecycle.java index 3d299edd9b56..0c4c89f2833b 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ComponentLifecycle.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/ComponentLifecycle.java @@ -59,4 +59,6 @@ Set scheduleComponents(URI exampleUri, String groupId, */ Set activateControllerServices(URI exampleUri, String groupId, Set servicesToUpdate, Set servicesRequiringDesiredState, ControllerServiceState desiredState, Pause pause, InvalidComponentAction invalidComponentAction) throws LifecycleManagementException; + + boolean waitForConnectionQueuesEmpty(URI exampleUri, Set connectionIds, Pause pause) throws LifecycleManagementException; } diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/LocalComponentLifecycle.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/LocalComponentLifecycle.java index efe346e3e9b4..f116c9b108f0 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/LocalComponentLifecycle.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/main/java/org/apache/nifi/web/util/LocalComponentLifecycle.java @@ -27,6 +27,7 @@ import org.apache.nifi.web.api.dto.AffectedComponentDTO; import org.apache.nifi.web.api.dto.ControllerServiceDTO; import org.apache.nifi.web.api.dto.DtoFactory; +import org.apache.nifi.web.api.dto.ListingRequestDTO; import org.apache.nifi.web.api.dto.ProcessorDTO; import org.apache.nifi.web.api.dto.ProcessorRunStatusDetailsDTO; import org.apache.nifi.web.api.entity.AffectedComponentEntity; @@ -40,10 +41,13 @@ import java.net.URI; import java.util.ArrayList; +import java.util.Comparator; import java.util.HashSet; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Set; +import java.util.UUID; import java.util.function.Function; import java.util.stream.Collectors; @@ -102,6 +106,56 @@ public Set activateControllerServices(final URI example .collect(Collectors.toSet()); } + @Override + public boolean waitForConnectionQueuesEmpty(final URI exampleUri, final Set connectionIds, final Pause pause) throws LifecycleManagementException { + if (connectionIds.isEmpty()) { + return true; + } + + final List orderedConnectionIds = connectionIds.stream() + .sorted(Comparator.naturalOrder()) + .toList(); + final Map queuedFlowFilesByConnection = new LinkedHashMap<>(); + + boolean continuePolling = true; + while (continuePolling) { + boolean allQueuesEmpty = true; + for (final String connectionId : orderedConnectionIds) { + final Integer queuedFlowFiles = getQueuedFlowFiles(connectionId); + queuedFlowFilesByConnection.put(connectionId, queuedFlowFiles); + if (queuedFlowFiles == null || queuedFlowFiles != 0) { + allQueuesEmpty = false; + break; + } + } + + if (allQueuesEmpty) { + return true; + } + + continuePolling = pause.pause(); + } + + logger.warn("Removed connection drain queue wait ended with remaining queues {}", queuedFlowFilesByConnection); + return false; + } + + private Integer getQueuedFlowFiles(final String connectionId) { + final String requestId = UUID.randomUUID().toString(); + try { + final ListingRequestDTO listingRequest = serviceFacade.createFlowFileListingRequest(connectionId, requestId); + final Integer queuedFlowFiles = listingRequest == null || listingRequest.getQueueSize() == null + ? null : listingRequest.getQueueSize().getObjectCount(); + logger.debug("Removed connection drain queue poll [connectionId={}, queuedFlowFiles={}]", connectionId, queuedFlowFiles); + return queuedFlowFiles; + } finally { + try { + serviceFacade.deleteFlowFileListingRequest(connectionId, requestId); + } catch (final Exception ignored) { + } + } + } + private void startComponents(final String processGroupId, final Map componentRevisions, final Map affectedComponents, final Pause pause, final InvalidComponentAction invalidComponentAction) throws LifecycleManagementException { diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/RemovedConnectionDrainClassifierTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/RemovedConnectionDrainClassifierTest.java new file mode 100644 index 000000000000..238224476d55 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/RemovedConnectionDrainClassifierTest.java @@ -0,0 +1,531 @@ +/* + * 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.web; + +import org.apache.nifi.components.validation.ValidationStatus; +import org.apache.nifi.connectable.ConnectableType; +import org.apache.nifi.controller.ScheduledState; +import org.junit.jupiter.api.Test; + +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class RemovedConnectionDrainClassifierTest { + private static final String ROOT_GROUP_ID = "root-group"; + + private final RemovedConnectionDrainClassifier classifier = new RemovedConnectionDrainClassifier(); + + @Test + public void testEmptyRemovedConnectionIsNoDrainRegardlessOfUnsupportedTopology() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-empty", "c-empty-v", ROOT_GROUP_ID, + "remote-source", "remote-source-v", ROOT_GROUP_ID, ConnectableType.REMOTE_OUTPUT_PORT, + "funnel-destination", "funnel-destination-v", ROOT_GROUP_ID, ConnectableType.FUNNEL, + RemovalReason.SOURCE_CHANGED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-empty", "remote-source", "funnel-destination", true) + .addConnectable("remote-source", ConnectableType.REMOTE_OUTPUT_PORT, ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID, true) + .addConnectable("funnel-destination", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify(createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertTrue(result.isSupported()); + assertEquals(RemovedConnectionDrainClassifier.Classification.NO_DRAIN, result.connectionResults().getFirst().classification()); + } + + @Test + public void testUnknownQueueStateDoesNotSkipDrain() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-unknown-queue", "c-unknown-queue-v", ROOT_GROUP_ID, + "source", "source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination", "destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-unknown-queue", "source", "destination", false) + .addProcessor("source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify(createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertTrue(result.isSupported()); + assertEquals(RemovedConnectionDrainClassifier.Classification.CANDIDATE, result.connectionResults().getFirst().classification()); + } + + @Test + public void testNonEmptySourceChangedRemovalIsUnsupported() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-source-changed", "c-source-changed-v", ROOT_GROUP_ID, + "source", "source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination", "destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.SOURCE_CHANGED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-source-changed", "source", "destination", false) + .addProcessor("source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify(createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertFalse(result.isSupported()); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.SOURCE_CHANGED_REMOVAL, + result.getFirstUnsupportedConnectionResult().orElseThrow().unsupportedReason()); + } + + @Test + public void testNullContainingGroupIsUnsupported() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-null-group", "c-null-group-v", null, + "source", "source-v", null, ConnectableType.PROCESSOR, + "destination", "destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-null-group", "source", "destination", false) + .addProcessor("source", null, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertFalse(result.isSupported()); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.CONNECTION_IN_REMOVED_GROUP, + result.getFirstUnsupportedConnectionResult().orElseThrow().unsupportedReason()); + } + + @Test + public void testMissingProducerUsesSourceAgnosticReason() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-missing-producer", "c-missing-producer-v", ROOT_GROUP_ID, + "missing-source", "missing-source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination", "destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-missing-producer", "missing-source", "destination", false) + .addProcessor("destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertFalse(result.isSupported()); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.NO_SUPPORTED_PRODUCER_FOUND, + result.getFirstUnsupportedConnectionResult().orElseThrow().unsupportedReason()); + } + + @Test + public void testRemovedGroupAndRemovedEndpointAreUnsupportedWhenQueueNotEmpty() { + final RemovedConnectionDescriptor removedInGroup = createDescriptor("c-removed-group", "c-removed-group-v", "child-group", + "source-a", "source-a-v", "child-group", ConnectableType.PROCESSOR, + "destination-a", "destination-a-v", "child-group", ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor removedEndpoint = createDescriptor("c-removed-endpoint", "c-removed-endpoint-v", ROOT_GROUP_ID, + "source-b", "source-b-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-b", "destination-b-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addGroup("child-group", ROOT_GROUP_ID) + .addConnection("c-removed-group", "source-a", "destination-a", false) + .addConnection("c-removed-endpoint", "source-b", "destination-b", false) + .addProcessor("source-a", "child-group", ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", "child-group", ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("source-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(List.of(removedInGroup, removedEndpoint), Set.of("child-group"), Set.of("destination-b")), context); + + final Map reasonsByConnection = result.connectionResults().stream() + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::unsupportedReason)); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.CONNECTION_IN_REMOVED_GROUP, reasonsByConnection.get("c-removed-group")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.DESTINATION_COMPONENT_REMOVED, reasonsByConnection.get("c-removed-endpoint")); + } + + @Test + public void testDirectSelfLoopRemovedSourceRemoteSourceAndFunnelDestinationAreUnsupported() { + final RemovedConnectionDescriptor selfLoop = createDescriptor("c-self-loop", "c-self-loop-v", ROOT_GROUP_ID, + "self-loop", "self-loop-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "self-loop", "self-loop-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor removedSource = createDescriptor("c-removed-source", "c-removed-source-v", ROOT_GROUP_ID, + "removed-source", "removed-source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-removed-source", "destination-removed-source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor remoteSource = createDescriptor("c-remote-source", "c-remote-source-v", ROOT_GROUP_ID, + "remote-source", "remote-source-v", ROOT_GROUP_ID, ConnectableType.REMOTE_OUTPUT_PORT, + "destination-remote-source", "destination-remote-source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor funnelDestination = createDescriptor("c-funnel-destination", "c-funnel-destination-v", ROOT_GROUP_ID, + "source-funnel-destination", "source-funnel-destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "funnel-destination", "funnel-destination-v", ROOT_GROUP_ID, ConnectableType.FUNNEL, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-self-loop", "self-loop", "self-loop", false) + .addConnection("c-removed-source", "removed-source", "destination-removed-source", false) + .addConnection("c-remote-source", "remote-source", "destination-remote-source", false) + .addConnection("c-funnel-destination", "source-funnel-destination", "funnel-destination", false) + .addProcessor("self-loop", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("removed-source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-removed-source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addConnectable("remote-source", ConnectableType.REMOTE_OUTPUT_PORT, ROOT_GROUP_ID, null, null, true) + .addProcessor("destination-remote-source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("source-funnel-destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addConnectable("funnel-destination", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(List.of(selfLoop, removedSource, remoteSource, funnelDestination), Set.of(), Set.of("removed-source")), context); + + final Map reasonsByConnection = result.connectionResults().stream() + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::unsupportedReason)); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.SELF_LOOP, reasonsByConnection.get("c-self-loop")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.SOURCE_COMPONENT_REMOVED, reasonsByConnection.get("c-removed-source")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.UNSUPPORTED_SOURCE_TYPE, reasonsByConnection.get("c-remote-source")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.UNSUPPORTED_DESTINATION_TYPE, reasonsByConnection.get("c-funnel-destination")); + } + + @Test + public void testProcessorDestinationMustBePhysicallyRunningAndValid() { + final RemovedConnectionDescriptor invalidDestination = createDescriptor("c-invalid-destination", "c-invalid-destination-v", ROOT_GROUP_ID, + "source-invalid", "source-invalid-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-invalid", "destination-invalid-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor startingDestination = createDescriptor("c-starting-destination", "c-starting-destination-v", ROOT_GROUP_ID, + "source-starting", "source-starting-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-starting", "destination-starting-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-invalid-destination", "source-invalid", "destination-invalid", false) + .addConnection("c-starting-destination", "source-starting", "destination-starting", false) + .addProcessor("source-invalid", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-invalid", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.INVALID) + .addProcessor("source-starting", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-starting", ROOT_GROUP_ID, ScheduledState.STARTING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify(createImpact(List.of(invalidDestination, startingDestination), Set.of(), Set.of()), context); + + final Map reasonsByConnection = result.connectionResults().stream() + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::unsupportedReason)); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.DESTINATION_NOT_VALID, reasonsByConnection.get("c-invalid-destination")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.DESTINATION_NOT_RUNNING, reasonsByConnection.get("c-starting-destination")); + } + + @Test + public void testPortDestinationUsesRunningState() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-port", "c-port-v", ROOT_GROUP_ID, + "source-port", "source-port-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-port", "destination-port-v", ROOT_GROUP_ID, ConnectableType.OUTPUT_PORT, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-port", "source-port", "destination-port", false) + .addProcessor("source-port", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addPort("destination-port", ConnectableType.OUTPUT_PORT, ROOT_GROUP_ID, false); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify(createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.DESTINATION_NOT_RUNNING, + result.getFirstUnsupportedConnectionResult().orElseThrow().unsupportedReason()); + } + + @Test + public void testSharedProducerIsDeduplicatedAndSharedDestinationSupported() { + final RemovedConnectionDescriptor sharedProducerA = createDescriptor("c-shared-producer-a", "c-shared-producer-a-v", ROOT_GROUP_ID, + "shared-source", "shared-source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-a", "destination-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor sharedProducerB = createDescriptor("c-shared-producer-b", "c-shared-producer-b-v", ROOT_GROUP_ID, + "shared-source", "shared-source-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-b", "destination-b-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor sharedDestination = createDescriptor("c-shared-destination", "c-shared-destination-v", ROOT_GROUP_ID, + "source-c", "source-c-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination-a", "destination-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-shared-producer-a", "shared-source", "destination-a", false) + .addConnection("c-shared-producer-b", "shared-source", "destination-b", false) + .addConnection("c-shared-destination", "source-c", "destination-a", false) + .addProcessor("shared-source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("source-c", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(List.of(sharedProducerA, sharedProducerB, sharedDestination), Set.of(), Set.of()), context); + + assertTrue(result.isSupported()); + assertEquals(Set.of("shared-source", "source-c"), result.producerBarrierComponentIds()); + assertTrue(result.connectionResults().stream().allMatch(r -> r.classification() == RemovedConnectionDrainClassifier.Classification.CANDIDATE)); + } + + @Test + public void testChainAndCycleTopologiesAreUnsupported() { + final RemovedConnectionDescriptor chainFirst = createDescriptor("c-chain-a", "c-chain-a-v", ROOT_GROUP_ID, + "producer-a", "producer-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "middle", "middle-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor chainSecond = createDescriptor("c-chain-b", "c-chain-b-v", ROOT_GROUP_ID, + "middle", "middle-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "consumer", "consumer-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor cycleFirst = createDescriptor("c-cycle-a", "c-cycle-a-v", ROOT_GROUP_ID, + "cycle-a", "cycle-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "cycle-b", "cycle-b-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor cycleSecond = createDescriptor("c-cycle-b", "c-cycle-b-v", ROOT_GROUP_ID, + "cycle-b", "cycle-b-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "cycle-a", "cycle-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-chain-a", "producer-a", "middle", false) + .addConnection("c-chain-b", "middle", "consumer", false) + .addConnection("c-cycle-a", "cycle-a", "cycle-b", false) + .addConnection("c-cycle-b", "cycle-b", "cycle-a", false) + .addProcessor("producer-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("middle", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("consumer", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("cycle-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("cycle-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(List.of(chainFirst, chainSecond, cycleFirst, cycleSecond), Set.of(), Set.of()), context); + + final Map classificationByConnection = result.connectionResults().stream() + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::classification)); + final Map reasonsByConnection = result.connectionResults().stream() + .filter(r -> r.classification() == RemovedConnectionDrainClassifier.Classification.UNSUPPORTED) + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::unsupportedReason)); + assertEquals(RemovedConnectionDrainClassifier.Classification.CANDIDATE, classificationByConnection.get("c-chain-a")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.PRODUCER_BARRIER_IS_REMOVED_DESTINATION, reasonsByConnection.get("c-chain-b")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.PRODUCER_BARRIER_IS_REMOVED_DESTINATION, reasonsByConnection.get("c-cycle-a")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.PRODUCER_BARRIER_IS_REMOVED_DESTINATION, reasonsByConnection.get("c-cycle-b")); + } + + @Test + public void testFunnelTraversalIsCycleSafeAndRejectsFunnelSelfLoop() { + final RemovedConnectionDescriptor supportedFunnel = createDescriptor("c-funnel", "c-funnel-v", ROOT_GROUP_ID, + "funnel", "funnel-v", ROOT_GROUP_ID, ConnectableType.FUNNEL, + "destination", "destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor selfLoopFunnel = createDescriptor("c-funnel-self-loop", "c-funnel-self-loop-v", ROOT_GROUP_ID, + "funnel-self-loop", "funnel-self-loop-v", ROOT_GROUP_ID, ConnectableType.FUNNEL, + "loop-destination", "loop-destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-funnel", "funnel", "destination", false) + .addConnection("c-funnel-self-loop", "funnel-self-loop", "loop-destination", false) + .addConnection("incoming-a", "upstream-a", "funnel", false) + .addConnection("incoming-cycle-1", "funnel-cycle", "funnel", false) + .addConnection("incoming-cycle-2", "funnel", "funnel-cycle", false) + .addConnection("loop-incoming", "loop-destination", "funnel-self-loop", false) + .addProcessor("upstream-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("loop-destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addConnectable("funnel", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false) + .addConnectable("funnel-cycle", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false) + .addConnectable("funnel-self-loop", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false); + + context.linkIncoming("funnel", Set.of("incoming-a", "incoming-cycle-1")); + context.linkOutgoing("funnel", Set.of("c-funnel", "incoming-cycle-2")); + context.linkIncoming("funnel-cycle", Set.of("incoming-cycle-2")); + context.linkOutgoing("funnel-cycle", Set.of("incoming-cycle-1")); + context.linkIncoming("funnel-self-loop", Set.of("loop-incoming")); + context.linkOutgoing("funnel-self-loop", Set.of("c-funnel-self-loop")); + context.linkOutgoing("upstream-a", Set.of("incoming-a")); + context.linkOutgoing("loop-destination", Set.of("loop-incoming")); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(List.of(supportedFunnel, selfLoopFunnel), Set.of(), Set.of()), context); + + final Map classificationByConnection = result.connectionResults().stream() + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::classification)); + final Map reasonsByConnection = result.connectionResults().stream() + .filter(r -> r.classification() == RemovedConnectionDrainClassifier.Classification.UNSUPPORTED) + .collect(Collectors.toMap(r -> r.connection().getConnectionInstanceId(), RemovedConnectionDrainClassifier.ConnectionResult::unsupportedReason)); + assertEquals(RemovedConnectionDrainClassifier.Classification.CANDIDATE, classificationByConnection.get("c-funnel")); + assertEquals(RemovedConnectionDrainClassifier.Classification.UNSUPPORTED, classificationByConnection.get("c-funnel-self-loop")); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.SELF_LOOP, reasonsByConnection.get("c-funnel-self-loop")); + assertEquals(Set.of("upstream-a"), result.connectionResults().stream() + .filter(r -> "c-funnel".equals(r.connection().getConnectionInstanceId())) + .findFirst().orElseThrow().producerBarrierComponentIds()); + } + + @Test + public void testRetainedFeedbackPathIsUnsupported() { + final RemovedConnectionDescriptor removedConnection = createDescriptor("c-feedback", "c-feedback-v", ROOT_GROUP_ID, + "producer", "producer-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "destination", "destination-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("c-feedback", "producer", "destination", false) + .addConnection("retained-1", "destination", "retained-mid", false) + .addConnection("retained-2", "retained-mid", "producer", false) + .addProcessor("producer", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("retained-mid", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + context.linkOutgoing("destination", Set.of("retained-1")); + context.linkOutgoing("retained-mid", Set.of("retained-2")); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify(createImpact(Set.of(removedConnection), Set.of(), Set.of()), context); + + assertFalse(result.isSupported()); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.RETAINED_FEEDBACK_PATH, + result.getFirstUnsupportedConnectionResult().orElseThrow().unsupportedReason()); + } + + @Test + public void testBatchUnsupportedResultIsDeterministic() { + final RemovedConnectionDescriptor laterUnsupported = createDescriptor("z-unsupported", "z-unsupported-v", ROOT_GROUP_ID, + "source-z", "source-z-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "dest-z", "dest-z-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor earlierUnsupported = createDescriptor("a-unsupported", "a-unsupported-v", ROOT_GROUP_ID, + "source-a", "source-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + "dest-a", "dest-a-v", ROOT_GROUP_ID, ConnectableType.PROCESSOR, + RemovalReason.SOURCE_CHANGED); + + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("z-unsupported", "source-z", "dest-z", false) + .addConnection("a-unsupported", "source-a", "dest-a", false) + .addProcessor("source-z", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("dest-z", ROOT_GROUP_ID, ScheduledState.STOPPED, ValidationStatus.VALID) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("dest-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final RemovedConnectionDrainClassifier.BatchResult result = classifier.classify( + createImpact(List.of(laterUnsupported, earlierUnsupported), Set.of(), Set.of()), context); + + assertEquals("a-unsupported", result.getFirstUnsupportedConnectionResult().orElseThrow().connection().getConnectionInstanceId()); + assertEquals(RemovedConnectionDrainClassifier.UnsupportedReason.SOURCE_CHANGED_REMOVAL, + result.getFirstUnsupportedConnectionResult().orElseThrow().unsupportedReason()); + } + + private FlowUpdateImpact createImpact(final Iterable removedConnections, + final Set removedProcessGroupIds, + final Set removedEndpointIds) { + final Set descriptors = new LinkedHashSet<>(); + removedConnections.forEach(descriptors::add); + return new FlowUpdateImpact(Set.of(), descriptors, removedProcessGroupIds, removedEndpointIds); + } + + private RemovedConnectionDescriptor createDescriptor(final String connectionInstanceId, final String connectionVersionedId, + final String containingProcessGroupId, + final String sourceInstanceId, final String sourceVersionedId, + final String sourceProcessGroupId, final ConnectableType sourceType, + final String destinationInstanceId, final String destinationVersionedId, + final String destinationProcessGroupId, final ConnectableType destinationType, + final RemovalReason removalReason) { + return new RemovedConnectionDescriptor(connectionInstanceId, connectionVersionedId, containingProcessGroupId, + sourceInstanceId, sourceVersionedId, sourceProcessGroupId, sourceType, + destinationInstanceId, destinationVersionedId, destinationProcessGroupId, destinationType, + removalReason); + } + + private static final class TestContext implements RemovedConnectionDrainClassifier.Context { + private final Map connections = new LinkedHashMap<>(); + private final Map connectables = new LinkedHashMap<>(); + private final Map groups = new LinkedHashMap<>(); + + TestContext addGroup(final String groupId, final String parentGroupId) { + groups.put(groupId, new RemovedConnectionDrainClassifier.LiveProcessGroup(groupId, parentGroupId)); + return this; + } + + TestContext addConnection(final String connectionId, final String sourceId, final String destinationId, final boolean queueEmpty) { + connections.put(connectionId, new RemovedConnectionDrainClassifier.LiveConnection(connectionId, sourceId, destinationId, queueEmpty)); + return this; + } + + TestContext addProcessor(final String connectableId, final String groupId, + final ScheduledState physicalScheduledState, final ValidationStatus validationStatus) { + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectableId, ConnectableType.PROCESSOR, + groupId, physicalScheduledState, validationStatus, false, Set.of(), Set.of())); + return this; + } + + TestContext addPort(final String connectableId, final ConnectableType connectableType, final String groupId, final boolean running) { + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectableId, connectableType, + groupId, null, null, running, Set.of(), Set.of())); + return this; + } + + TestContext addConnectable(final String connectableId, final ConnectableType connectableType, final String groupId, + final ScheduledState physicalScheduledState, final ValidationStatus validationStatus, + final boolean running) { + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectableId, connectableType, + groupId, physicalScheduledState, validationStatus, running, Set.of(), Set.of())); + return this; + } + + TestContext linkIncoming(final String connectableId, final Set incomingConnectionIds) { + final RemovedConnectionDrainClassifier.LiveConnectable connectable = connectables.get(connectableId); + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectable.id(), connectable.type(), connectable.processGroupId(), + connectable.physicalScheduledState(), connectable.validationStatus(), connectable.running(), incomingConnectionIds, + connectable.outgoingConnectionIds())); + return this; + } + + TestContext linkOutgoing(final String connectableId, final Set outgoingConnectionIds) { + final RemovedConnectionDrainClassifier.LiveConnectable connectable = connectables.get(connectableId); + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectable.id(), connectable.type(), connectable.processGroupId(), + connectable.physicalScheduledState(), connectable.validationStatus(), connectable.running(), connectable.incomingConnectionIds(), + outgoingConnectionIds)); + return this; + } + + @Override + public RemovedConnectionDrainClassifier.LiveConnection getConnection(final String connectionId) { + return connections.get(connectionId); + } + + @Override + public RemovedConnectionDrainClassifier.LiveConnectable getConnectable(final String connectableId) { + return connectables.get(connectableId); + } + + @Override + public RemovedConnectionDrainClassifier.LiveProcessGroup getProcessGroup(final String processGroupId) { + return groups.get(processGroupId); + } + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/RemovedConnectionDrainCoordinatorTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/RemovedConnectionDrainCoordinatorTest.java new file mode 100644 index 000000000000..61a327d95790 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/RemovedConnectionDrainCoordinatorTest.java @@ -0,0 +1,899 @@ +/* + * 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.web; + +import org.apache.nifi.components.validation.ValidationStatus; +import org.apache.nifi.connectable.ConnectableType; +import org.apache.nifi.controller.ScheduledState; +import org.apache.nifi.web.api.dto.AffectedComponentDTO; +import org.apache.nifi.web.api.entity.AffectedComponentEntity; +import org.apache.nifi.web.util.ComponentLifecycle; +import org.apache.nifi.web.util.InvalidComponentAction; +import org.apache.nifi.web.util.LifecycleManagementException; +import org.apache.nifi.web.util.Pause; +import org.junit.jupiter.api.Test; + +import java.net.URI; +import java.time.Duration; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +class RemovedConnectionDrainCoordinatorTest { + private static final String ROOT_GROUP_ID = "root-group"; + private static final URI REQUEST_URI = URI.create("http://localhost:8080/nifi-api"); + + @Test + void testCoordinateDrainReturnsImmediatelyWhenNoCandidateConnections() throws Exception { + final FlowUpdateImpact impact = createImpact(Set.of(), Set.of(affectedProcessor("source", "Running", 1))); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + final TestCancellationHandle cancellationHandle = new TestCancellationHandle(); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, new TestContext().addGroup(ROOT_GROUP_ID, null), lifecycle, REQUEST_URI, ROOT_GROUP_ID, cancellationHandle); + + assertTrue(result.candidateConnectionIds().isEmpty()); + assertTrue(result.drainStoppedComponents().isEmpty()); + assertFalse(result.cancelled()); + assertTrue(lifecycle.scheduleCalls.isEmpty()); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + @Test + void testCoordinateDrainUsesSingleSharedProducerBarrierAndLeavesDestinationRunning() throws Exception { + final RemovedConnectionDescriptor first = descriptor("connection-a", "shared-source", ConnectableType.PROCESSOR, "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor second = descriptor("connection-b", "shared-source", ConnectableType.PROCESSOR, "destination-b", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(first, second), Set.of( + affectedProcessor("shared-source", "Running", 1), + affectedProcessor("destination-a", "Running", 2), + affectedProcessor("destination-b", "Running", 3))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "shared-source", "destination-a", false) + .addConnection("connection-b", "shared-source", "destination-b", false) + .addProcessor("shared-source", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("shared-source", affectedProcessor("shared-source", "Stopped", 0)); + lifecycle.queueWaitResult = true; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle()); + + assertFalse(result.cancelled()); + assertEquals(Set.of("connection-a", "connection-b"), result.candidateConnectionIds()); + assertEquals(List.of(new ScheduleCall(ScheduledState.STOPPED, Set.of("shared-source"))), lifecycle.scheduleCalls); + assertEquals(List.of(Set.of("connection-a", "connection-b")), lifecycle.queueWaitCalls); + assertEquals(Set.of("shared-source"), ids(result.drainStoppedComponents())); + } + + @Test + void testCoordinateDrainCancelsDuringProducerStopAndRestoresOnlyActuallyStoppedComponents() throws Exception { + final RemovedConnectionDescriptor first = descriptor("connection-a", "source-a", ConnectableType.PROCESSOR, "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final RemovedConnectionDescriptor second = descriptor("connection-b", "source-b", ConnectableType.PROCESSOR, "destination-b", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(first, second), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("source-b", "Running", 1), + affectedProcessor("destination-a", "Running", 1), + affectedProcessor("destination-b", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addConnection("connection-b", "source-b", "destination-b", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("source-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.cancelAfterStop = true; + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.stoppedResultById.put("source-b", affectedProcessor("source-b", "Running", 1)); + lifecycle.runningResultById.put("source-a", affectedProcessor("source-a", "Running", 1)); + lifecycle.runningResultById.put("source-b", affectedProcessor("source-b", "Running", 1)); + final TestCancellationHandle cancellationHandle = new TestCancellationHandle(); + lifecycle.cancellationHandle = cancellationHandle; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, cancellationHandle); + + assertTrue(result.cancelled()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a", "source-b")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + @Test + void testCoordinateDrainCancelsDuringQueueWaitAndRestoresStoppedComponents() throws Exception { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.runningResultById.put("source-a", affectedProcessor("source-a", "Running", 1)); + lifecycle.cancelDuringQueueWait = true; + final TestCancellationHandle cancellationHandle = new TestCancellationHandle(); + lifecycle.cancellationHandle = cancellationHandle; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, cancellationHandle); + + assertTrue(result.cancelled()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + assertEquals(List.of(Set.of("connection-a")), lifecycle.queueWaitCalls); + } + + @Test + void testCoordinateDrainCancelsWhenQueueWaitReportsSuccessAndRestoresStoppedComponents() throws Exception { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.runningResultById.put("source-a", affectedProcessor("source-a", "Running", 1)); + lifecycle.cancelDuringQueueWait = true; + lifecycle.queueWaitResult = true; + final TestCancellationHandle cancellationHandle = new TestCancellationHandle(); + lifecycle.cancellationHandle = cancellationHandle; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, cancellationHandle); + + assertTrue(result.cancelled()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + } + + @Test + void testCoordinateDrainCancellationRetainsRestorationFailureWithoutRetry() throws Exception { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.cancelDuringQueueWait = true; + lifecycle.restoreException = new LifecycleManagementException("Failed to restore stopped producers"); + final TestCancellationHandle cancellationHandle = new TestCancellationHandle(); + lifecycle.cancellationHandle = cancellationHandle; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, cancellationHandle); + + assertTrue(result.cancelled()); + assertEquals("Failed to restore stopped producers", result.restorationFailure().getMessage()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + } + + @Test + void testCoordinateDrainCancelDuringStopThrowReturnsCancellationResultWithRestorationFailure() throws Exception { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "funnel-source", ConnectableType.FUNNEL, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("upstream-a", "Running", 1), + affectedProcessor("upstream-b", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "funnel-source", "destination-a", false) + .addConnection("incoming-a", "upstream-a", "funnel-source", false) + .addConnection("incoming-b", "upstream-b", "funnel-source", false) + .addProcessor("upstream-a", ROOT_GROUP_ID, ScheduledState.STOPPED, ValidationStatus.VALID) + .addProcessor("upstream-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addConnectable("funnel-source", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false); + context.linkIncoming("funnel-source", Set.of("incoming-a", "incoming-b")); + context.linkOutgoing("upstream-a", Set.of("incoming-a")); + context.linkOutgoing("upstream-b", Set.of("incoming-b")); + + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.cancelAfterStop = true; + lifecycle.stopException = new LifecycleManagementException("Failed while waiting for components to transition to state of STOPPED"); + lifecycle.restoreException = new LifecycleManagementException("Failed to restore stopped producers"); + final TestCancellationHandle cancellationHandle = new TestCancellationHandle(); + lifecycle.cancellationHandle = cancellationHandle; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, cancellationHandle); + + assertTrue(result.cancelled()); + assertEquals(Set.of("connection-a"), result.candidateConnectionIds()); + assertEquals("Failed to restore stopped producers", result.restorationFailure().getMessage()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("upstream-a", "upstream-b")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("upstream-a")) + ), lifecycle.scheduleCalls); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + @Test + void testCoordinateDrainTimesOutAcrossStopAndQueueWaitUsingSingleDeadline() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final SequencedPauseFactory pauseFactory = new SequencedPauseFactory(); + final RecordingDeadlinePause deadlinePause = new RecordingDeadlinePause(List.of(true, false)); + pauseFactory.drainPause = deadlinePause; + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.queueWaitUsesPause = true; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), pauseFactory, Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertTrue(exception.getMessage().contains("Removed connection drain timed out")); + assertTrue(exception.getMessage().contains("connectionIds=[connection-a]")); + assertSame(deadlinePause, lifecycle.queueWaitPause); + assertEquals(3, deadlinePause.pauseInvocations); + assertEquals(1, pauseFactory.restorationPauseCreations); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + } + + @Test + void testCoordinateDrainRestoresStoppedComponentsWhenQueueWaitFails() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.runningResultById.put("source-a", affectedProcessor("source-a", "Running", 1)); + lifecycle.queueWaitException = new LifecycleManagementException("Interrupted while waiting for connection queues to empty"); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertTrue(exception.getMessage().contains("Removed connection drain failed while waiting for connections [connection-a]")); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + } + + @Test + void testCoordinateDrainAddsRestorationFailureAsSuppressed() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.queueWaitResult = false; + lifecycle.restoreException = new LifecycleManagementException("Failed to restore stopped producers"); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertEquals(1, exception.getSuppressed().length); + assertEquals("Failed to restore stopped producers", exception.getSuppressed()[0].getMessage()); + assertTrue(exception.getMessage().contains("Removed connection drain timed out")); + } + + @Test + void testCoordinateDrainRestoresStoppedComponentsWhenQueueWaitThrowsRuntimeException() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.runningResultById.put("source-a", affectedProcessor("source-a", "Running", 1)); + lifecycle.queueWaitRuntimeException = new IllegalStateException("Cluster membership changed"); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertSame(lifecycle.queueWaitRuntimeException, exception.getCause()); + assertTrue(exception.getMessage().contains("Removed connection drain failed for connections [connection-a]")); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + } + + @Test + void testCoordinateDrainRetainsRestorationFailureWhenQueueWaitThrowsRuntimeException() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stoppedResultById.put("source-a", affectedProcessor("source-a", "Stopped", 0)); + lifecycle.queueWaitRuntimeException = new IllegalStateException("Cluster membership changed"); + lifecycle.restoreException = new LifecycleManagementException("Failed to restore stopped producers"); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertEquals(1, exception.getSuppressed().length); + assertEquals("Failed to restore stopped producers", exception.getSuppressed()[0].getMessage()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + } + + @Test + void testCoordinateDrainRestoresOnlyLiveStoppedProducerWhenStopThrowsAfterPartialTransition() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "funnel-source", ConnectableType.FUNNEL, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("upstream-a", "Running", 1), + affectedProcessor("upstream-b", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "funnel-source", "destination-a", false) + .addConnection("incoming-a", "upstream-a", "funnel-source", false) + .addConnection("incoming-b", "upstream-b", "funnel-source", false) + .addProcessor("upstream-a", ROOT_GROUP_ID, ScheduledState.STOPPED, ValidationStatus.VALID) + .addProcessor("upstream-b", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addConnectable("funnel-source", ConnectableType.FUNNEL, ROOT_GROUP_ID, null, null, false); + context.linkIncoming("funnel-source", Set.of("incoming-a", "incoming-b")); + context.linkOutgoing("upstream-a", Set.of("incoming-a")); + context.linkOutgoing("upstream-b", Set.of("incoming-b")); + + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stopException = new LifecycleManagementException("Failed while waiting for components to transition to state of STOPPED"); + lifecycle.runningResultById.put("upstream-a", affectedProcessor("upstream-a", "Running", 1)); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertEquals("Removed connection drain failed while stopping producer barriers [upstream-a, upstream-b]: Failed while waiting for components to transition to state of STOPPED", + exception.getMessage()); + assertSame(lifecycle.stopException, exception.getCause()); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("upstream-a", "upstream-b")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("upstream-a")) + ), lifecycle.scheduleCalls); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + @Test + void testCoordinateDrainRestoresStoppingProcessorWhenStopThrowsAfterPartialTransition() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.STOPPING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stopException = new LifecycleManagementException("Failed while waiting for components to transition to state of STOPPED"); + lifecycle.runningResultById.put("source-a", affectedProcessor("source-a", "Running", 1)); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertTrue(exception.getMessage().contains("Removed connection drain failed while stopping producer barriers [source-a]")); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("source-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("source-a")) + ), lifecycle.scheduleCalls); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + @Test + void testCoordinateDrainRestoresStoppedInputPortWhenStopThrowsAfterPartialTransition() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "input-port-a", ConnectableType.INPUT_PORT, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedPort("input-port-a", AffectedComponentDTO.COMPONENT_TYPE_INPUT_PORT, "Running"), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "input-port-a", "destination-a", false) + .addConnectable("input-port-a", ConnectableType.INPUT_PORT, ROOT_GROUP_ID, null, null, false) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.stopException = new LifecycleManagementException("Failed while waiting for components to transition to state of STOPPED"); + lifecycle.runningResultById.put("input-port-a", affectedPort("input-port-a", AffectedComponentDTO.COMPONENT_TYPE_INPUT_PORT, "Running")); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertTrue(exception.getMessage().contains("Removed connection drain failed while stopping producer barriers [input-port-a]")); + assertEquals(List.of( + new ScheduleCall(ScheduledState.STOPPED, Set.of("input-port-a")), + new ScheduleCall(ScheduledState.RUNNING, Set.of("input-port-a")) + ), lifecycle.scheduleCalls); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + @Test + void testAlreadyStoppedProducerIsNotStoppedOrRestarted() throws Exception { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.COMPONENT_REMOVED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Stopped", 0), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.STOPPED, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + lifecycle.queueWaitResult = true; + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final RemovedConnectionDrainCoordinator.DrainResult result = coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle()); + + assertTrue(lifecycle.scheduleCalls.isEmpty()); + assertEquals(Set.of("connection-a"), result.candidateConnectionIds()); + assertTrue(result.drainStoppedComponents().isEmpty()); + } + + @Test + void testUnsupportedClassificationFailsWithoutAnyMutation() { + final RemovedConnectionDescriptor removedConnection = descriptor( + "connection-a", "source-a", ConnectableType.PROCESSOR, + "destination-a", ConnectableType.PROCESSOR, RemovalReason.SOURCE_CHANGED); + final FlowUpdateImpact impact = createImpact(Set.of(removedConnection), Set.of( + affectedProcessor("source-a", "Running", 1), + affectedProcessor("destination-a", "Running", 1))); + final TestContext context = new TestContext() + .addGroup(ROOT_GROUP_ID, null) + .addConnection("connection-a", "source-a", "destination-a", false) + .addProcessor("source-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID) + .addProcessor("destination-a", ROOT_GROUP_ID, ScheduledState.RUNNING, ValidationStatus.VALID); + final TestComponentLifecycle lifecycle = new TestComponentLifecycle(); + + final RemovedConnectionDrainCoordinator coordinator = new RemovedConnectionDrainCoordinator( + new RemovedConnectionDrainClassifier(), new TestPauseFactory(), Duration.ofSeconds(30)); + + final LifecycleManagementException exception = assertThrows(LifecycleManagementException.class, () -> coordinator.coordinateDrain( + impact, context, lifecycle, REQUEST_URI, ROOT_GROUP_ID, new TestCancellationHandle())); + + assertTrue(exception.getMessage().contains("connection-a[reason=SOURCE_CHANGED_REMOVAL]")); + assertTrue(lifecycle.scheduleCalls.isEmpty()); + assertTrue(lifecycle.queueWaitCalls.isEmpty()); + } + + private FlowUpdateImpact createImpact(final Set removedConnections, final Set affectedComponents) { + return new FlowUpdateImpact(affectedComponents, removedConnections, Set.of(), Set.of()); + } + + private RemovedConnectionDescriptor descriptor(final String connectionId, final String sourceId, final ConnectableType sourceType, + final String destinationId, final ConnectableType destinationType, + final RemovalReason removalReason) { + return new RemovedConnectionDescriptor(connectionId, connectionId + "-v", ROOT_GROUP_ID, + sourceId, sourceId + "-v", ROOT_GROUP_ID, sourceType, + destinationId, destinationId + "-v", ROOT_GROUP_ID, destinationType, + removalReason); + } + + private AffectedComponentEntity affectedProcessor(final String id, final String state, final int activeThreadCount) { + final AffectedComponentDTO dto = new AffectedComponentDTO(); + dto.setId(id); + dto.setName(id); + dto.setProcessGroupId(ROOT_GROUP_ID); + dto.setReferenceType(AffectedComponentDTO.COMPONENT_TYPE_PROCESSOR); + dto.setState(state); + dto.setActiveThreadCount(activeThreadCount); + + final AffectedComponentEntity entity = new AffectedComponentEntity(); + entity.setId(id); + entity.setReferenceType(AffectedComponentDTO.COMPONENT_TYPE_PROCESSOR); + entity.setComponent(dto); + return entity; + } + + private AffectedComponentEntity affectedPort(final String id, final String referenceType, final String state) { + final AffectedComponentDTO dto = new AffectedComponentDTO(); + dto.setId(id); + dto.setName(id); + dto.setProcessGroupId(ROOT_GROUP_ID); + dto.setReferenceType(referenceType); + dto.setState(state); + + final AffectedComponentEntity entity = new AffectedComponentEntity(); + entity.setId(id); + entity.setReferenceType(referenceType); + entity.setComponent(dto); + return entity; + } + + private Set ids(final Set components) { + final Set ids = new LinkedHashSet<>(); + for (final AffectedComponentEntity component : components) { + ids.add(component.getId()); + } + return ids; + } + + private record ScheduleCall(ScheduledState state, Set componentIds) { + ScheduleCall { + componentIds = Set.copyOf(componentIds); + } + } + + private static final class TestComponentLifecycle implements ComponentLifecycle { + private final List scheduleCalls = new ArrayList<>(); + private final List> queueWaitCalls = new ArrayList<>(); + private final Map stoppedResultById = new LinkedHashMap<>(); + private final Map runningResultById = new LinkedHashMap<>(); + private final Set initiallyEmptyConnectionIds = new LinkedHashSet<>(); + private boolean queueWaitResult; + private boolean cancelAfterStop; + private boolean cancelDuringQueueWait; + private boolean queueWaitUsesPause; + private LifecycleManagementException stopException; + private LifecycleManagementException queueWaitException; + private RuntimeException queueWaitRuntimeException; + private LifecycleManagementException restoreException; + private TestCancellationHandle cancellationHandle; + private Pause queueWaitPause; + + @Override + public Set scheduleComponents(final URI exampleUri, final String groupId, final Set components, + final ScheduledState desiredState, final Pause pause, + final InvalidComponentAction invalidComponentAction) throws LifecycleManagementException { + final Set componentIds = ids(components); + scheduleCalls.add(new ScheduleCall(desiredState, componentIds)); + + if (desiredState == ScheduledState.RUNNING && restoreException != null) { + throw restoreException; + } + + if (desiredState == ScheduledState.STOPPED && cancelAfterStop && cancellationHandle != null) { + cancellationHandle.cancel(); + } + + final Map resultsById = desiredState == ScheduledState.STOPPED ? stoppedResultById : runningResultById; + final Set results = new LinkedHashSet<>(); + for (final AffectedComponentEntity component : components) { + results.add(resultsById.getOrDefault(component.getId(), component)); + } + + if (desiredState == ScheduledState.STOPPED && stopException != null) { + throw stopException; + } + + return results; + } + + @Override + public Set activateControllerServices(final URI exampleUri, final String groupId, + final Set servicesToUpdate, + final Set servicesRequiringDesiredState, + final org.apache.nifi.controller.service.ControllerServiceState desiredState, + final Pause pause, + final InvalidComponentAction invalidComponentAction) { + throw new UnsupportedOperationException(); + } + + @Override + public boolean waitForConnectionQueuesEmpty(final URI exampleUri, final Set connectionIds, final Pause pause) throws LifecycleManagementException { + if (connectionIds.size() == 1 && !pause.pause()) { + return initiallyEmptyConnectionIds.contains(connectionIds.iterator().next()); + } + + queueWaitCalls.add(Set.copyOf(connectionIds)); + queueWaitPause = pause; + + if (queueWaitException != null) { + throw queueWaitException; + } + + if (queueWaitRuntimeException != null) { + throw queueWaitRuntimeException; + } + + if (cancelDuringQueueWait && cancellationHandle != null) { + cancellationHandle.cancel(); + } + + if (queueWaitUsesPause) { + pause.pause(); + return pause.pause(); + } + + return queueWaitResult; + } + + private Set ids(final Set components) { + final Set ids = new LinkedHashSet<>(); + for (final AffectedComponentEntity component : components) { + ids.add(component.getId()); + } + return ids; + } + } + + private static final class TestCancellationHandle implements RemovedConnectionDrainCoordinator.CancellationHandle { + private boolean cancelled; + private Runnable cancelCallback; + + @Override + public boolean isCancelled() { + return cancelled; + } + + @Override + public void setCancelCallback(final Runnable runnable) { + this.cancelCallback = runnable; + } + + private void cancel() { + cancelled = true; + if (cancelCallback != null) { + cancelCallback.run(); + } + } + } + + private static class TestPauseFactory implements RemovedConnectionDrainCoordinator.PauseFactory { + @Override + public RemovedConnectionDrainCoordinator.DeadlinePause createDrainPause(final Duration timeout) { + return new RecordingDeadlinePause(List.of(true)); + } + + @Override + public Pause createRestorationPause() { + return () -> true; + } + } + + private static final class SequencedPauseFactory extends TestPauseFactory { + private RecordingDeadlinePause drainPause; + private int restorationPauseCreations; + + @Override + public RemovedConnectionDrainCoordinator.DeadlinePause createDrainPause(final Duration timeout) { + return drainPause; + } + + @Override + public Pause createRestorationPause() { + restorationPauseCreations++; + return () -> true; + } + } + + private static final class RecordingDeadlinePause implements RemovedConnectionDrainCoordinator.DeadlinePause { + private final List decisions; + private int index; + private int pauseInvocations; + private boolean cancelled; + + private RecordingDeadlinePause(final List decisions) { + this.decisions = decisions; + } + + @Override + public void cancel() { + cancelled = true; + } + + @Override + public boolean pause() { + pauseInvocations++; + if (cancelled) { + return false; + } + + if (index >= decisions.size()) { + return false; + } + + return decisions.get(index++); + } + } + + private static final class TestContext implements RemovedConnectionDrainClassifier.Context { + private final Map connections = new LinkedHashMap<>(); + private final Map connectables = new LinkedHashMap<>(); + private final Map groups = new LinkedHashMap<>(); + + private TestContext addGroup(final String groupId, final String parentGroupId) { + groups.put(groupId, new RemovedConnectionDrainClassifier.LiveProcessGroup(groupId, parentGroupId)); + return this; + } + + private TestContext addConnection(final String connectionId, final String sourceId, final String destinationId, final boolean queueEmpty) { + connections.put(connectionId, new RemovedConnectionDrainClassifier.LiveConnection(connectionId, sourceId, destinationId, queueEmpty)); + return this; + } + + private TestContext addProcessor(final String connectableId, final String groupId, + final ScheduledState physicalScheduledState, final ValidationStatus validationStatus) { + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectableId, ConnectableType.PROCESSOR, + groupId, physicalScheduledState, validationStatus, false, Set.of(), Set.of())); + return this; + } + + private TestContext addConnectable(final String connectableId, final ConnectableType connectableType, final String groupId, + final ScheduledState physicalScheduledState, final ValidationStatus validationStatus, + final boolean running) { + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectableId, connectableType, + groupId, physicalScheduledState, validationStatus, running, Set.of(), Set.of())); + return this; + } + + private TestContext linkIncoming(final String connectableId, final Set incomingConnectionIds) { + final RemovedConnectionDrainClassifier.LiveConnectable connectable = connectables.get(connectableId); + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectable.id(), connectable.type(), connectable.processGroupId(), + connectable.physicalScheduledState(), connectable.validationStatus(), connectable.running(), incomingConnectionIds, + connectable.outgoingConnectionIds())); + return this; + } + + private TestContext linkOutgoing(final String connectableId, final Set outgoingConnectionIds) { + final RemovedConnectionDrainClassifier.LiveConnectable connectable = connectables.get(connectableId); + connectables.put(connectableId, new RemovedConnectionDrainClassifier.LiveConnectable(connectable.id(), connectable.type(), connectable.processGroupId(), + connectable.physicalScheduledState(), connectable.validationStatus(), connectable.running(), connectable.incomingConnectionIds(), + outgoingConnectionIds)); + return this; + } + + @Override + public RemovedConnectionDrainClassifier.LiveConnection getConnection(final String connectionId) { + return connections.get(connectionId); + } + + @Override + public RemovedConnectionDrainClassifier.LiveConnectable getConnectable(final String connectableId) { + return connectables.get(connectableId); + } + + @Override + public RemovedConnectionDrainClassifier.LiveProcessGroup getProcessGroup(final String processGroupId) { + return groups.get(processGroupId); + } + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/StandardNiFiServiceFacadeTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/StandardNiFiServiceFacadeTest.java index bc322addf402..d241df4379f0 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/StandardNiFiServiceFacadeTest.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/StandardNiFiServiceFacadeTest.java @@ -28,7 +28,10 @@ import org.apache.nifi.authorization.AuthorizationResult; import org.apache.nifi.authorization.Authorizer; import org.apache.nifi.authorization.ComponentAuthorizable; +import org.apache.nifi.authorization.ConnectionAuthorizable; import org.apache.nifi.authorization.Group; +import org.apache.nifi.authorization.ProcessGroupAuthorizable; +import org.apache.nifi.authorization.RequestAction; import org.apache.nifi.authorization.Resource; import org.apache.nifi.authorization.User; import org.apache.nifi.authorization.resource.Authorizable; @@ -53,6 +56,8 @@ import org.apache.nifi.components.state.StateManagerProvider; import org.apache.nifi.components.state.StateMap; import org.apache.nifi.components.validation.ValidationStatus; +import org.apache.nifi.connectable.Connectable; +import org.apache.nifi.connectable.ConnectableType; import org.apache.nifi.controller.ClusterTopologyProvider; import org.apache.nifi.controller.ControllerService; import org.apache.nifi.controller.Counter; @@ -63,10 +68,13 @@ import org.apache.nifi.controller.service.ControllerServiceNode; import org.apache.nifi.controller.service.ControllerServiceProvider; import org.apache.nifi.controller.status.ProcessGroupStatus; +import org.apache.nifi.flow.ConnectableComponent; +import org.apache.nifi.flow.ConnectableComponentType; import org.apache.nifi.flow.ExecutionEngine; import org.apache.nifi.flow.ExternalControllerServiceReference; import org.apache.nifi.flow.ParameterProviderReference; import org.apache.nifi.flow.VersionedComponent; +import org.apache.nifi.flow.VersionedConnection; import org.apache.nifi.flow.VersionedControllerService; import org.apache.nifi.flow.VersionedParameterContext; import org.apache.nifi.flow.VersionedProcessGroup; @@ -105,8 +113,12 @@ import org.apache.nifi.registry.flow.diff.FlowComparison; import org.apache.nifi.registry.flow.diff.StandardComparableDataFlow; import org.apache.nifi.registry.flow.diff.StandardFlowComparator; +import org.apache.nifi.registry.flow.diff.StandardFlowDifference; import org.apache.nifi.registry.flow.diff.StaticDifferenceDescriptor; import org.apache.nifi.registry.flow.mapping.FlowMappingOptions; +import org.apache.nifi.registry.flow.mapping.InstantiatedConnectableComponent; +import org.apache.nifi.registry.flow.mapping.InstantiatedVersionedConnection; +import org.apache.nifi.registry.flow.mapping.InstantiatedVersionedPort; import org.apache.nifi.registry.flow.mapping.InstantiatedVersionedProcessGroup; import org.apache.nifi.registry.flow.mapping.VersionedComponentFlowMapper; import org.apache.nifi.reporting.Bulletin; @@ -194,6 +206,7 @@ import java.util.HashSet; import java.util.Iterator; import java.util.LinkedHashMap; +import java.util.LinkedHashSet; import java.util.List; import java.util.Map; import java.util.Optional; @@ -270,6 +283,7 @@ public class StandardNiFiServiceFacadeTest { private StandardNiFiServiceFacade serviceFacade; private Authorizer authorizer; + private AuthorizableLookup authorizableLookup; private FlowController flowController; private ProcessGroupDAO processGroupDAO; private ConnectorManagedComponentLookup connectorManagedComponentLookup; @@ -298,7 +312,7 @@ public void setUp() throws Exception { }); // authorizable lookup - final AuthorizableLookup authorizableLookup = mock(AuthorizableLookup.class); + authorizableLookup = mock(AuthorizableLookup.class); final Answer processorLookupAnswer = getProcessorInvocation -> { final String processorId = getProcessorInvocation.getArgument(0); @@ -375,14 +389,19 @@ public Resource getResource() { ruleViolationsManager = mock(RuleViolationsManager.class); connectorManagedComponentLookup = mock(ConnectorManagedComponentLookup.class); + final DtoFactory dtoFactory = new DtoFactory(); + dtoFactory.setAuthorizer(authorizer); + dtoFactory.setEntityFactory(new EntityFactory()); + serviceFacade = new StandardNiFiServiceFacade(); serviceFacade.setAuditService(auditService); serviceFacade.setAuthorizableLookup(authorizableLookup); serviceFacade.setAuthorizer(authorizer); serviceFacade.setEntityFactory(new EntityFactory()); - serviceFacade.setDtoFactory(new DtoFactory()); + serviceFacade.setDtoFactory(dtoFactory); serviceFacade.setControllerFacade(controllerFacade); serviceFacade.setProcessGroupDAO(processGroupDAO); + serviceFacade.setRevisionManager(new NaiveRevisionManager()); serviceFacade.setRuleViolationsManager(ruleViolationsManager); serviceFacade.setConnectorManagedComponentLookup(connectorManagedComponentLookup); @@ -445,6 +464,143 @@ public void testGetComponentsAffectedByFlowUpdate_WithNewStatelessProcessGroup_R assertTrue(affected.isEmpty(), "No local components should be affected for added Stateless group"); } + @Test + public void testGetFlowUpdateImpactExtractsRemovedConnectionMetadataAndAffectedProjection() { + final String rootGroupId = "root-group-instance"; + final String removedGroupId = "removed-group-instance"; + + final ProcessGroup rootGroup = mock(ProcessGroup.class); + when(processGroupDAO.getProcessGroup(rootGroupId)).thenReturn(rootGroup); + + final ProcessGroup rootGroupMetadata = mock(ProcessGroup.class); + when(rootGroupMetadata.getIdentifier()).thenReturn(rootGroupId); + when(rootGroupMetadata.getName()).thenReturn("Root Group"); + + final ProcessGroup removedGroup = mock(ProcessGroup.class); + when(removedGroup.getIdentifier()).thenReturn(removedGroupId); + when(processGroupDAO.getProcessGroup(removedGroupId)).thenReturn(removedGroup); + when(removedGroup.findAllProcessors()).thenReturn(List.of()); + when(removedGroup.findAllFunnels()).thenReturn(List.of()); + when(removedGroup.findAllInputPorts()).thenReturn(List.of()); + when(removedGroup.findAllOutputPorts()).thenReturn(List.of()); + when(removedGroup.findAllRemoteProcessGroups()).thenReturn(List.of()); + when(removedGroup.findAllControllerServices()).thenReturn(Set.of()); + + final ProcessorNode removedSource = createLocalConnectableProcessor( + "removed-source-instance", "removed-source-versioned", ConnectableType.PROCESSOR, rootGroupMetadata); + final org.apache.nifi.connectable.Port removedDestination = createLocalConnectablePort( + "removed-destination-instance", "removed-destination-versioned", ConnectableType.OUTPUT_PORT, rootGroupMetadata); + final org.apache.nifi.connectable.Port changedSource = createLocalConnectablePort( + "changed-source-instance", "changed-source-versioned", ConnectableType.INPUT_PORT, rootGroupMetadata); + final ProcessorNode changedDestination = createLocalConnectableProcessor( + "changed-destination-instance", "changed-destination-versioned", ConnectableType.PROCESSOR, rootGroupMetadata); + + when(rootGroup.findAllProcessors()).thenReturn(List.of(removedSource, changedDestination)); + when(rootGroup.findAllFunnels()).thenReturn(List.of()); + when(rootGroup.findAllInputPorts()).thenReturn(List.of(changedSource)); + when(rootGroup.findAllOutputPorts()).thenReturn(List.of(removedDestination)); + when(rootGroup.findAllRemoteProcessGroups()).thenReturn(List.of()); + + stubLocalConnectableAuthorizable(removedSource.getIdentifier()); + stubLocalConnectableAuthorizable(removedDestination.getIdentifier()); + stubLocalConnectableAuthorizable(changedSource.getIdentifier()); + stubLocalConnectableAuthorizable(changedDestination.getIdentifier()); + stubConnectionAuthorizable("removed-connection-instance"); + stubConnectionAuthorizable("changed-connection-instance"); + stubProcessGroupAuthorizable(removedGroupId, removedGroup); + stubInputPortAuthorizable("removed-endpoint-instance"); + + final FlowComparison comparison = mock(FlowComparison.class); + when(comparison.getDifferences()).thenReturn(Set.of( + new StandardFlowDifference(DifferenceType.COMPONENT_REMOVED, createRemovedConnection(), null, null, null, "Removed connection"), + new StandardFlowDifference(DifferenceType.SOURCE_CHANGED, createSourceChangedConnection(), createReplacementConnection(), null, null, "Source changed connection"), + new StandardFlowDifference(DifferenceType.COMPONENT_REMOVED, createRemovedGroup(removedGroupId), null, null, null, "Removed group"), + new StandardFlowDifference(DifferenceType.COMPONENT_REMOVED, createRemovedEndpoint(removedGroupId), null, null, null, "Removed endpoint") + )); + + final StandardNiFiServiceFacade serviceFacadeSpy = spy(serviceFacade); + doReturn(comparison).when(serviceFacadeSpy).compareFlowUpdate(eq(rootGroup), any(RegisteredFlowSnapshot.class)); + + final RegisteredFlowSnapshot updatedSnapshot = new RegisteredFlowSnapshot(); + final FlowUpdateImpact impact = serviceFacadeSpy.getFlowUpdateImpact(rootGroupId, updatedSnapshot); + + assertNotNull(impact); + assertEquals(Set.of(removedGroupId), impact.getRemovedProcessGroupIds()); + assertEquals(Set.of("removed-endpoint-instance"), impact.getRemovedEndpointIds()); + + final Map removedConnectionByInstanceId = impact.getRemovedConnections().stream() + .collect(Collectors.toMap(RemovedConnectionDescriptor::getConnectionInstanceId, Function.identity())); + assertEquals(Set.of("removed-connection-instance", "changed-connection-instance"), removedConnectionByInstanceId.keySet()); + + final RemovedConnectionDescriptor pureRemoval = removedConnectionByInstanceId.get("removed-connection-instance"); + assertEquals(RemovalReason.COMPONENT_REMOVED, pureRemoval.getRemovalReason()); + assertEquals("removed-connection-versioned", pureRemoval.getConnectionVersionedId()); + assertEquals("root-group-instance", pureRemoval.getContainingProcessGroupId()); + assertEquals("removed-source-instance", pureRemoval.getSourceInstanceId()); + assertEquals("removed-source-versioned", pureRemoval.getSourceVersionedId()); + assertEquals("removed-source-group-instance", pureRemoval.getSourceProcessGroupId()); + assertEquals(ConnectableType.PROCESSOR, pureRemoval.getSourceType()); + assertEquals("removed-destination-instance", pureRemoval.getDestinationInstanceId()); + assertEquals("removed-destination-versioned", pureRemoval.getDestinationVersionedId()); + assertEquals("removed-destination-group-instance", pureRemoval.getDestinationProcessGroupId()); + assertEquals(ConnectableType.OUTPUT_PORT, pureRemoval.getDestinationType()); + + final RemovedConnectionDescriptor sourceChanged = removedConnectionByInstanceId.get("changed-connection-instance"); + assertEquals(RemovalReason.SOURCE_CHANGED, sourceChanged.getRemovalReason()); + assertEquals("changed-source-instance", sourceChanged.getSourceInstanceId()); + assertEquals("changed-destination-instance", sourceChanged.getDestinationInstanceId()); + + final Set affectedIdsFromImpact = impact.getAffectedComponents().stream() + .map(AffectedComponentEntity::getId) + .collect(Collectors.toCollection(LinkedHashSet::new)); + assertTrue(affectedIdsFromImpact.contains("removed-source-instance")); + assertTrue(affectedIdsFromImpact.contains("removed-destination-instance")); + assertTrue(affectedIdsFromImpact.contains("changed-source-instance")); + assertTrue(affectedIdsFromImpact.contains("changed-destination-instance")); + + final Set affectedIdsFromProjection = serviceFacadeSpy.getComponentsAffectedByFlowUpdate(rootGroupId, updatedSnapshot).stream() + .map(AffectedComponentEntity::getId) + .collect(Collectors.toCollection(LinkedHashSet::new)); + assertEquals(affectedIdsFromImpact, affectedIdsFromProjection); + } + + @Test + public void testGetFlowUpdateImpactUsesInstanceIdentifiersForRemovedRuntimeIdentities() { + final String rootGroupId = "root-group-instance"; + final ProcessGroup rootGroup = mock(ProcessGroup.class); + when(processGroupDAO.getProcessGroup(rootGroupId)).thenReturn(rootGroup); + when(rootGroup.findAllProcessors()).thenReturn(List.of()); + when(rootGroup.findAllFunnels()).thenReturn(List.of()); + when(rootGroup.findAllInputPorts()).thenReturn(List.of()); + when(rootGroup.findAllOutputPorts()).thenReturn(List.of()); + when(rootGroup.findAllRemoteProcessGroups()).thenReturn(List.of()); + + final InstantiatedVersionedProcessGroup removedGroup = createRemovedGroup("runtime-group-id"); + removedGroup.setIdentifier("versioned-group-id"); + + final InstantiatedVersionedPort removedEndpoint = createRemovedEndpoint("runtime-group-id"); + removedEndpoint.setIdentifier("versioned-endpoint-id"); + + stubProcessGroupAuthorizable("runtime-group-id", mock(ProcessGroup.class)); + stubInputPortAuthorizable("removed-endpoint-instance"); + + final FlowComparison comparison = mock(FlowComparison.class); + when(comparison.getDifferences()).thenReturn(Set.of( + new StandardFlowDifference(DifferenceType.COMPONENT_REMOVED, removedGroup, null, null, null, "Removed group"), + new StandardFlowDifference(DifferenceType.COMPONENT_REMOVED, removedEndpoint, null, null, null, "Removed endpoint") + )); + + final StandardNiFiServiceFacade serviceFacadeSpy = spy(serviceFacade); + doReturn(comparison).when(serviceFacadeSpy).compareFlowUpdate(eq(rootGroup), any(RegisteredFlowSnapshot.class)); + + final FlowUpdateImpact impact = serviceFacadeSpy.getFlowUpdateImpact(rootGroupId, new RegisteredFlowSnapshot()); + + assertEquals(Set.of("runtime-group-id"), impact.getRemovedProcessGroupIds()); + assertEquals(Set.of("removed-endpoint-instance"), impact.getRemovedEndpointIds()); + assertFalse(impact.getRemovedProcessGroupIds().contains("versioned-group-id")); + assertFalse(impact.getRemovedEndpointIds().contains("versioned-endpoint-id")); + } + private FlowChangeAction getAction(final Integer actionId, final String processorId) { final FlowChangeAction action = new FlowChangeAction(); action.setId(actionId); @@ -454,6 +610,116 @@ private FlowChangeAction getAction(final Integer actionId, final String processo return action; } + private ProcessorNode createLocalConnectableProcessor(final String instanceId, final String versionedId, + final ConnectableType connectableType, final ProcessGroup processGroup) { + final ProcessorNode connectable = mock(ProcessorNode.class); + stubLocalConnectable(connectable, instanceId, versionedId, connectableType, processGroup); + return connectable; + } + + private org.apache.nifi.connectable.Port createLocalConnectablePort(final String instanceId, final String versionedId, + final ConnectableType connectableType, final ProcessGroup processGroup) { + final org.apache.nifi.connectable.Port connectable = mock(org.apache.nifi.connectable.Port.class); + stubLocalConnectable(connectable, instanceId, versionedId, connectableType, processGroup); + return connectable; + } + + private void stubLocalConnectable(final Connectable connectable, final String instanceId, final String versionedId, + final ConnectableType connectableType, final ProcessGroup processGroup) { + final String processGroupId = processGroup.getIdentifier(); + when(connectable.getIdentifier()).thenReturn(instanceId); + when(connectable.getVersionedComponentId()).thenReturn(Optional.of(versionedId)); + when(connectable.getConnectableType()).thenReturn(connectableType); + when(connectable.getProcessGroup()).thenReturn(processGroup); + when(connectable.getProcessGroupIdentifier()).thenReturn(processGroupId); + when(connectable.getScheduledState()).thenReturn(org.apache.nifi.controller.ScheduledState.RUNNING); + when(connectable.getName()).thenReturn(instanceId); + } + + private void stubLocalConnectableAuthorizable(final String componentId) { + final Authorizable authorizable = mock(Authorizable.class); + when(authorizable.isAuthorized(any(Authorizer.class), any(RequestAction.class), any())).thenReturn(true); + when(authorizableLookup.getLocalConnectable(componentId)).thenReturn(authorizable); + } + + private void stubInputPortAuthorizable(final String componentId) { + final Authorizable authorizable = mock(Authorizable.class); + when(authorizable.isAuthorized(any(Authorizer.class), any(RequestAction.class), any())).thenReturn(true); + when(authorizableLookup.getInputPort(componentId)).thenReturn(authorizable); + } + + private void stubConnectionAuthorizable(final String connectionId) { + final ConnectionAuthorizable connectionAuthorizable = mock(ConnectionAuthorizable.class); + final Authorizable authorizable = mock(Authorizable.class); + when(authorizable.isAuthorized(any(Authorizer.class), any(RequestAction.class), any())).thenReturn(true); + when(connectionAuthorizable.getAuthorizable()).thenReturn(authorizable); + when(authorizableLookup.getConnection(connectionId)).thenReturn(connectionAuthorizable); + } + + private void stubProcessGroupAuthorizable(final String groupId, final ProcessGroup processGroup) { + final ProcessGroupAuthorizable processGroupAuthorizable = mock(ProcessGroupAuthorizable.class); + final Authorizable authorizable = mock(Authorizable.class); + when(authorizable.isAuthorized(any(Authorizer.class), any(RequestAction.class), any())).thenReturn(true); + when(processGroupAuthorizable.getAuthorizable()).thenReturn(authorizable); + when(processGroupAuthorizable.getProcessGroup()).thenReturn(processGroup); + when(authorizableLookup.getProcessGroup(groupId)).thenReturn(processGroupAuthorizable); + } + + private InstantiatedVersionedConnection createRemovedConnection() { + final InstantiatedVersionedConnection connection = new InstantiatedVersionedConnection("removed-connection-instance", "root-group-instance"); + connection.setIdentifier("removed-connection-versioned"); + connection.setSource(createConnectableComponent("removed-source-instance", "removed-source-versioned", + "removed-source-group-instance", "removed-source-group-versioned", ConnectableType.PROCESSOR)); + connection.setDestination(createConnectableComponent("removed-destination-instance", "removed-destination-versioned", + "removed-destination-group-instance", "removed-destination-group-versioned", ConnectableType.OUTPUT_PORT)); + return connection; + } + + private InstantiatedVersionedConnection createSourceChangedConnection() { + final InstantiatedVersionedConnection connection = new InstantiatedVersionedConnection("changed-connection-instance", "root-group-instance"); + connection.setIdentifier("changed-connection-versioned"); + connection.setSource(createConnectableComponent("changed-source-instance", "changed-source-versioned", + "changed-source-group-instance", "changed-source-group-versioned", ConnectableType.INPUT_PORT)); + connection.setDestination(createConnectableComponent("changed-destination-instance", "changed-destination-versioned", + "changed-destination-group-instance", "changed-destination-group-versioned", ConnectableType.PROCESSOR)); + return connection; + } + + private VersionedConnection createReplacementConnection() { + final VersionedConnection connection = new VersionedConnection(); + connection.setIdentifier("replacement-connection-versioned"); + connection.setSource(createConnectableComponent("replacement-source-instance", "replacement-source-versioned", + "replacement-source-group-instance", "replacement-source-group-versioned", ConnectableType.PROCESSOR)); + connection.setDestination(createConnectableComponent("changed-destination-instance", "changed-destination-versioned", + "changed-destination-group-instance", "changed-destination-group-versioned", ConnectableType.PROCESSOR)); + return connection; + } + + private InstantiatedVersionedProcessGroup createRemovedGroup(final String instanceId) { + final InstantiatedVersionedProcessGroup group = new InstantiatedVersionedProcessGroup(instanceId, "root-group-instance"); + group.setIdentifier("removed-group-versioned"); + group.setComponentType(org.apache.nifi.flow.ComponentType.PROCESS_GROUP); + return group; + } + + private InstantiatedVersionedPort createRemovedEndpoint(final String parentGroupId) { + final InstantiatedVersionedPort port = new InstantiatedVersionedPort("removed-endpoint-instance", parentGroupId); + port.setIdentifier("removed-endpoint-versioned"); + port.setGroupIdentifier("removed-endpoint-group-versioned"); + port.setComponentType(org.apache.nifi.flow.ComponentType.INPUT_PORT); + return port; + } + + private ConnectableComponent createConnectableComponent(final String instanceId, final String versionedId, + final String instanceGroupId, final String versionedGroupId, + final ConnectableType connectableType) { + final InstantiatedConnectableComponent component = new InstantiatedConnectableComponent(instanceId, instanceGroupId); + component.setId(versionedId); + component.setGroupId(versionedGroupId); + component.setType(ConnectableComponentType.valueOf(connectableType.name())); + return component; + } + @Test public void testGetUnknownAction() { assertThrows(ResourceNotFoundException.class, () -> serviceFacade.getAction(UNKNOWN_ACTION_ID)); diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/FlowUpdateResourceTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/FlowUpdateResourceTest.java new file mode 100644 index 000000000000..c86d100764ee --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/FlowUpdateResourceTest.java @@ -0,0 +1,57 @@ +/* + * 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.web.api; + +import org.apache.nifi.web.api.concurrent.UpdateStep; +import org.junit.jupiter.api.Test; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +class FlowUpdateResourceTest { + private static final List STANDARD_STEPS = List.of( + "Stopping Affected Processors", + "Disabling Affected Controller Services", + "Updating Flow", + "Re-Enabling Controller Services", + "Restarting Affected Processors"); + + @Test + void testRegistryUpdateHasRemovedConnectionDrainStep() { + assertEquals(List.of( + "Draining Removed Connections", + "Stopping Affected Processors", + "Disabling Affected Controller Services", + "Updating Flow", + "Re-Enabling Controller Services", + "Restarting Affected Processors"), getStepDescriptions(FlowUpdateResource.UPDATE_REQUEST_TYPE)); + } + + @Test + void testNonUpdateRequestsRetainStandardSteps() { + assertEquals(STANDARD_STEPS, getStepDescriptions("revert-requests")); + assertEquals(STANDARD_STEPS, getStepDescriptions("rebase-requests")); + assertEquals(STANDARD_STEPS, getStepDescriptions("replace-requests")); + } + + private List getStepDescriptions(final String requestType) { + return FlowUpdateResource.getUpdateFlowSteps(requestType).stream() + .map(UpdateStep::getDescription) + .toList(); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequestTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequestTest.java index 80157dbcd3c8..2953a8c9965a 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequestTest.java +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/api/concurrent/StandardAsynchronousWebRequestTest.java @@ -55,6 +55,39 @@ public void testCancelInvokesCancelCallback() { assertEquals("Request cancelled by user", request.getFailureReason()); } + @Test + public void testAppendFailureDetailAfterCancelRetainsCancellationReason() { + final StandardAsynchronousWebRequest request = createRequest(); + + request.cancel(); + request.appendFailureDetail("restoration failed: component could not be started"); + + assertEquals("Request cancelled by user; restoration failed: component could not be started", request.getFailureReason()); + assertTrue(request.isCancelled()); + assertTrue(request.isComplete()); + } + + @Test + public void testAppendFailureDetailEstablishesFailureReason() { + final StandardAsynchronousWebRequest request = createRequest(); + + request.appendFailureDetail("component could not be started"); + + assertEquals("component could not be started", request.getFailureReason()); + assertTrue(request.isComplete()); + assertFalse(request.isCancelled()); + } + + @Test + public void testAppendFailureDetailAppendsToExistingFailureReason() { + final StandardAsynchronousWebRequest request = createRequest(); + + request.fail("operation failed"); + request.appendFailureDetail("component could not be started"); + + assertEquals("operation failed; component could not be started", request.getFailureReason()); + } + @Test public void testFailBeforeCancelSetsFailureReason() { final StandardAsynchronousWebRequest request = createRequest(); diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycleTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycleTest.java new file mode 100644 index 000000000000..766adfbef4f8 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/util/ClusterReplicationComponentLifecycleTest.java @@ -0,0 +1,404 @@ +/* + * 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.web.util; + +import jakarta.ws.rs.HttpMethod; +import jakarta.ws.rs.core.Response; +import org.apache.nifi.authorization.user.NiFiUser; +import org.apache.nifi.authorization.user.NiFiUserDetails; +import org.apache.nifi.authorization.user.StandardNiFiUser; +import org.apache.nifi.cluster.coordination.ClusterCoordinator; +import org.apache.nifi.cluster.coordination.http.replication.AsyncClusterResponse; +import org.apache.nifi.cluster.coordination.http.replication.RequestReplicator; +import org.apache.nifi.cluster.coordination.node.NodeConnectionState; +import org.apache.nifi.cluster.manager.NodeResponse; +import org.apache.nifi.cluster.protocol.NodeIdentifier; +import org.apache.nifi.web.api.dto.ListingRequestDTO; +import org.apache.nifi.web.api.dto.QueueSizeDTO; +import org.apache.nifi.web.api.entity.ListingRequestEntity; +import org.apache.nifi.web.security.token.NiFiAuthenticationToken; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import org.springframework.security.core.context.SecurityContextHolder; + +import java.net.URI; +import java.util.Collections; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class ClusterReplicationComponentLifecycleTest { + private static final NodeIdentifier NODE_1 = new NodeIdentifier("node-1", "localhost", 8081, "localhost", 9081, "localhost", 10081, 11081, false); + private static final NodeIdentifier NODE_2 = new NodeIdentifier("node-2", "localhost", 8082, "localhost", 9082, "localhost", 10082, 11082, false); + private static final Set EXPECTED_NODES = Set.of(NODE_1, NODE_2); + private static final URI EXAMPLE_URI = URI.create("http://localhost:8080/nifi-api/flow/connections/connection-a/status"); + + @Mock + private ClusterCoordinator clusterCoordinator; + @Mock + private RequestReplicator requestReplicator; + private NiFiUser user; + + @BeforeEach + void setUpCurrentUser() { + user = new StandardNiFiUser.Builder().identity("unit-test-user").build(); + SecurityContextHolder.getContext().setAuthentication(new NiFiAuthenticationToken(new NiFiUserDetails(user))); + } + + @AfterEach + void clearCurrentUser() { + SecurityContextHolder.clearContext(); + } + + @Test + void testWaitForConnectionQueuesEmptyReturnsTrueOnlyWhenAllConnectionsAreZeroAcrossTargetNodes() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(true); + final AsyncClusterResponse firstConnectionNonEmpty = asyncResponse(mergedNodeResponse("connection-a", 1), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse firstConnectionDelete = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + final AsyncClusterResponse firstConnectionEmpty = asyncResponse(mergedNodeResponse("connection-a", 0), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse firstConnectionSecondDelete = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + final AsyncClusterResponse secondConnectionEmpty = asyncResponse(mergedNodeResponse("connection-b", 0), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse secondConnectionDelete = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, firstConnectionNonEmpty, firstConnectionEmpty, secondConnectionEmpty); + stubReplicate(HttpMethod.DELETE, firstConnectionDelete, firstConnectionSecondDelete, secondConnectionDelete); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a", "connection-b"), pause); + + assertTrue(result); + verify(clusterCoordinator).getNodeIdentifiers(NodeConnectionState.CONNECTED); + verifyReplicate(HttpMethod.POST, 3); + verifyReplicate(HttpMethod.DELETE, 3); + verify(requestReplicator, never()).forwardToCoordinator(eq(NODE_1), eq(user), any(), any(URI.class), any(), any()); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsMissingNodeCoverage() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 0), completedResponses(successfulNodeResponse(NODE_1))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsExtraNodeCoverage() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final NodeIdentifier extraNode = new NodeIdentifier("node-3", "localhost", 8083, "localhost", 9083, "localhost", 10083, 11083, false); + final AsyncClusterResponse createResponse = asyncResponse( + mergedNodeResponse("connection-a", 0), + completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2), successfulNodeResponse(extraNode))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2), deleteNodeResponse(extraNode))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsDuplicateNodeCoverage() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final NodeIdentifier duplicateNode = new NodeIdentifier("node-1", "localhost", 8083, "localhost", 9083, "localhost", 10083, 11083, false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 0), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(duplicateNode))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(duplicateNode))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsProblematicNodeResponses() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 0), completedResponses(successfulNodeResponse(NODE_1), failedNodeResponse(NODE_2, 500))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsIncompleteAsyncResponse() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse( + mergedNodeResponse("connection-a", 0), + completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2)), false); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsCompletedIdentifiersThatMissExpectedNodeId() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final NodeIdentifier unexpectedNode = new NodeIdentifier("node-3", "localhost", 8083, "localhost", 9083, "localhost", 10083, 11083, false); + final AsyncClusterResponse createResponse = asyncResponse( + mergedNodeResponse("connection-a", 0), + Set.of(NODE_1, unexpectedNode), + completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyRejectsNonEmptyAggregateEvenWithFullCoverage() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 1), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + } + + @Test + void testWaitForConnectionQueuesEmptyReturnsFalseWhenPauseCancelsPolling() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 1), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertFalse(result); + assertTrue(pause.wasInvoked()); + } + + @Test + void testWaitForConnectionQueuesEmptyReplicatesDirectlyToSnapshotNodesWhenCoordinatorActive() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 0), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse deleteResponse = asyncResponse(deleteResponse(), completedResponses(deleteNodeResponse(NODE_1), deleteNodeResponse(NODE_2))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertTrue(result); + verifyReplicate(HttpMethod.POST, 1); + verifyReplicate(HttpMethod.DELETE, 1); + } + + @Test + void testWaitForConnectionQueuesEmptyTreatsDeleteCleanupAsBestEffort() throws Exception { + final ClusterReplicationComponentLifecycle lifecycle = createLifecycle(); + final TestPause pause = new TestPause(false); + final AsyncClusterResponse createResponse = asyncResponse(mergedNodeResponse("connection-a", 0), completedResponses(successfulNodeResponse(NODE_1), successfulNodeResponse(NODE_2))); + final AsyncClusterResponse deleteResponse = asyncResponse(failedDeleteResponse(), completedResponses(deleteNodeResponse(NODE_1), failedDeleteNodeResponse(NODE_2, 404))); + + when(clusterCoordinator.getNodeIdentifiers(NodeConnectionState.CONNECTED)).thenReturn(EXPECTED_NODES); + stubReplicate(HttpMethod.POST, createResponse); + stubReplicate(HttpMethod.DELETE, deleteResponse); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(EXAMPLE_URI, Set.of("connection-a"), pause); + + assertTrue(result); + verifyReplicate(HttpMethod.POST, 1); + verifyReplicate(HttpMethod.DELETE, 1); + } + + private ClusterReplicationComponentLifecycle createLifecycle() { + final ClusterReplicationComponentLifecycle lifecycle = new ClusterReplicationComponentLifecycle(); + lifecycle.setClusterCoordinator(clusterCoordinator); + lifecycle.setRequestReplicator(requestReplicator); + lenient().when(clusterCoordinator.getElectedActiveCoordinatorNode()).thenReturn(NODE_1); + return lifecycle; + } + + private void stubReplicate(final String method, final AsyncClusterResponse... responses) { + when(requestReplicator.replicate(eq(EXPECTED_NODES), eq(user), eq(method), any(URI.class), eq(Collections.emptyMap()), eq(Collections.emptyMap()), eq(true), eq(true))) + .thenReturn(responses[0], java.util.Arrays.copyOfRange(responses, 1, responses.length)); + } + + private void verifyReplicate(final String method, final int times) { + verify(requestReplicator, times(times)).replicate(eq(EXPECTED_NODES), eq(user), eq(method), any(URI.class), eq(Collections.emptyMap()), eq(Collections.emptyMap()), eq(true), eq(true)); + } + + private AsyncClusterResponse asyncResponse(final NodeResponse mergedResponse, final Set completedResponses) throws Exception { + final Set completedNodeIdentifiers = completedResponses.stream() + .map(NodeResponse::getNodeId) + .collect(Collectors.toUnmodifiableSet()); + return asyncResponse(mergedResponse, completedNodeIdentifiers, completedResponses, true); + } + + private AsyncClusterResponse asyncResponse(final NodeResponse mergedResponse, final Set completedResponses, final boolean complete) throws Exception { + final Set completedNodeIdentifiers = completedResponses.stream() + .map(NodeResponse::getNodeId) + .collect(Collectors.toUnmodifiableSet()); + return asyncResponse(mergedResponse, completedNodeIdentifiers, completedResponses, complete); + } + + private AsyncClusterResponse asyncResponse(final NodeResponse mergedResponse, final Set completedNodeIdentifiers, + final Set completedResponses) throws Exception { + return asyncResponse(mergedResponse, completedNodeIdentifiers, completedResponses, true); + } + + private AsyncClusterResponse asyncResponse(final NodeResponse mergedResponse, final Set completedNodeIdentifiers, + final Set completedResponses, final boolean complete) throws Exception { + final AsyncClusterResponse asyncResponse = mock(AsyncClusterResponse.class); + when(asyncResponse.awaitMergedResponse()).thenReturn(mergedResponse); + lenient().when(asyncResponse.getNodesInvolved()).thenReturn(completedNodeIdentifiers); + lenient().when(asyncResponse.getCompletedNodeIdentifiers()).thenReturn(completedNodeIdentifiers); + lenient().when(asyncResponse.getCompletedNodeResponses()).thenReturn(completedResponses); + lenient().when(asyncResponse.isComplete()).thenReturn(complete); + return asyncResponse; + } + + private Set completedResponses(final NodeResponse... nodeResponses) { + final Set completedResponses = new LinkedHashSet<>(); + for (final NodeResponse nodeResponse : nodeResponses) { + completedResponses.add(nodeResponse); + } + return completedResponses; + } + + private NodeResponse mergedNodeResponse(final String connectionId, final int aggregateQueued) { + final QueueSizeDTO queueSize = new QueueSizeDTO(); + queueSize.setObjectCount(aggregateQueued); + + final ListingRequestDTO listingRequest = new ListingRequestDTO(); + listingRequest.setId(connectionId + "-request"); + listingRequest.setQueueSize(queueSize); + + final ListingRequestEntity entity = new ListingRequestEntity(); + entity.setListingRequest(listingRequest); + + final Response response = Response.accepted(entity).build(); + final NodeResponse nodeResponse = new NodeResponse(NODE_1, HttpMethod.POST, EXAMPLE_URI, response, 0L, connectionId); + return new NodeResponse(nodeResponse, entity); + } + + private NodeResponse deleteResponse() { + final Response response = Response.ok().build(); + return new NodeResponse(NODE_1, HttpMethod.DELETE, EXAMPLE_URI, response, 0L, "delete"); + } + + private NodeResponse failedDeleteResponse() { + final Response response = Response.status(404).build(); + return new NodeResponse(NODE_1, HttpMethod.DELETE, EXAMPLE_URI, response, 0L, "delete"); + } + + private NodeResponse successfulNodeResponse(final NodeIdentifier nodeIdentifier) { + return new NodeResponse(nodeIdentifier, HttpMethod.POST, EXAMPLE_URI, Response.accepted().build(), 0L, nodeIdentifier.getId()); + } + + private NodeResponse deleteNodeResponse(final NodeIdentifier nodeIdentifier) { + return new NodeResponse(nodeIdentifier, HttpMethod.DELETE, EXAMPLE_URI, Response.ok().build(), 0L, nodeIdentifier.getId() + "-delete"); + } + + private NodeResponse failedDeleteNodeResponse(final NodeIdentifier nodeIdentifier, final int status) { + return new NodeResponse(nodeIdentifier, HttpMethod.DELETE, EXAMPLE_URI, Response.status(status).build(), 0L, nodeIdentifier.getId() + "-delete"); + } + + private NodeResponse failedNodeResponse(final NodeIdentifier nodeIdentifier, final int status) { + return new NodeResponse(nodeIdentifier, HttpMethod.POST, EXAMPLE_URI, Response.status(status).build(), 0L, nodeIdentifier.getId()); + } + + private static final class TestPause implements Pause { + private final List decisions; + private int index = 0; + private boolean invoked; + + private TestPause(final Boolean... decisions) { + this.decisions = List.of(decisions); + } + + @Override + public boolean pause() { + invoked = true; + if (index >= decisions.size()) { + return false; + } + + return decisions.get(index++); + } + + private boolean wasInvoked() { + return invoked; + } + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/util/LocalComponentLifecycleTest.java b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/util/LocalComponentLifecycleTest.java new file mode 100644 index 000000000000..8b7edf9dafd9 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-web/nifi-web-api/src/test/java/org/apache/nifi/web/util/LocalComponentLifecycleTest.java @@ -0,0 +1,141 @@ +/* + * 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.web.util; + +import org.apache.nifi.web.NiFiServiceFacade; +import org.apache.nifi.web.api.dto.ListingRequestDTO; +import org.apache.nifi.web.api.dto.QueueSizeDTO; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.net.URI; +import java.util.List; +import java.util.Set; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class LocalComponentLifecycleTest { + @Mock + private NiFiServiceFacade serviceFacade; + + @Test + void testWaitForConnectionQueuesEmptyReturnsTrueImmediatelyWhenEveryQueueIsEmpty() throws LifecycleManagementException { + final LocalComponentLifecycle lifecycle = new LocalComponentLifecycle(); + lifecycle.setServiceFacade(serviceFacade); + + when(serviceFacade.createFlowFileListingRequest(eq("connection-a"), anyString())).thenReturn(listingRequest(0)); + when(serviceFacade.createFlowFileListingRequest(eq("connection-b"), anyString())).thenReturn(listingRequest(0)); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(URI.create("http://localhost:8080/nifi-api"), Set.of("connection-a", "connection-b"), new TestPause(true)); + + assertTrue(result); + verify(serviceFacade).createFlowFileListingRequest(eq("connection-a"), anyString()); + verify(serviceFacade).createFlowFileListingRequest(eq("connection-b"), anyString()); + verify(serviceFacade).deleteFlowFileListingRequest(eq("connection-a"), anyString()); + verify(serviceFacade).deleteFlowFileListingRequest(eq("connection-b"), anyString()); + } + + @Test + void testWaitForConnectionQueuesEmptyPollsUntilEveryQueueIsEmpty() throws LifecycleManagementException { + final LocalComponentLifecycle lifecycle = new LocalComponentLifecycle(); + lifecycle.setServiceFacade(serviceFacade); + final TestPause pause = new TestPause(true, true); + + when(serviceFacade.createFlowFileListingRequest(eq("connection-a"), anyString())).thenReturn(listingRequest(1), listingRequest(0)); + when(serviceFacade.createFlowFileListingRequest(eq("connection-b"), anyString())).thenReturn(listingRequest(0)); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(URI.create("http://localhost:8080/nifi-api"), Set.of("connection-a", "connection-b"), pause); + + assertTrue(result); + assertTrue(pause.wasInvoked()); + verify(serviceFacade, times(2)).createFlowFileListingRequest(eq("connection-a"), anyString()); + verify(serviceFacade).createFlowFileListingRequest(eq("connection-b"), anyString()); + verify(serviceFacade, times(2)).deleteFlowFileListingRequest(eq("connection-a"), anyString()); + verify(serviceFacade).deleteFlowFileListingRequest(eq("connection-b"), anyString()); + } + + @Test + void testWaitForConnectionQueuesEmptyReturnsFalseWhenPauseStopsBeforeDrainCompletes() throws LifecycleManagementException { + final LocalComponentLifecycle lifecycle = new LocalComponentLifecycle(); + lifecycle.setServiceFacade(serviceFacade); + final TestPause pause = new TestPause(false); + + when(serviceFacade.createFlowFileListingRequest(eq("connection-a"), anyString())).thenReturn(listingRequest(2)); + final boolean result = lifecycle.waitForConnectionQueuesEmpty(URI.create("http://localhost:8080/nifi-api"), Set.of("connection-a", "connection-b"), pause); + + assertFalse(result); + assertTrue(pause.wasInvoked()); + verify(serviceFacade).deleteFlowFileListingRequest(eq("connection-a"), anyString()); + } + + @Test + void testWaitForConnectionQueuesEmptyTreatsUnacknowledgedFlowFilesAsNotEmpty() throws LifecycleManagementException { + final LocalComponentLifecycle lifecycle = new LocalComponentLifecycle(); + lifecycle.setServiceFacade(serviceFacade); + final TestPause pause = new TestPause(false); + + when(serviceFacade.createFlowFileListingRequest(eq("connection-a"), anyString())).thenReturn(listingRequest(1)); + + final boolean result = lifecycle.waitForConnectionQueuesEmpty(URI.create("http://localhost:8080/nifi-api"), Set.of("connection-a"), pause); + + assertFalse(result); + verify(serviceFacade).deleteFlowFileListingRequest(eq("connection-a"), anyString()); + } + + private ListingRequestDTO listingRequest(final int flowFilesQueued) { + final QueueSizeDTO queueSize = new QueueSizeDTO(); + queueSize.setObjectCount(flowFilesQueued); + + final ListingRequestDTO listingRequest = new ListingRequestDTO(); + listingRequest.setQueueSize(queueSize); + return listingRequest; + } + + private static final class TestPause implements Pause { + private final List decisions; + private int index = 0; + private boolean invoked; + + private TestPause(final Boolean... decisions) { + this.decisions = List.of(decisions); + } + + @Override + public boolean pause() { + invoked = true; + if (index >= decisions.size()) { + return false; + } + + return decisions.get(index++); + } + + private boolean wasInvoked() { + return invoked; + } + } +} diff --git a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/GatedPassThrough.java b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/GatedPassThrough.java new file mode 100644 index 000000000000..3ab4ab29906f --- /dev/null +++ b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/java/org/apache/nifi/processors/tests/system/GatedPassThrough.java @@ -0,0 +1,73 @@ +/* + * 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.processors.tests.system; + +import org.apache.nifi.annotation.behavior.InputRequirement; +import org.apache.nifi.components.PropertyDescriptor; +import org.apache.nifi.flowfile.FlowFile; +import org.apache.nifi.processor.AbstractProcessor; +import org.apache.nifi.processor.ProcessContext; +import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; +import org.apache.nifi.processor.exception.ProcessException; +import org.apache.nifi.processor.util.StandardValidators; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Set; + +@InputRequirement(InputRequirement.Requirement.INPUT_ALLOWED) +public class GatedPassThrough extends AbstractProcessor { + static final PropertyDescriptor GATE_FILE = new PropertyDescriptor.Builder() + .name("Gate File") + .description("Absolute path of a file whose presence permits input FlowFiles to be transferred") + .required(true) + .addValidator(StandardValidators.NON_EMPTY_VALIDATOR) + .build(); + + static final Relationship REL_SUCCESS = new Relationship.Builder() + .name("success") + .description("FlowFiles transferred after the gate opens") + .build(); + + @Override + protected List getSupportedPropertyDescriptors() { + return List.of(GATE_FILE); + } + + @Override + public Set getRelationships() { + return Set.of(REL_SUCCESS); + } + + @Override + public void onTrigger(final ProcessContext context, final ProcessSession session) throws ProcessException { + final Path gateFile = Path.of(context.getProperty(GATE_FILE).getValue()); + if (!Files.exists(gateFile)) { + context.yield(); + return; + } + + final FlowFile flowFile = session.get(); + if (flowFile != null) { + session.transfer(flowFile, REL_SUCCESS); + } else { + context.yield(); + } + } +} diff --git a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor index b3167115e443..d64c56133b74 100644 --- a/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor +++ b/nifi-system-tests/nifi-system-test-extensions-bundle/nifi-system-test-extensions/src/main/resources/META-INF/services/org.apache.nifi.processor.Processor @@ -32,6 +32,7 @@ org.apache.nifi.processors.tests.system.FakeProcessor org.apache.nifi.processors.tests.system.FakeDynamicPropertiesProcessor org.apache.nifi.processors.tests.system.GenerateAndCountCallbacks org.apache.nifi.processors.tests.system.GenerateFlowFile +org.apache.nifi.processors.tests.system.GatedPassThrough org.apache.nifi.processors.tests.system.GenerateTruncatableFlowFiles org.apache.nifi.processors.tests.system.HoldInput org.apache.nifi.processors.tests.system.IngestFile diff --git a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java index faef5b9a5245..e40cc9c52fe7 100644 --- a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java +++ b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/NiFiClientUtil.java @@ -2771,20 +2771,7 @@ public VersionedFlowUpdateRequestEntity changeFlowVersion(final String processGr logger.info("Submitting Change Flow Version request to change Group with ID {} to Version {}", processGroupId, version); try { - final ProcessGroupEntity groupEntity = nifiClient.getProcessGroupClient().getProcessGroup(processGroupId); - final ProcessGroupDTO groupDto = groupEntity.getComponent(); - final VersionControlInformationDTO vciDto = groupDto.getVersionControlInformation(); - if (vciDto == null) { - throw new IllegalArgumentException("Process Group with ID " + processGroupId + " is not under Version Control"); - } - - vciDto.setVersion(version); - - final VersionControlInformationEntity requestEntity = new VersionControlInformationEntity(); - requestEntity.setProcessGroupRevision(groupEntity.getRevision()); - requestEntity.setVersionControlInformation(vciDto); - - final VersionedFlowUpdateRequestEntity result = nifiClient.getVersionsClient().updateVersionControlInfo(processGroupId, requestEntity); + final VersionedFlowUpdateRequestEntity result = initiateFlowVersionChange(processGroupId, version); return waitForVersionFlowUpdateComplete(result.getRequest().getRequestId(), throwOnFailure); } catch (final Exception e) { logger.error("Failed to change flow version for Process Group {} to version {}", processGroupId, version); @@ -2792,6 +2779,23 @@ public VersionedFlowUpdateRequestEntity changeFlowVersion(final String processGr } } + public VersionedFlowUpdateRequestEntity initiateFlowVersionChange(final String processGroupId, final String version) + throws NiFiClientException, IOException { + final ProcessGroupEntity groupEntity = nifiClient.getProcessGroupClient().getProcessGroup(processGroupId); + final ProcessGroupDTO groupDto = groupEntity.getComponent(); + final VersionControlInformationDTO vciDto = groupDto.getVersionControlInformation(); + if (vciDto == null) { + throw new IllegalArgumentException("Process Group with ID " + processGroupId + " is not under Version Control"); + } + + vciDto.setVersion(version); + + final VersionControlInformationEntity requestEntity = new VersionControlInformationEntity(); + requestEntity.setProcessGroupRevision(groupEntity.getRevision()); + requestEntity.setVersionControlInformation(vciDto); + return nifiClient.getVersionsClient().updateVersionControlInfo(processGroupId, requestEntity); + } + public VersionedFlowUpdateRequestEntity waitForVersionFlowUpdateComplete(final String updateRequestId, final boolean throwOnFailure) throws NiFiClientException, IOException, InterruptedException { while (true) { final VersionedFlowUpdateRequestEntity result = nifiClient.getVersionsClient().getUpdateRequest(updateRequestId); diff --git a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/registry/RegistryClientIT.java b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/registry/RegistryClientIT.java index 0d37673225b4..17a797b37ed6 100644 --- a/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/registry/RegistryClientIT.java +++ b/nifi-system-tests/nifi-system-test-suite/src/test/java/org/apache/nifi/tests/system/registry/RegistryClientIT.java @@ -28,11 +28,14 @@ import org.apache.nifi.web.api.dto.VersionControlInformationDTO; import org.apache.nifi.web.api.dto.flow.FlowDTO; import org.apache.nifi.web.api.dto.flow.ProcessGroupFlowDTO; +import org.apache.nifi.web.api.dto.status.ProcessGroupStatusSnapshotDTO; import org.apache.nifi.web.api.entity.ConnectionEntity; import org.apache.nifi.web.api.entity.ControllerServiceEntity; import org.apache.nifi.web.api.entity.FlowRegistryClientEntity; import org.apache.nifi.web.api.entity.PortEntity; +import org.apache.nifi.web.api.entity.PortStatusSnapshotEntity; import org.apache.nifi.web.api.entity.ProcessGroupEntity; +import org.apache.nifi.web.api.entity.ProcessGroupStatusSnapshotEntity; import org.apache.nifi.web.api.entity.ProcessorEntity; import org.apache.nifi.web.api.entity.SnippetEntity; import org.apache.nifi.web.api.entity.VersionControlInformationEntity; @@ -41,6 +44,8 @@ import java.io.File; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; import java.util.Collections; import java.util.HashMap; import java.util.List; @@ -58,6 +63,153 @@ public class RegistryClientIT extends NiFiSystemIT { public static final String FIRST_FLOW_ID = "first-flow"; + @Test + public void testChangeVersionDrainsRemovedConnectionBeforeUpdate() throws Exception { + final Path gateFile = Path.of(System.getProperty("java.io.tmpdir"), "nifi-removed-connection-drain-" + System.nanoTime()); + final RemovedConnectionFixture fixture = createRemovedConnectionFixture(gateFile, false, true); + final NiFiClientUtil util = getClientUtil(); + + final VersionedFlowUpdateRequestEntity initiated = util.initiateFlowVersionChange(fixture.groupId(), "2"); + final String requestId = initiated.getRequest().getRequestId(); + waitFor(() -> "Draining Removed Connections".equals(getNifiClient().getVersionsClient().getUpdateRequest(requestId).getRequest().getState())); + if (getNumberOfNodes() > 1) { + final List queuedByNode = getNifiClient().getFlowClient().getConnectionStatus(fixture.connectionId(), true) + .getConnectionStatus().getNodeSnapshots().stream() + .map(node -> node.getStatusSnapshot().getFlowFilesQueued()) + .sorted() + .toList(); + assertEquals(List.of(0, 1), queuedByNode); + } + assertEquals("1", getNifiClient().getProcessGroupClient().getProcessGroup(fixture.groupId()) + .getComponent().getVersionControlInformation().getVersion()); + + Files.createFile(gateFile); + final VersionedFlowUpdateRequestEntity completed = util.waitForVersionFlowUpdateComplete(requestId, true); + assertTrue(completed.getRequest().isComplete()); + assertNull(completed.getRequest().getFailureReason()); + assertEquals("2", completed.getRequest().getVersionControlInformation().getVersion()); + util.waitForRunningProcessor(fixture.sourceId()); + util.waitForRunningProcessor(fixture.destinationId()); + assertTrue(getConnections(fixture.groupId()).stream().noneMatch(connection -> fixture.connectionId().equals(connection.getId()))); + + Files.deleteIfExists(gateFile); + } + + @Test + public void testCancelledRemovedConnectionDrainRestoresOriginalFlow() throws Exception { + final Path gateFile = Path.of(System.getProperty("java.io.tmpdir"), "nifi-removed-connection-cancel-" + System.nanoTime()); + final RemovedConnectionFixture fixture = createRemovedConnectionFixture(gateFile, false, true); + final NiFiClientUtil util = getClientUtil(); + + final VersionedFlowUpdateRequestEntity initiated = util.initiateFlowVersionChange(fixture.groupId(), "2"); + final String requestId = initiated.getRequest().getRequestId(); + waitFor(() -> "Draining Removed Connections".equals(getNifiClient().getVersionsClient().getUpdateRequest(requestId).getRequest().getState())); + + final VersionedFlowUpdateRequestEntity cancelled = getNifiClient().getVersionsClient().deleteUpdateRequest(requestId); + assertEquals("Request cancelled by user", cancelled.getRequest().getFailureReason()); + assertOriginalFlowRestored(fixture); + assertTrue(getConnectionQueueSize(fixture.connectionId()) >= 1); + } + + @Test + public void testRemovedConnectionDrainTimeoutRestoresOriginalFlow() throws Exception { + final Path gateFile = Path.of(System.getProperty("java.io.tmpdir"), "nifi-removed-connection-timeout-" + System.nanoTime()); + final RemovedConnectionFixture fixture = createRemovedConnectionFixture(gateFile, false, true); + final NiFiClientUtil util = getClientUtil(); + + final VersionedFlowUpdateRequestEntity initiated = util.initiateFlowVersionChange(fixture.groupId(), "2"); + final VersionedFlowUpdateRequestEntity completed = util.waitForVersionFlowUpdateComplete(initiated.getRequest().getRequestId(), false); + + assertTrue(completed.getRequest().getFailureReason().contains("Removed connection drain timed out")); + assertOriginalFlowRestored(fixture); + assertTrue(getConnectionQueueSize(fixture.connectionId()) >= 1); + } + + @Test + public void testUnsupportedNonEmptyRemovalFailsBeforeRuntimeStateChanges() throws Exception { + final Path gateFile = Path.of(System.getProperty("java.io.tmpdir"), "nifi-unsupported-removal-" + System.nanoTime()); + final RemovedConnectionFixture fixture = createRemovedConnectionFixture(gateFile, true, true); + + final VersionedFlowUpdateRequestEntity completed = getClientUtil().changeFlowVersion(fixture.groupId(), "2", false); + + assertTrue(completed.getRequest().getFailureReason().contains("DESTINATION_COMPONENT_REMOVED")); + assertOriginalFlowRestored(fixture); + assertEquals(1, getConnectionQueueSize(fixture.connectionId())); + } + + @Test + public void testUnsupportedEmptyRemovalRetainsSuccessfulUpdateBehavior() throws Exception { + final Path gateFile = Path.of(System.getProperty("java.io.tmpdir"), "nifi-empty-unsupported-removal-" + System.nanoTime()); + final RemovedConnectionFixture fixture = createRemovedConnectionFixture(gateFile, true, false); + + final VersionedFlowUpdateRequestEntity completed = getClientUtil().changeFlowVersion(fixture.groupId(), "2", true); + + assertNull(completed.getRequest().getFailureReason()); + assertEquals("2", completed.getRequest().getVersionControlInformation().getVersion()); + assertTrue(getConnections(fixture.groupId()).stream().noneMatch(connection -> fixture.connectionId().equals(connection.getId()))); + assertTrue(getNifiClient().getFlowClient().getProcessGroup(fixture.groupId()).getProcessGroupFlow().getFlow().getProcessors().stream() + .noneMatch(processor -> fixture.destinationId().equals(processor.getId()))); + } + + @Test + public void testRemovedConnectionToInputPortStopsPortBeforeRemoval() throws Exception { + final FlowRegistryClientEntity clientEntity = registerClient(); + final NiFiClientUtil util = getClientUtil(); + final ProcessGroupEntity group = util.createProcessGroup("Removed Port Connection", "root"); + final ProcessGroupEntity child = util.createProcessGroup("Port Child", group.getId()); + final ProcessorEntity generate = util.createProcessor("GenerateFlowFile", group.getId()); + util.updateProcessorProperties(generate, Map.of("Max FlowFiles", "1", "State Scope", "CLUSTER")); + final PortEntity inputPort = util.createInputPort("Drain Destination", child.getId()); + final ProcessorEntity terminate = util.createProcessor("TerminateFlowFile", child.getId()); + final ConnectionEntity removedConnection = util.createConnection(generate, inputPort, "success"); + final ConnectionEntity portOutput = util.createConnection(inputPort, terminate); + util.updateConnectionBackpressure(portOutput, 1, 1_000_000); + + final VersionControlInformationEntity v1 = util.startVersionControl(group, clientEntity, TEST_FLOWS_BUCKET, "removed-port-" + System.nanoTime()); + getNifiClient().getConnectionClient().deleteConnection(removedConnection); + util.setAutoTerminatedRelationships(generate, "success"); + inputPort.getComponent().setComments("Updated in version 2"); + getNifiClient().getInputPortClient().updateInputPort(inputPort); + util.saveFlowVersion(group, clientEntity, v1); + util.changeFlowVersion(group.getId(), "1"); + + final ConnectionEntity restoredConnection = getConnections(group.getId()).stream() + .filter(connection -> "Drain Destination".equals(connection.getComponent().getDestination().getName())) + .filter(connection -> generate.getId().equals(connection.getComponent().getSource().getId())) + .findFirst() + .orElseThrow(); + final PortEntity restoredPort = getNifiClient().getInputPortClient().getInputPort(restoredConnection.getComponent().getDestination().getId()); + getNifiClient().getInputPortClient().startInputPort(restoredPort); + util.waitForValidProcessor(restoredConnection.getComponent().getSource().getId()); + util.startProcessor(getNifiClient().getProcessorClient().getProcessor(restoredConnection.getComponent().getSource().getId())); + final ConnectionEntity restoredPortOutput = getConnections(restoredPort.getComponent().getParentGroupId()).stream() + .filter(connection -> restoredPort.getId().equals(connection.getComponent().getSource().getId())) + .findFirst() + .orElseThrow(); + waitForQueueCount(restoredPortOutput.getId(), getNumberOfNodes()); + final ProcessorEntity restoredSource = getNifiClient().getProcessorClient().getProcessor(restoredConnection.getComponent().getSource().getId()); + util.stopProcessor(restoredSource); + util.startProcessor(restoredSource); + waitForQueueCount(restoredConnection.getId(), getNumberOfNodes()); + + final VersionedFlowUpdateRequestEntity initiated = util.initiateFlowVersionChange(group.getId(), "2"); + final String requestId = initiated.getRequest().getRequestId(); + waitFor(() -> "Draining Removed Connections".equals(getNifiClient().getVersionsClient().getUpdateRequest(requestId).getRequest().getState())); + util.startProcessor(getNifiClient().getProcessorClient().getProcessor(terminate.getId())); + final VersionedFlowUpdateRequestEntity completed = util.waitForVersionFlowUpdateComplete(requestId, false); + + assertNotNull(completed.getRequest().getFailureReason()); + assertTrue(completed.getRequest().getFailureReason().contains("Input Port")); + assertTrue(completed.getRequest().getFailureReason().contains("Port has no incoming connections")); + assertEquals("2", completed.getRequest().getVersionControlInformation().getVersion()); + assertTrue(getConnections(group.getId()).stream().noneMatch(connection -> restoredConnection.getId().equals(connection.getId()))); + final String updatedPortState = getNifiClient().getInputPortClient().getInputPort(restoredPort.getId()).getComponent().getState(); + assertEquals("STOPPED", updatedPortState); + final ProcessGroupStatusSnapshotDTO groupStatus = getNifiClient().getFlowClient().getProcessGroupStatus(group.getId(), true) + .getProcessGroupStatus().getAggregateSnapshot(); + assertEquals(0, getInputPortActiveThreadCount(groupStatus, restoredPort.getId())); + } + /** * Test a scenario where we have Parent Process Group with a child process group. The child group is under Version Control. * Then the parent is placed under Version Control. Then modify a Processor in child. Register snapshot for child, then for parent. @@ -739,4 +891,82 @@ public void testStopVersionControlThenSetVersionControlInfo() throws NiFiClientE assertEquals(vci.getVersionControlInformation().getVersion(), groupAfterSetVersionInfo.getComponent().getVersionControlInformation().getVersion()); assertEquals("UP_TO_DATE", groupAfterSetVersionInfo.getComponent().getVersionControlInformation().getState()); } + + private RemovedConnectionFixture createRemovedConnectionFixture(final Path gateFile, final boolean removeDestination, + final boolean queueFlowFile) throws Exception { + Files.deleteIfExists(gateFile); + final FlowRegistryClientEntity clientEntity = registerClient(); + final NiFiClientUtil util = getClientUtil(); + final ProcessGroupEntity group = util.createProcessGroup("Removed Connection Drain", "root"); + final ProcessorEntity generate = util.createProcessor("GenerateFlowFile", group.getId()); + util.updateProcessorProperties(generate, Map.of("Max FlowFiles", "1", "State Scope", "CLUSTER")); + if (getNumberOfNodes() > 1) { + util.updateProcessorExecutionNode(generate, ExecutionNode.PRIMARY); + } + + final ProcessorEntity gated = util.createProcessor("GatedPassThrough", group.getId()); + util.updateProcessorProperties(gated, Map.of("Gate File", gateFile.toString())); + util.setAutoTerminatedRelationships(gated, "success"); + final ProcessorEntity terminate = util.createProcessor("TerminateFlowFile", group.getId()); + final ConnectionEntity removedConnection = util.createConnection(generate, gated, "success"); + + final VersionControlInformationEntity v1 = util.startVersionControl(group, clientEntity, TEST_FLOWS_BUCKET, "removed-connection-" + System.nanoTime()); + getNifiClient().getConnectionClient().deleteConnection(removedConnection); + util.createConnection(generate, terminate, "success"); + if (removeDestination) { + getNifiClient().getProcessorClient().deleteProcessor(gated); + } + util.saveFlowVersion(group, clientEntity, v1); + util.changeFlowVersion(group.getId(), "1"); + final ConnectionEntity restoredConnection = getConnections(group.getId()).stream() + .filter(connection -> "GenerateFlowFile".equals(connection.getComponent().getSource().getName())) + .filter(connection -> "GatedPassThrough".equals(connection.getComponent().getDestination().getName())) + .findFirst() + .orElseThrow(); + final String restoredConnectionId = restoredConnection.getId(); + final String restoredSourceId = restoredConnection.getComponent().getSource().getId(); + final String restoredDestinationId = restoredConnection.getComponent().getDestination().getId(); + + if (queueFlowFile) { + util.waitForValidProcessor(restoredDestinationId); + util.startProcessor(getNifiClient().getProcessorClient().getProcessor(restoredDestinationId)); + util.waitForValidProcessor(restoredSourceId); + util.startProcessor(getNifiClient().getProcessorClient().getProcessor(restoredSourceId)); + waitForQueueCount(restoredConnectionId, 1); + } + + return new RemovedConnectionFixture(group.getId(), restoredSourceId, restoredDestinationId, restoredConnectionId); + } + + private Set getConnections(final String groupId) throws NiFiClientException, IOException { + return getNifiClient().getFlowClient().getProcessGroup(groupId).getProcessGroupFlow().getFlow().getConnections(); + } + + private int getInputPortActiveThreadCount(final ProcessGroupStatusSnapshotDTO groupStatus, final String portId) { + for (final PortStatusSnapshotEntity portStatus : groupStatus.getInputPortStatusSnapshots()) { + if (portId.equals(portStatus.getId())) { + return portStatus.getPortStatusSnapshot().getActiveThreadCount(); + } + } + + for (final ProcessGroupStatusSnapshotEntity childStatus : groupStatus.getProcessGroupStatusSnapshots()) { + final int activeThreadCount = getInputPortActiveThreadCount(childStatus.getProcessGroupStatusSnapshot(), portId); + if (activeThreadCount >= 0) { + return activeThreadCount; + } + } + + return -1; + } + + private void assertOriginalFlowRestored(final RemovedConnectionFixture fixture) throws Exception { + assertEquals("1", getNifiClient().getProcessGroupClient().getProcessGroup(fixture.groupId()) + .getComponent().getVersionControlInformation().getVersion()); + assertTrue(getConnections(fixture.groupId()).stream().anyMatch(connection -> fixture.connectionId().equals(connection.getId()))); + getClientUtil().waitForRunningProcessor(fixture.sourceId()); + getClientUtil().waitForRunningProcessor(fixture.destinationId()); + } + + private record RemovedConnectionFixture(String groupId, String sourceId, String destinationId, String connectionId) { + } }