diff --git a/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ComponentMetricReporter.java b/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ComponentMetricReporter.java index aeeb61649ce1..c6ce22109c81 100644 --- a/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ComponentMetricReporter.java +++ b/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ComponentMetricReporter.java @@ -59,4 +59,22 @@ default void close() { */ default void recordProcessSessionEvent(ProcessSessionEvent processSessionEvent) { } + + /** + * Record current Connection Status for operations completed during a Process Session + * + * @param connectionStatusEvent Connection Status Event containing component context and captured status + */ + default void recordConnectionStatusEvent(ConnectionStatusEvent connectionStatusEvent) { + } + + /** + * Status indicator for recording Connection Status Event defaults to disabled and requires overriding along with + * recordConnectionStatusEvent() to handle Connection Status Events + * + * @return Enabled or disabled status for recording Connection Status Events + */ + default boolean isRecordConnectionStatusEventEnabled() { + return false; + } } diff --git a/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ConnectionStatusEvent.java b/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ConnectionStatusEvent.java new file mode 100644 index 000000000000..d372489d59f2 --- /dev/null +++ b/nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/ConnectionStatusEvent.java @@ -0,0 +1,74 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.controller.metrics; + +import org.apache.nifi.controller.status.FlowFileAvailability; +import org.apache.nifi.controller.status.LoadBalanceStatus; + +/** + * Event abstraction for Connection Status metrics collected during ProcessSession operations + */ +public interface ConnectionStatusEvent { + /** + * Get Component Metric Context describing the Component associated with the recorded ProcessSession operations + * + * @return Component Metric Context + */ + ComponentMetricContext getComponentMetricContext(); + + /** + * Get configured Back Pressure Bytes Threshold + * + * @return Back Pressure Bytes Threshold + */ + long getBackPressureBytesThreshold(); + + /** + * Get configured Back Pressure Object Threshold + * + * @return Back Pressure Object Threshold + */ + long getBackPressureObjectThreshold(); + + /** + * Get bytes from queued FlowFiles for the Connection + * + * @return Queued Bytes + */ + long getQueuedBytes(); + + /** + * Get count of FlowFiles queued for the Connection + * + * @return Queued FlowFiles + */ + int getQueuedCount(); + + /** + * Get Load Balance Status for the Connection + * + * @return Load Balance Status + */ + LoadBalanceStatus getLoadBalanceStatus(); + + /** + * Get FlowFile Availability for the Connection + * + * @return FlowFile Availability + */ + FlowFileAvailability getFlowFileAvailability(); +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/AbstractRepositoryContext.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/AbstractRepositoryContext.java index c1e911828ae0..f92390352ae0 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/AbstractRepositoryContext.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/AbstractRepositoryContext.java @@ -24,6 +24,7 @@ import org.apache.nifi.controller.ProcessorNode; import org.apache.nifi.controller.metrics.ComponentMetricContext; import org.apache.nifi.controller.metrics.ComponentMetricReporter; +import org.apache.nifi.controller.metrics.ConnectionStatusEvent; import org.apache.nifi.controller.metrics.CounterRecord; import org.apache.nifi.controller.metrics.GaugeRecord; import org.apache.nifi.controller.metrics.ProcessSessionEvent; @@ -182,6 +183,16 @@ public void recordProcessSessionEvent(final ProcessSessionEvent event) { componentMetricReporter.recordProcessSessionEvent(event); } + @Override + public void recordConnectionStatusEvent(final ConnectionStatusEvent event) { + componentMetricReporter.recordConnectionStatusEvent(event); + } + + @Override + public boolean isRecordConnectionStatusEventEnabled() { + return componentMetricReporter.isRecordConnectionStatusEventEnabled(); + } + @Override public ContentRepository getContentRepository() { return contentRepo; diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/RepositoryContext.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/RepositoryContext.java index 212ab3447758..a26a39bfff9b 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/RepositoryContext.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/RepositoryContext.java @@ -21,6 +21,7 @@ import org.apache.nifi.connectable.Connectable; import org.apache.nifi.connectable.Connection; import org.apache.nifi.controller.metrics.ComponentMetricContext; +import org.apache.nifi.controller.metrics.ConnectionStatusEvent; import org.apache.nifi.controller.metrics.GaugeRecord; import org.apache.nifi.controller.metrics.ProcessSessionEvent; import org.apache.nifi.controller.repository.claim.ContentClaimWriteCache; @@ -70,6 +71,10 @@ public interface RepositoryContext { void recordProcessSessionEvent(ProcessSessionEvent event); + void recordConnectionStatusEvent(ConnectionStatusEvent event); + + boolean isRecordConnectionStatusEventEnabled(); + ProvenanceEventBuilder createProvenanceEventBuilder(); StateManager getStateManager(); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java index db0a419db3b9..bd60a85b5937 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java @@ -25,9 +25,11 @@ import org.apache.nifi.controller.ProcessorNode; import org.apache.nifi.controller.lifecycle.TaskTermination; import org.apache.nifi.controller.metrics.ComponentMetricContext; +import org.apache.nifi.controller.metrics.ConnectionStatusEvent; import org.apache.nifi.controller.metrics.GaugeRecord; import org.apache.nifi.controller.metrics.ProcessSessionEvent; import org.apache.nifi.controller.queue.FlowFileQueue; +import org.apache.nifi.controller.queue.LoadBalanceStrategy; import org.apache.nifi.controller.queue.PollStrategy; import org.apache.nifi.controller.queue.QueueSize; import org.apache.nifi.controller.repository.claim.ContentClaim; @@ -41,12 +43,16 @@ import org.apache.nifi.controller.repository.io.LimitedInputStream; import org.apache.nifi.controller.repository.io.TaskTerminationInputStream; import org.apache.nifi.controller.repository.io.TaskTerminationOutputStream; +import org.apache.nifi.controller.repository.metrics.ConnectionStatusEventBuilder; import org.apache.nifi.controller.repository.metrics.PerformanceTracker; import org.apache.nifi.controller.repository.metrics.PerformanceTrackingInputStream; import org.apache.nifi.controller.repository.metrics.ProcessSessionEventBuilder; import org.apache.nifi.controller.state.StandardStateMap; +import org.apache.nifi.controller.status.FlowFileAvailability; +import org.apache.nifi.controller.status.LoadBalanceStatus; import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.flowfile.attributes.CoreAttributes; +import org.apache.nifi.processor.DataUnit; import org.apache.nifi.processor.FlowFileFilter; import org.apache.nifi.processor.ProcessSession; import org.apache.nifi.processor.Relationship; @@ -145,6 +151,7 @@ public class StandardProcessSession implements ProcessSession, ProvenanceEventEn private final Map records = new ConcurrentHashMap<>(); private final Map connectionCounts = new ConcurrentHashMap<>(); + private final Map processedConnections = new ConcurrentHashMap<>(); private final Map connectionMetricContexts = new ConcurrentHashMap<>(); private final Map> unacknowledgedFlowFiles = new ConcurrentHashMap<>(); private final Map appendableStreams = new ConcurrentHashMap<>(); @@ -807,11 +814,52 @@ private void updateEventRepository(final Checkpoint checkpoint) { context.getFlowFileEventRepository().updateRepository(connectionSessionEvent); context.recordProcessSessionEvent(connectionSessionEvent); } + + recordConnectionStatusEvents(checkpoint); } catch (final IOException ioe) { LOG.error("FlowFile Event Repository failed to update", ioe); } } + private void recordConnectionStatusEvents(final Checkpoint checkpoint) { + // Check enabled status to avoid building objects and calling methods when not used + if (context.isRecordConnectionStatusEventEnabled()) { + for (final Connection connection : checkpoint.processedConnections.values()) { + final ComponentMetricContext connectionMetricContext = checkpoint.connectionMetricContexts.get(connection.getIdentifier()); + final FlowFileQueue flowFileQueue = connection.getFlowFileQueue(); + final QueueSize queueSize = flowFileQueue.size(); + final long backPressureBytesThreshold = DataUnit.parseDataSize(flowFileQueue.getBackPressureDataSizeThreshold(), DataUnit.B).longValue(); + final LoadBalanceStatus loadBalanceStatus = getLoadBalanceStatus(flowFileQueue); + final FlowFileAvailability flowFileAvailability = flowFileQueue.getFlowFileAvailability(); + + final ConnectionStatusEvent connectionStatusEvent = ConnectionStatusEventBuilder.forComponent(connectionMetricContext) + .backPressureBytesThreshold(backPressureBytesThreshold) + .backPressureObjectThreshold(flowFileQueue.getBackPressureObjectThreshold()) + .queuedBytes(queueSize.getByteCount()) + .queuedCount(queueSize.getObjectCount()) + .loadBalanceStatus(loadBalanceStatus) + .flowFileAvailability(flowFileAvailability) + .build(); + context.recordConnectionStatusEvent(connectionStatusEvent); + } + } + } + + private LoadBalanceStatus getLoadBalanceStatus(final FlowFileQueue flowFileQueue) { + final LoadBalanceStatus loadBalanceStatus; + + final LoadBalanceStrategy loadBalanceStrategy = flowFileQueue.getLoadBalanceStrategy(); + if (loadBalanceStrategy == LoadBalanceStrategy.DO_NOT_LOAD_BALANCE) { + loadBalanceStatus = LoadBalanceStatus.LOAD_BALANCE_NOT_CONFIGURED; + } else if (flowFileQueue.isActivelyLoadBalancing()) { + loadBalanceStatus = LoadBalanceStatus.LOAD_BALANCE_ACTIVE; + } else { + loadBalanceStatus = LoadBalanceStatus.LOAD_BALANCE_INACTIVE; + } + + return loadBalanceStatus; + } + private Map combineCounters(final Map first, final Map second) { final boolean firstEmpty = first == null || first.isEmpty(); final boolean secondEmpty = second == null || second.isEmpty(); @@ -1453,6 +1501,7 @@ private void resetState() { bytesRead = 0L; bytesWritten = 0L; connectionCounts.clear(); + processedConnections.clear(); connectionMetricContexts.clear(); createdFlowFiles.clear(); createdFlowFilesWithoutLineage.clear(); @@ -1652,8 +1701,18 @@ private synchronized void migrate(final StandardProcessSession newOwner, Collect final FlowFileQueue inputQueue = repoRecord.getOriginalQueue(); if (inputQueue != null) { final String connectionId = inputQueue.getIdentifier(); - incrementConnectionOutputCounts(connectionId, -1, -repoRecord.getOriginal().getSize()); - newOwner.incrementConnectionOutputCounts(connectionId, 1, repoRecord.getOriginal().getSize()); + final long originalSize = repoRecord.getOriginal().getSize(); + + // A FlowFile that has an original queue was dequeued within this session, so the source Connection is tracked and can be + // reused to record Connection Status for the destination session. Adjust counts by identifier only when it is not tracked. + final Connection connection = processedConnections.get(connectionId); + if (connection == null) { + incrementConnectionOutputCounts(connectionId, -1, -originalSize); + newOwner.incrementConnectionOutputCounts(connectionId, 1, originalSize); + } else { + incrementConnectionOutputCounts(connection, -1, -originalSize); + newOwner.incrementConnectionOutputCounts(connection, 1, originalSize); + } unacknowledgedFlowFiles.get(inputQueue).remove(flowFile); newOwner.unacknowledgedFlowFiles.computeIfAbsent(inputQueue, queue -> new HashSet<>()).add(flowFileRecord); @@ -1838,15 +1897,14 @@ private void formatNanos(final long nanos, final StringBuilder sb) { } private void incrementConnectionInputCounts(final Connection connection, final RepositoryRecord record) { - incrementConnectionInputCounts(connection, 1, record.getCurrent().getSize()); - } - - private void incrementConnectionInputCounts(final Connection connection, final int flowFileCount, final long bytes) { final String connectionId = connection.getIdentifier(); cacheConnectionMetricContext(connection); final ProcessSessionEventBuilder connectionEvent = connectionCounts.computeIfAbsent( connectionId, id -> ProcessSessionEventBuilder.forComponent(getConnectionMetricContext(connectionId))); - connectionEvent.addFlowFilesIn(flowFileCount).addContentSizeIn(bytes); + final long bytes = record.getCurrent().getSize(); + connectionEvent.addFlowFilesIn(1).addContentSizeIn(bytes); + + processedConnections.put(connectionId, connection); } private void incrementConnectionOutputCounts(final Connection connection, final FlowFileRecord record) { @@ -1859,6 +1917,8 @@ private void incrementConnectionOutputCounts(final Connection connection, final final ProcessSessionEventBuilder connectionEvent = connectionCounts.computeIfAbsent( connectionId, id -> ProcessSessionEventBuilder.forComponent(getConnectionMetricContext(connectionId))); connectionEvent.addFlowFilesOut(flowFileCount).addContentSizeOut(bytes); + + processedConnections.put(connectionId, connection); } private void incrementConnectionOutputCounts(final String connectionId, final int flowFileCount, final long bytes) { @@ -3994,6 +4054,7 @@ protected static class Checkpoint { private Map records; private Map connectionCounts; + private Map processedConnections; private Map connectionMetricContexts; private Map countersOnCommit; @@ -4034,6 +4095,7 @@ private void initializeForCopy() { records = new ConcurrentHashMap<>(); connectionCounts = new ConcurrentHashMap<>(); + processedConnections = new ConcurrentHashMap<>(); connectionMetricContexts = new ConcurrentHashMap<>(); countersOnCommit = new HashMap<>(); @@ -4070,6 +4132,7 @@ private void directCheckpoint(final StandardProcessSession session, final List

destination.merge(toMerge.build())); + mergeMapsWithMutableValue(this.connectionCounts, session.connectionCounts, + (destination, toMerge) -> destination.merge(toMerge.build()) + ); + mergeMaps(this.processedConnections, session.processedConnections, (existing, incoming) -> existing); mergeMaps(this.connectionMetricContexts, session.connectionMetricContexts, (existing, incoming) -> existing); mergeMaps(this.countersOnCommit, session.countersOnCommit, Long::sum); mergeMaps(this.immediateCounters, session.immediateCounters, Long::sum); diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/metrics/ConnectionStatusEventBuilder.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/metrics/ConnectionStatusEventBuilder.java new file mode 100644 index 000000000000..8c48a5e8e1a6 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/metrics/ConnectionStatusEventBuilder.java @@ -0,0 +1,89 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.controller.repository.metrics; + +import org.apache.nifi.controller.metrics.ComponentMetricContext; +import org.apache.nifi.controller.metrics.ConnectionStatusEvent; +import org.apache.nifi.controller.status.FlowFileAvailability; +import org.apache.nifi.controller.status.LoadBalanceStatus; + +import java.util.Objects; + +/** + * Builder for Connection Status Events with Component Metric Context required + */ +public class ConnectionStatusEventBuilder { + + private final ComponentMetricContext componentMetricContext; + + private long backPressureBytesThreshold; + private long backPressureObjectThreshold; + private long queuedBytes; + private int queuedCount; + private LoadBalanceStatus loadBalanceStatus = LoadBalanceStatus.LOAD_BALANCE_NOT_CONFIGURED; + private FlowFileAvailability flowFileAvailability = FlowFileAvailability.ACTIVE_QUEUE_EMPTY; + + private ConnectionStatusEventBuilder(final ComponentMetricContext componentMetricContext) { + this.componentMetricContext = Objects.requireNonNull(componentMetricContext, "Component Metric Context required"); + } + + public static ConnectionStatusEventBuilder forComponent(final ComponentMetricContext componentMetricContext) { + return new ConnectionStatusEventBuilder(componentMetricContext); + } + + public ConnectionStatusEventBuilder backPressureBytesThreshold(final long backPressureBytesThreshold) { + this.backPressureBytesThreshold = backPressureBytesThreshold; + return this; + } + + public ConnectionStatusEventBuilder backPressureObjectThreshold(final long backPressureObjectThreshold) { + this.backPressureObjectThreshold = backPressureObjectThreshold; + return this; + } + + public ConnectionStatusEventBuilder queuedBytes(final long queuedBytes) { + this.queuedBytes = queuedBytes; + return this; + } + + public ConnectionStatusEventBuilder queuedCount(final int queuedCount) { + this.queuedCount = queuedCount; + return this; + } + + public ConnectionStatusEventBuilder loadBalanceStatus(final LoadBalanceStatus loadBalanceStatus) { + this.loadBalanceStatus = Objects.requireNonNull(loadBalanceStatus, "Load Balance Status required"); + return this; + } + + public ConnectionStatusEventBuilder flowFileAvailability(final FlowFileAvailability flowFileAvailability) { + this.flowFileAvailability = Objects.requireNonNull(flowFileAvailability, "FlowFile Availability required"); + return this; + } + + public ConnectionStatusEvent build() { + return new StandardConnectionStatusEvent( + componentMetricContext, + backPressureBytesThreshold, + backPressureObjectThreshold, + queuedBytes, + queuedCount, + loadBalanceStatus, + flowFileAvailability + ); + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/metrics/StandardConnectionStatusEvent.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/metrics/StandardConnectionStatusEvent.java new file mode 100644 index 000000000000..4f4b32a20a85 --- /dev/null +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/metrics/StandardConnectionStatusEvent.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.nifi.controller.repository.metrics; + +import org.apache.nifi.controller.metrics.ComponentMetricContext; +import org.apache.nifi.controller.metrics.ConnectionStatusEvent; +import org.apache.nifi.controller.status.FlowFileAvailability; +import org.apache.nifi.controller.status.LoadBalanceStatus; + +/** + * Standard record representation of Connection Status Event with package-private visibility for Builder + */ +record StandardConnectionStatusEvent( + ComponentMetricContext componentMetricContext, + long backPressureBytesThreshold, + long backPressureObjectThreshold, + long queuedBytes, + int queuedCount, + LoadBalanceStatus loadBalanceStatus, + FlowFileAvailability flowFileAvailability +) implements ConnectionStatusEvent { + + @Override + public ComponentMetricContext getComponentMetricContext() { + return componentMetricContext; + } + + @Override + public long getBackPressureBytesThreshold() { + return backPressureBytesThreshold; + } + + @Override + public long getBackPressureObjectThreshold() { + return backPressureObjectThreshold; + } + + @Override + public long getQueuedBytes() { + return queuedBytes; + } + + @Override + public int getQueuedCount() { + return queuedCount; + } + + @Override + public LoadBalanceStatus getLoadBalanceStatus() { + return loadBalanceStatus; + } + + @Override + public FlowFileAvailability getFlowFileAvailability() { + return flowFileAvailability; + } +} diff --git a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionTest.java b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionTest.java index 18d468b10e77..79ce3f98a0fb 100644 --- a/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionTest.java +++ b/nifi-framework-bundle/nifi-framework/nifi-framework-components/src/test/java/org/apache/nifi/controller/repository/StandardProcessSessionTest.java @@ -17,15 +17,24 @@ package org.apache.nifi.controller.repository; import org.apache.nifi.connectable.Connectable; +import org.apache.nifi.connectable.Connection; import org.apache.nifi.connectable.FlowFileActivity; import org.apache.nifi.controller.lifecycle.TaskTermination; +import org.apache.nifi.controller.metrics.ComponentMetricContext; +import org.apache.nifi.controller.metrics.ConnectionStatusEvent; import org.apache.nifi.controller.metrics.GaugeRecord; import org.apache.nifi.controller.metrics.ProcessSessionEvent; +import org.apache.nifi.controller.queue.FlowFileQueue; +import org.apache.nifi.controller.queue.LoadBalanceStrategy; +import org.apache.nifi.controller.queue.QueueSize; import org.apache.nifi.controller.repository.claim.ContentClaim; import org.apache.nifi.controller.repository.claim.ContentClaimWriteCache; import org.apache.nifi.controller.repository.metrics.PerformanceTracker; +import org.apache.nifi.controller.status.FlowFileAvailability; +import org.apache.nifi.controller.status.LoadBalanceStatus; import org.apache.nifi.flowfile.FlowFile; import org.apache.nifi.processor.ProcessSession; +import org.apache.nifi.processor.Relationship; import org.apache.nifi.processor.metrics.CommitTiming; import org.apache.nifi.provenance.InternalProvenanceReporter; import org.apache.nifi.provenance.ProvenanceRepository; @@ -44,16 +53,20 @@ import java.io.IOException; import java.nio.file.Files; import java.nio.file.Path; +import java.util.List; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNotEquals; import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anySet; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.ArgumentMatchers.isA; import static org.mockito.ArgumentMatchers.isNull; 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; @@ -76,6 +89,11 @@ class StandardProcessSessionTest { private static final double GAUGE_VALUE = 64.5; + private static final String INPUT_CONNECTION_ID = "input-connection-id"; + private static final String OUTPUT_CONNECTION_ID = "output-connection-id"; + private static final String BACK_PRESSURE_DATA_SIZE_THRESHOLD = "1 MB"; + private static final long BACK_PRESSURE_BYTES_THRESHOLD = 1048576; + @Mock RepositoryContext repositoryContext; @@ -118,6 +136,9 @@ class StandardProcessSessionTest { @Captor ArgumentCaptor gaugeRecordCaptor; + @Captor + ArgumentCaptor connectionStatusEventCaptor; + StandardProcessSession session; @BeforeEach @@ -132,6 +153,163 @@ void setSession() { session = new StandardProcessSession(repositoryContext, taskTermination, performanceTracker); } + @Test + void testGetTransferConnectionStatusEventsDisabled() { + setRepositoryContext(); + when(repositoryContext.getContentRepository()).thenReturn(contentRepository); + when(repositoryContext.isRecordConnectionStatusEventEnabled()).thenReturn(false); + + final Connection connection = mock(Connection.class); + when(repositoryContext.getPollableConnections()).thenReturn(List.of(connection)); + final FlowFileRecord flowFileRecord = mock(FlowFileRecord.class); + when(connection.poll(anySet())).thenReturn(flowFileRecord); + final FlowFileQueue flowFileQueue = mock(FlowFileQueue.class); + when(connection.getFlowFileQueue()).thenReturn(flowFileQueue); + when(connection.getIdentifier()).thenReturn(INPUT_CONNECTION_ID); + + final FlowFile flowFile = session.get(); + assertNotNull(flowFile); + session.transfer(flowFile); + session.commit(); + + verify(repositoryContext, never()).recordConnectionStatusEvent(connectionStatusEventCaptor.capture()); + } + + @Test + void testGetTransferConnectionStatusEvents() { + setRepositoryContext(); + when(repositoryContext.getContentRepository()).thenReturn(contentRepository); + when(repositoryContext.isRecordConnectionStatusEventEnabled()).thenReturn(true); + + final Connection connection = mock(Connection.class); + when(repositoryContext.getPollableConnections()).thenReturn(List.of(connection)); + final FlowFileRecord flowFileRecord = mock(FlowFileRecord.class); + when(connection.poll(anySet())).thenReturn(flowFileRecord); + final FlowFileQueue flowFileQueue = mock(FlowFileQueue.class); + when(connection.getFlowFileQueue()).thenReturn(flowFileQueue); + when(connection.getIdentifier()).thenReturn(INPUT_CONNECTION_ID); + + final FlowFile flowFile = session.get(); + assertNotNull(flowFile); + + final Connection outputConnection = mock(Connection.class); + when(outputConnection.getIdentifier()).thenReturn(OUTPUT_CONNECTION_ID); + final FlowFileQueue outputFlowFileQueue = mock(FlowFileQueue.class); + when(outputFlowFileQueue.getBackPressureDataSizeThreshold()).thenReturn(BACK_PRESSURE_DATA_SIZE_THRESHOLD); + when(outputConnection.getFlowFileQueue()).thenReturn(outputFlowFileQueue); + final QueueSize outputQueueSize = mock(QueueSize.class); + when(outputFlowFileQueue.size()).thenReturn(outputQueueSize); + when(outputFlowFileQueue.getFlowFileAvailability()).thenReturn(FlowFileAvailability.FLOWFILE_AVAILABLE); + + final Relationship relationship = new Relationship.Builder().name(Relationship.class.getSimpleName()).build(); + when(repositoryContext.getConnections(eq(relationship))).thenReturn(List.of(outputConnection)); + session.transfer(flowFile, relationship); + + when(flowFileQueue.getBackPressureDataSizeThreshold()).thenReturn(BACK_PRESSURE_DATA_SIZE_THRESHOLD); + final QueueSize queueSize = mock(QueueSize.class); + final int objectCount = Integer.MAX_VALUE; + when(queueSize.getObjectCount()).thenReturn(objectCount); + final long byteCount = Long.MAX_VALUE; + when(queueSize.getByteCount()).thenReturn(byteCount); + when(flowFileQueue.size()).thenReturn(queueSize); + when(flowFileQueue.getLoadBalanceStrategy()).thenReturn(LoadBalanceStrategy.ROUND_ROBIN); + when(flowFileQueue.getFlowFileAvailability()).thenReturn(FlowFileAvailability.FLOWFILE_AVAILABLE); + + session.commit(); + + verify(repositoryContext, times(2)).recordConnectionStatusEvent(connectionStatusEventCaptor.capture()); + final List events = connectionStatusEventCaptor.getAllValues(); + + final ConnectionStatusEvent firstConnectionStatusEvent = events.getFirst(); + final ComponentMetricContext componentMetricContext = firstConnectionStatusEvent.getComponentMetricContext(); + assertEquals(INPUT_CONNECTION_ID, componentMetricContext.id()); + assertEquals(BACK_PRESSURE_BYTES_THRESHOLD, firstConnectionStatusEvent.getBackPressureBytesThreshold()); + assertEquals(objectCount, firstConnectionStatusEvent.getQueuedCount()); + assertEquals(byteCount, firstConnectionStatusEvent.getQueuedBytes()); + assertEquals(LoadBalanceStatus.LOAD_BALANCE_INACTIVE, firstConnectionStatusEvent.getLoadBalanceStatus()); + + final ConnectionStatusEvent secondConnectionStatusEvent = events.getLast(); + final ComponentMetricContext secondComponentMetricContext = secondConnectionStatusEvent.getComponentMetricContext(); + assertEquals(OUTPUT_CONNECTION_ID, secondComponentMetricContext.id()); + } + + @Test + void testBatchedCheckpointRetainsConnectionMetricContext() { + setRepositoryContext(); + when(repositoryContext.getContentRepository()).thenReturn(contentRepository); + when(repositoryContext.isRecordConnectionStatusEventEnabled()).thenReturn(true); + + final Connection connection = mock(Connection.class); + when(repositoryContext.getPollableConnections()).thenReturn(List.of(connection)); + final FlowFileRecord flowFileRecord = mock(FlowFileRecord.class); + when(connection.poll(anySet())).thenReturn(flowFileRecord); + final FlowFileQueue flowFileQueue = mock(FlowFileQueue.class); + when(connection.getFlowFileQueue()).thenReturn(flowFileQueue); + when(connection.getIdentifier()).thenReturn(INPUT_CONNECTION_ID); + when(connection.getName()).thenReturn("Connection Name"); + + final FlowFile flowFile = session.get(); + assertNotNull(flowFile); + + when(flowFileQueue.getBackPressureDataSizeThreshold()).thenReturn(BACK_PRESSURE_DATA_SIZE_THRESHOLD); + final QueueSize queueSize = mock(QueueSize.class); + when(flowFileQueue.size()).thenReturn(queueSize); + when(flowFileQueue.getLoadBalanceStrategy()).thenReturn(LoadBalanceStrategy.DO_NOT_LOAD_BALANCE); + when(flowFileQueue.getFlowFileAvailability()).thenReturn(FlowFileAvailability.FLOWFILE_AVAILABLE); + + session.remove(flowFile); + session.checkpoint(); + session.commit(); + + verify(repositoryContext, times(1)).recordConnectionStatusEvent(connectionStatusEventCaptor.capture()); + final ConnectionStatusEvent connectionStatusEvent = connectionStatusEventCaptor.getValue(); + assertEquals("Connection Name", connectionStatusEvent.getComponentMetricContext().name()); + } + + @Test + void testMigrateTracksConnectionStatusEventForNewOwner() { + setRepositoryContext(); + when(repositoryContext.getContentRepository()).thenReturn(contentRepository); + when(repositoryContext.isRecordConnectionStatusEventEnabled()).thenReturn(true); + + final Connection connection = mock(Connection.class); + when(repositoryContext.getPollableConnections()).thenReturn(List.of(connection)); + final FlowFileRecord flowFileRecord = mock(FlowFileRecord.class); + when(connection.poll(anySet())).thenReturn(flowFileRecord); + final FlowFileQueue flowFileQueue = mock(FlowFileQueue.class); + when(connection.getFlowFileQueue()).thenReturn(flowFileQueue); + when(connection.getIdentifier()).thenReturn(INPUT_CONNECTION_ID); + when(flowFileQueue.getIdentifier()).thenReturn(INPUT_CONNECTION_ID); + + final FlowFile flowFile = session.get(); + assertNotNull(flowFile); + + final StandardProcessSession newOwner = new StandardProcessSession(repositoryContext, taskTermination, performanceTracker); + session.migrate(newOwner); + + when(flowFileQueue.getBackPressureDataSizeThreshold()).thenReturn(BACK_PRESSURE_DATA_SIZE_THRESHOLD); + final QueueSize queueSize = mock(QueueSize.class); + final int objectCount = Integer.MAX_VALUE; + when(queueSize.getObjectCount()).thenReturn(objectCount); + final long byteCount = Long.MAX_VALUE; + when(queueSize.getByteCount()).thenReturn(byteCount); + when(flowFileQueue.size()).thenReturn(queueSize); + when(flowFileQueue.getLoadBalanceStrategy()).thenReturn(LoadBalanceStrategy.DO_NOT_LOAD_BALANCE); + when(flowFileQueue.getFlowFileAvailability()).thenReturn(FlowFileAvailability.FLOWFILE_AVAILABLE); + + newOwner.remove(flowFile); + newOwner.commit(); + + verify(repositoryContext, times(1)).recordConnectionStatusEvent(connectionStatusEventCaptor.capture()); + final ConnectionStatusEvent connectionStatusEvent = connectionStatusEventCaptor.getValue(); + assertEquals(INPUT_CONNECTION_ID, connectionStatusEvent.getComponentMetricContext().id()); + assertEquals(BACK_PRESSURE_BYTES_THRESHOLD, connectionStatusEvent.getBackPressureBytesThreshold()); + assertEquals(objectCount, connectionStatusEvent.getQueuedCount()); + assertEquals(byteCount, connectionStatusEvent.getQueuedBytes()); + assertEquals(LoadBalanceStatus.LOAD_BALANCE_NOT_CONFIGURED, connectionStatusEvent.getLoadBalanceStatus()); + assertEquals(FlowFileAvailability.FLOWFILE_AVAILABLE, connectionStatusEvent.getFlowFileAvailability()); + } + @Test void testExportToPathFlowFileEventBytes() throws IOException { setRepositoryContext();