NIFI-16123 Add Connection Status reporting to Process Session

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;
+    }
 }

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

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;

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

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<Long, StandardRepositoryRecord> records = new ConcurrentHashMap<>();
     private final Map<String, ProcessSessionEventBuilder> connectionCounts = new ConcurrentHashMap<>();
+    private final Map<String, Connection> processedConnections = new ConcurrentHashMap<>();
     private final Map<String, ComponentMetricContext> connectionMetricContexts = new ConcurrentHashMap<>();
     private final Map<FlowFileQueue, Set<FlowFileRecord>> unacknowledgedFlowFiles = new ConcurrentHashMap<>();
     private final Map<ContentClaim, ByteCountingOutputStream> 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<String, Long> combineCounters(final Map<String, Long> first, final Map<String, Long> 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<Long, StandardRepositoryRecord> records;
         private Map<String, ProcessSessionEventBuilder> connectionCounts;
+        private Map<String, Connection> processedConnections;
         private Map<String, ComponentMetricContext> connectionMetricContexts;
 
         private Map<String, Long> 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<P
             this.records = session.records;
 
             this.connectionCounts = session.connectionCounts;
+            this.processedConnections = session.processedConnections;
             this.connectionMetricContexts = session.connectionMetricContexts;
             this.countersOnCommit = session.countersOnCommit == null ? Collections.emptyMap() : session.countersOnCommit;
             this.immediateCounters = session.immediateCounters == null ? Collections.emptyMap() : session.immediateCounters;
@@ -4118,7 +4181,10 @@ private void copyCheckpoint(final StandardProcessSession session, final List<Pro
 
             this.records.putAll(session.records);
 
-            mergeMapsWithMutableValue(this.connectionCounts, session.connectionCounts, (destination, toMerge) -> 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);