NIFI-16321 Add Attribute Handling for Counters and Gauges

nifi-extension-bundles/nifi-groovyx-bundle/nifi-groovyx-processors/src/main/java/org/apache/nifi/processors/groovyx/flow/ProcessSessionWrap.java

@@ -312,11 +312,21 @@ public void adjustCounter(String name, long delta, boolean immediate) {
         session.adjustCounter(name, delta, immediate);
     }
 
+    @Override
+    public void adjustCounter(final String name, final long delta, final Map<String, String> attributes, final CommitTiming commitTiming) {
+        session.adjustCounter(name, delta, attributes, commitTiming);
+    }
+
     @Override
     public void recordGauge(final String name, final double value, final CommitTiming commitTiming) {
         session.recordGauge(name, value, commitTiming);
     }
 
+    @Override
+    public void recordGauge(final String name, final double value, final Map<String, String> attributes, final CommitTiming commitTiming) {
+        session.recordGauge(name, value, attributes, commitTiming);
+    }
+
     /**
      * @return FlowFile that is next highest priority FlowFile to process. Otherwise returns null.
      */

nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/CounterRecord.java

@@ -17,19 +17,49 @@
 package org.apache.nifi.controller.metrics;
 
 import java.time.Instant;
+import java.util.Map;
+import java.util.Objects;
 
 /**
  * Single measurement for a named Counter recorded during processing
  *
  * @param name Counter Name
  * @param value Counter Value
+ * @param attributes Map of keys and values associated with the Counter measurement, which may be empty but not null
  * @param recorded Timestamp when the Component recorded the Counter value
  * @param componentMetricContext Context for Component Metric record
  */
 public record CounterRecord(
         String name,
         long value,
+        Map<String, String> attributes,
         Instant recorded,
         ComponentMetricContext componentMetricContext
 ) {
+    public CounterRecord {
+        attributes = Map.copyOf(Objects.requireNonNull(attributes, "Attributes required"));
+    }
+
+    /**
+     * Counter Record constructor for compatibility with earlier versions
+     *
+     * @param name Counter Name
+     * @param value Counter Value
+     * @param recorded Timestamp when the Processor recorded the Counter value
+     * @param componentMetricContext Context for Component Metric record
+     */
+    public CounterRecord(
+            final String name,
+            final long value,
+            final Instant recorded,
+            final ComponentMetricContext componentMetricContext
+    ) {
+        this(
+                name,
+                value,
+                Map.of(),
+                recorded,
+                componentMetricContext
+        );
+    }
 }

nifi-framework-api/src/main/java/org/apache/nifi/controller/metrics/GaugeRecord.java

@@ -17,19 +17,49 @@
 package org.apache.nifi.controller.metrics;
 
 import java.time.Instant;
+import java.util.Map;
+import java.util.Objects;
 
 /**
  * Single measurement for a named Gauge recorded during processing
  *
  * @param name Gauge Name
  * @param value Gauge Value
+ * @param attributes Map of keys and values associated with the Gauge measurement, which may be empty but not null
  * @param recorded Timestamp when the Processor recorded the Gauge value
  * @param componentMetricContext Context for Component Metric record
  */
 public record GaugeRecord(
         String name,
         double value,
+        Map<String, String> attributes,
         Instant recorded,
         ComponentMetricContext componentMetricContext
 ) {
+    public GaugeRecord {
+        attributes = Map.copyOf(Objects.requireNonNull(attributes, "Attributes required"));
+    }
+
+    /**
+     * Gauge Record constructor for compatibility with earlier versions
+     *
+     * @param name Gauge Name
+     * @param value Gauge Value
+     * @param recorded Timestamp when the Processor recorded the Gauge value
+     * @param componentMetricContext Context for Component Metric record
+     */
+    public GaugeRecord(
+            final String name,
+            final double value,
+            final Instant recorded,
+            final ComponentMetricContext componentMetricContext
+    ) {
+        this(
+                name,
+                value,
+                Map.of(),
+                recorded,
+                componentMetricContext
+        );
+    }
 }

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/AbstractRepositoryContext.java

@@ -42,6 +42,7 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Predicate;
 
@@ -172,11 +173,11 @@ private boolean pollFromSelfLoopsOnly() {
     }
 
     @Override
-    public void adjustCounter(final String name, final long delta) {
+    public void adjustCounter(final String name, final long delta, final Map<String, String> attributes) {
         counterRepo.adjustCounter(componentNameCounterContext, name, delta);
         counterRepo.adjustCounter(componentTypeCounterContext, name, delta);
 
-        final CounterRecord counterRecord = new CounterRecord(name, delta, Instant.now(), componentMetricContext);
+        final CounterRecord counterRecord = new CounterRecord(name, delta, attributes, Instant.now(), componentMetricContext);
         componentMetricReporter.recordCounter(counterRecord);
     }
 

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/RepositoryContext.java

@@ -34,6 +34,7 @@
 
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
 import java.util.function.Predicate;
 
 public interface RepositoryContext {
@@ -69,7 +70,7 @@ public interface RepositoryContext {
 
     long getNextFlowFileSequence();
 
-    void adjustCounter(String name, long delta);
+    void adjustCounter(String name, long delta, Map<String, String> attributes);
 
     void recordGauge(GaugeRecord gaugeRecord);
 

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/StandardProcessSession.java

@@ -165,8 +165,8 @@ public class StandardProcessSession implements ProcessSession, ProvenanceEventEn
     private final String connectableDescription;
     private final PerformanceTracker performanceTracker;
 
-    private Map<String, Long> countersOnCommit;
-    private Map<String, Long> immediateCounters;
+    private Map<CounterKey, Long> countersOnCommit;
+    private Map<CounterKey, Long> immediateCounters;
     private List<GaugeRecord> gaugeRecordsSessionCommitted;
 
     private final Set<String> removedFlowFiles = new HashSet<>();
@@ -692,8 +692,9 @@ protected void commit(final Checkpoint checkpoint, final boolean asynchronous) {
                 }
             }
 
-            for (final Map.Entry<String, Long> entry : checkpoint.countersOnCommit.entrySet()) {
-                context.adjustCounter(entry.getKey(), entry.getValue());
+            for (final Map.Entry<CounterKey, Long> entry : checkpoint.countersOnCommit.entrySet()) {
+                final CounterKey counterKey = entry.getKey();
+                context.adjustCounter(counterKey.name(), entry.getValue(), counterKey.attributes());
             }
 
             for (final GaugeRecord gaugeRecord : checkpoint.gaugeRecordsSessionCommitted) {
@@ -889,23 +890,36 @@ private LoadBalanceStatus getLoadBalanceStatus(final FlowFileQueue flowFileQueue
         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();
+    private Map<String, Long> combineCounters(final Map<CounterKey, Long> first, final Map<CounterKey, Long> second) {
+        final Map<String, Long> firstValues = getCounterValues(first);
+        final Map<String, Long> secondValues = getCounterValues(second);
 
-        if (firstEmpty && secondEmpty) {
-            return null;
+        if (firstValues == null) {
+            return secondValues;
         }
-        if (firstEmpty) {
-            return second;
+        if (secondValues == null) {
+            return firstValues;
         }
-        if (secondEmpty) {
-            return first;
+
+        secondValues.forEach((name, value) -> firstValues.merge(name, value, Long::sum));
+        return firstValues;
+    }
+
+    /**
+     * Reduce Counter measurements to values keyed by Counter name, summing the measurements recorded for a name with
+     * differing attributes, since FlowFile Events track Counter values by name alone.
+     *
+     * @param counters Counter measurements which may be null or empty
+     * @return Counter values keyed by Counter name, or null when no measurements were recorded
+     */
+    private Map<String, Long> getCounterValues(final Map<CounterKey, Long> counters) {
+        if (counters == null || counters.isEmpty()) {
+            return null;
         }
 
-        final Map<String, Long> combined = new HashMap<>(first);
-        second.forEach((key, value) -> combined.merge(key, value, Long::sum));
-        return combined;
+        final Map<String, Long> counterValues = new HashMap<>();
+        counters.forEach((counterKey, value) -> counterValues.merge(counterKey.name(), value, Long::sum));
+        return counterValues;
     }
 
     private void addEventType(final Map<String, BitSet> map, final String id, final ProvenanceEventType eventType) {
@@ -1410,7 +1424,7 @@ protected synchronized void rollback(final boolean penalize, final boolean rollb
         final ProcessSessionEvent flowFileEvent = ProcessSessionEventBuilder.forComponent(context.getComponentMetricContext())
                 .bytesRead(bytesRead)
                 .bytesWritten(bytesWritten)
-                .counters(immediateCounters)
+                .counters(getCounterValues(immediateCounters))
                 .build();
 
         // update event repository
@@ -2009,11 +2023,17 @@ private void handleConflictingId(final FlowFileRecord flowFile, final Connection
 
     @Override
     public void recordGauge(final String name, final double value, final CommitTiming commitTiming) {
+        recordGauge(name, value, Map.of(), commitTiming);
+    }
+
+    @Override
+    public void recordGauge(final String name, final double value, final Map<String, String> attributes, final CommitTiming commitTiming) {
         Objects.requireNonNull(name, "Gauge Name required");
+        Objects.requireNonNull(attributes, "Gauge Attributes required");
         Objects.requireNonNull(commitTiming, "Commit Timing required");
 
         final Instant recorded = Instant.now();
-        final GaugeRecord gaugeRecord = new GaugeRecord(name, value, recorded, context.getComponentMetricContext());
+        final GaugeRecord gaugeRecord = new GaugeRecord(name, value, Map.copyOf(attributes), recorded, context.getComponentMetricContext());
 
         if (CommitTiming.NOW == commitTiming) {
             context.recordGauge(gaugeRecord);
@@ -2027,14 +2047,25 @@ public void recordGauge(final String name, final double value, final CommitTimin
 
     @Override
     public void adjustCounter(final String name, final long delta, final boolean immediate) {
+        adjustCounter(name, delta, Map.of(), immediate ? CommitTiming.NOW : CommitTiming.SESSION_COMMITTED);
+    }
+
+    @Override
+    public void adjustCounter(final String name, final long delta, final Map<String, String> attributes, final CommitTiming commitTiming) {
+        Objects.requireNonNull(name, "Counter Name required");
+        Objects.requireNonNull(attributes, "Counter Attributes required");
+        Objects.requireNonNull(commitTiming, "Commit Timing required");
+
+        final boolean immediate = CommitTiming.NOW == commitTiming;
+
         // If we are adjusting the counter immediately, allow it even if the task is terminated. The contract states:
         // "the counter will be updated immediately, without regard to whether the session is committed or rolled back"
         // so we need to ensure that we allow adjusting the counter even after the task is terminated.
         if (!immediate) {
             verifyTaskActive();
         }
 
-        final Map<String, Long> counters;
+        final Map<CounterKey, Long> counters;
         if (immediate) {
             if (immediateCounters == null) {
                 immediateCounters = new HashMap<>();
@@ -2047,13 +2078,17 @@ public void adjustCounter(final String name, final long delta, final boolean imm
             counters = countersOnCommit;
         }
 
+        // Measurements are aggregated for each distinct combination of Counter name and attributes
+        final Map<String, String> counterAttributes = Map.copyOf(attributes);
+        final CounterKey counterKey = new CounterKey(name, counterAttributes);
+
         // Set current value or adjust when found
-        counters.compute(name, (currentName, currentValue) ->
+        counters.compute(counterKey, (currentKey, currentValue) ->
             currentValue == null ? delta : currentValue + delta
         );
 
         if (immediate) {
-            context.adjustCounter(name, delta);
+            context.adjustCounter(name, delta, counterAttributes);
         }
     }
 
@@ -4071,6 +4106,15 @@ private interface ConnectionPoller {
         List<FlowFileRecord> poll(Connection connection, Set<FlowFileRecord> expiredRecords);
     }
 
+    /**
+     * Key for aggregating Counter measurements recorded under the same Counter name with the same attributes
+     *
+     * @param name Counter name
+     * @param attributes Immutable Map of keys and values associated with the Counter measurement
+     */
+    private record CounterKey(String name, Map<String, String> attributes) {
+    }
+
     protected static class Checkpoint {
 
         private long processingTime = 0L;
@@ -4085,8 +4129,8 @@ protected static class Checkpoint {
         private Map<String, Connection> processedConnections;
         private Map<String, ComponentMetricContext> connectionMetricContexts;
 
-        private Map<String, Long> countersOnCommit;
-        private Map<String, Long> immediateCounters;
+        private Map<CounterKey, Long> countersOnCommit;
+        private Map<CounterKey, Long> immediateCounters;
 
         private List<GaugeRecord> gaugeRecordsSessionCommitted;
 

nifi-framework-bundle/nifi-framework/nifi-framework-components/src/main/java/org/apache/nifi/controller/repository/WeakHashMapProcessSessionFactory.java

@@ -145,11 +145,21 @@ public void adjustCounter(final String name, final long delta, final boolean imm
             delegate.adjustCounter(name, delta, immediate);
         }
 
+        @Override
+        public void adjustCounter(final String name, final long delta, final Map<String, String> attributes, final CommitTiming commitTiming) {
+            delegate.adjustCounter(name, delta, attributes, commitTiming);
+        }
+
         @Override
         public void recordGauge(final String name, final double value, final CommitTiming commitTiming) {
             delegate.recordGauge(name, value, commitTiming);
         }
 
+        @Override
+        public void recordGauge(final String name, final double value, final Map<String, String> attributes, final CommitTiming commitTiming) {
+            delegate.recordGauge(name, value, attributes, commitTiming);
+        }
+
         @Override
         public FlowFile get() {
             return delegate.get();