kafka-clients-metrics: Refactor node topic metrics

instrumentation/kafka-clients-metrics-3.0.0/src/main/java/com/nr/instrumentation/kafka/ClientType.java

@@ -0,0 +1,23 @@
+/*
+ *
+ *  * Copyright 2025 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.nr.instrumentation.kafka;
+
+public enum ClientType {
+    CONSUMER("Consume"),
+    PRODUCER("Produce");
+
+    private final String operation;
+
+    ClientType(String operation) {
+        this.operation = operation;
+    }
+
+    public String getOperation() {
+        return operation;
+    }
+}

instrumentation/kafka-clients-metrics-3.0.0/src/main/java/com/nr/instrumentation/kafka/MetricsScheduler.java

@@ -78,15 +78,11 @@ public void run() {
                     }
                 }
 
-                for (NewRelicMetricsReporter.NodeMetricNames consumerNodeMetricNames : nrMetricsReporter.getNodes().values()) {
+                for (String nodeTopicName : nrMetricsReporter.getNodeTopicNames()) {
                     if (METRICS_AS_EVENTS) {
-                        for (String eventName : consumerNodeMetricNames.getEventNames()) {
-                            eventData.put(eventName, 1f);
-                        }
+                        eventData.put(nodeTopicName, 1f);
                     } else {
-                        for (String metricName : consumerNodeMetricNames.getMetricNames()) {
-                            NewRelic.recordMetric(metricName, 1f);
-                        }
+                        NewRelic.recordMetric(nodeTopicName, 1f);
                     }
                 }
 

instrumentation/kafka-clients-metrics-3.0.0/src/main/java/com/nr/instrumentation/kafka/NewRelicMetricsReporter.java

@@ -8,44 +8,35 @@
 package com.nr.instrumentation.kafka;
 
 import com.newrelic.agent.bridge.AgentBridge;
+import org.apache.kafka.common.Node;
 import org.apache.kafka.common.metrics.KafkaMetric;
 import org.apache.kafka.common.metrics.MetricsReporter;
 
-import java.util.Collections;
-import java.util.HashSet;
+import java.util.Collection;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.logging.Level;
 
 import static com.nr.instrumentation.kafka.MetricsConstants.KAFKA_METRICS_DEBUG;
-import static com.nr.instrumentation.kafka.MetricsConstants.NODE_PREFIX;
 
 public class NewRelicMetricsReporter implements MetricsReporter {
 
 
     private final Map<String, KafkaMetric> metrics = new ConcurrentHashMap<>();
 
-    private final Map<String, NodeMetricNames> nodes;
+    private final NodeTopicRegistry nodeTopicRegistry;
 
-    public NewRelicMetricsReporter() {
-        this.nodes = Collections.emptyMap();
-    }
-
-    public NewRelicMetricsReporter(Set<String> nodes, Mode mode) {
-        this.nodes = new ConcurrentHashMap<>(nodes.size());
-        for(String node: nodes) {
-            this.nodes.put(node, new NodeMetricNames(node, mode));
-        }
+    public NewRelicMetricsReporter(ClientType clientType, Collection<Node> nodes) {
+        this.nodeTopicRegistry = new NodeTopicRegistry(clientType, nodes);
     }
 
     public Map<String, KafkaMetric> getMetrics() {
         return this.metrics;
     }
 
-    public Map<String, NodeMetricNames> getNodes() {
-        return nodes;
+    public Collection<String> getNodeTopicNames() {
+        return this.nodeTopicRegistry.getNodeTopicNames();
     }
 
     @Override
@@ -56,13 +47,15 @@ public void init(final List<KafkaMetric> initMetrics) {
                 AgentBridge.getAgent().getLogger().log(Level.FINEST, "init(): {0} = {1}", metricGroupAndName, kafkaMetric.metricName());
             }
             metrics.put(metricGroupAndName, kafkaMetric);
+            nodeTopicRegistry.register(kafkaMetric);
         }
         MetricsScheduler.addMetricsReporter(this);
     }
 
     @Override
     public void metricChange(final KafkaMetric metric) {
         String metricGroupAndName = getMetricGroupAndName(metric);
+        nodeTopicRegistry.register(metric);
         if (KAFKA_METRICS_DEBUG) {
             AgentBridge.getAgent().getLogger().log(Level.FINEST, "metricChange(): {0} = {1}", metricGroupAndName, metric.metricName());
         }
@@ -81,109 +74,20 @@ public void metricRemoval(final KafkaMetric metric) {
     private String getMetricGroupAndName(final KafkaMetric metric) {
         if (metric.metricName().tags().containsKey("topic")) {
             String topic = metric.metricName().tags().get("topic");
-            addTopicToNodeMetrics(topic);
-
             // Special case for handling topic names in metrics
             return metric.metricName().group() + "/" + topic + "/" + metric.metricName().name();
         }
         return metric.metricName().group() + "/" + metric.metricName().name();
     }
 
-    private void addTopicToNodeMetrics(String topic) {
-        for (NodeMetricNames nodeMetricNames : nodes.values()) {
-            nodeMetricNames.addMetricNameForTopic(topic);
-        }
-    }
-
     @Override
     public void close() {
         MetricsScheduler.removeMetricsReporter(this);
         metrics.clear();
+        nodeTopicRegistry.close();
     }
 
     @Override
     public void configure(final Map<String, ?> configs) {
     }
-
-    /**
-     * This class is used to track all the metric names that are related to a specific node:
-     *
-     * - MessageBroker/Kafka/Nodes/host:port
-     * - MessageBroker/Kafka/Nodes/host:port/Consume/topicName
-     * - MessageBroker/Kafka/Nodes/host:port/Produce/topicName
-     *
-     * At initialization time we only have the node and the mode (is this a metrics reporter
-     * for a Kafka consumer or for a Kafka producer?).
-     *
-     * Then, as topics are discovered through the metricChange method, the topic metric names are
-     * generated. This is the best way we have to get track of the topics since they're not
-     * available when the KafkaConsumer/KafkaProducer is initialized.
-     *
-     * For KafkaConsumer, the SubscriptionState doesn't contain the topics and partitions
-     * at initialization time because it takes time for the rebalance to happen.
-     *
-     * For KafkaProducer, topics are dynamic since a producer could send records to any
-     * topic and the concept of subscription doesn't exist there.
-     *
-     * Alternatively we could get the topics from the records in KafkaProducer.doSend or
-     * KafkaConsumer.poll, and call NewRelicMetricsReporter.addTopicToNodeMetrics from there.
-     * This approach would have a small impact in performance, and getting the topics from the
-     * KafkaMetrics is a good enough solution.
-     */
-    public static class NodeMetricNames {
-
-        private final String node;
-        private final Mode mode;
-
-        private final Set<String> topics = new HashSet<>();
-
-        private final Set<String> metricNames = new HashSet<>();
-        private final Set<String> eventNames = new HashSet<>();
-
-        public NodeMetricNames(String node, Mode mode) {
-            this.node = node;
-            this.mode = mode;
-
-            String nodeMetricName = NODE_PREFIX + node;
-            metricNames.add(nodeMetricName);
-            eventNames.add(getEventNameForMetric(nodeMetricName));
-        }
-
-        private void addMetricNameForTopic(String topic) {
-            if (!topics.contains(topic)) {
-                String nodeTopicMetricName = NODE_PREFIX + node + "/" + mode.getMetricSegmentName() + "/" + topic;
-                metricNames.add(nodeTopicMetricName);
-                eventNames.add(getEventNameForMetric(nodeTopicMetricName));
-
-                topics.add(topic);
-            }
-        }
-
-        private String getEventNameForMetric(String metricName) {
-            return metricName.replace('/', '.');
-        }
-
-        public Set<String> getMetricNames() {
-            return metricNames;
-        }
-
-        public Set<String> getEventNames() {
-            return eventNames;
-        }
-    }
-
-    public enum Mode {
-        CONSUMER("Consume"),
-        PRODUCER("Produce");
-
-        private final String metricSegmentName;
-
-        Mode(String metricSegmentName) {
-            this.metricSegmentName = metricSegmentName;
-        }
-
-        public String getMetricSegmentName() {
-            return metricSegmentName;
-        }
-    }
 }

instrumentation/kafka-clients-metrics-3.0.0/src/main/java/com/nr/instrumentation/kafka/NodeTopicRegistry.java

@@ -0,0 +1,92 @@
+/*
+ *
+ *  * Copyright 2025 New Relic Corporation. All rights reserved.
+ *  * SPDX-License-Identifier: Apache-2.0
+ *
+ */
+
+package com.nr.instrumentation.kafka;
+
+import org.apache.kafka.common.Node;
+import org.apache.kafka.common.metrics.KafkaMetric;
+
+import java.util.Collection;
+import java.util.HashSet;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static com.nr.instrumentation.kafka.MetricsConstants.NODE_PREFIX;
+
+/**
+ * This class is used to track all the metric names that are related to a specific node:
+ *
+ * - MessageBroker/Kafka/Nodes/host:port
+ * - MessageBroker/Kafka/Nodes/host:port/Consume/topicName
+ * - MessageBroker/Kafka/Nodes/host:port/Produce/topicName
+ *
+ * At initialization time we only have the node and the mode (is this a metrics reporter
+ * for a Kafka consumer or for a Kafka producer?).
+ *
+ * Then, as topics are discovered through the metricChange method, the topic metric names are
+ * generated. This is the best way we have to get track of the topics since they're not
+ * available when the KafkaConsumer/KafkaProducer is initialized.
+ *
+ * For KafkaConsumer, the SubscriptionState doesn't contain the topics and partitions
+ * at initialization time because it takes time for the rebalance to happen.
+ *
+ * For KafkaProducer, topics are dynamic since a producer could send records to any
+ * topic and the concept of subscription doesn't exist there.
+ *
+ * Alternatively we could get the topics from the records in KafkaProducer.doSend or
+ * KafkaConsumer.poll, and call NewRelicMetricsReporter.addTopicToNodeMetrics from there.
+ * This approach would have a small impact in performance, and getting the topics from the
+ * KafkaMetrics is a good enough solution.
+ */
+public class NodeTopicRegistry {
+    private final Set<String> recordedTopics = ConcurrentHashMap.newKeySet();
+    // contains the registered metric or event names, according to config
+    private final Set<String> convertedNames = ConcurrentHashMap.newKeySet();
+    private final Set<String> nodes = new HashSet<>();
+    private final ClientType clientType;
+
+    public NodeTopicRegistry(ClientType clientType, Collection<Node> nodes) {
+        this.clientType = clientType;
+        for (Node node : nodes) {
+            String nodeName = node.host() + ":" + node.port();
+            this.nodes.add(nodeName);
+            this.convertedNames.add(convertName(NODE_PREFIX + nodeName));
+        }
+    }
+
+    /**
+     * @return true if the metric contains a topic and it was registered
+     */
+    public boolean register(KafkaMetric metric) {
+        String topic = metric.metricName().tags().get("topic");
+        if (topic != null && recordedTopics.add(topic)) {
+            for (String node : nodes) {
+                String metricName = NODE_PREFIX + node + "/" + clientType.getOperation() + "/" + topic;
+                this.convertedNames.add(convertName(metricName));
+            }
+            return true;
+        }
+        return false;
+    }
+
+    public Collection<String> getNodeTopicNames() {
+        return this.convertedNames;
+    }
+
+    private String convertName(String metricName) {
+        if (MetricsConstants.METRICS_AS_EVENTS) {
+            return metricName.replace('/', '.');
+        }
+        return metricName;
+    }
+
+    public void close() {
+        recordedTopics.clear();
+        convertedNames.clear();
+        nodes.clear();
+    }
+}

instrumentation/kafka-clients-metrics-3.0.0/src/main/java/org/apache/kafka/clients/consumer/KafkaConsumer_Instrumentation.java

@@ -11,6 +11,7 @@
 import java.util.List;
 import java.util.Set;
 
+import com.nr.instrumentation.kafka.ClientType;
 import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
 import org.apache.kafka.common.errors.WakeupException;
 import org.apache.kafka.common.metrics.Metrics;
@@ -41,11 +42,7 @@ public class KafkaConsumer_Instrumentation<K, V> {
     public KafkaConsumer_Instrumentation() {
         if (!initialized) {
             List<Node> nodes = metadata.fetch().nodes();
-            Set<String> nodeNames = new HashSet<>(nodes.size());
-            for (Node node : nodes) {
-                nodeNames.add(node.host() + ":" + node.port());
-            }
-            metrics.addReporter(new NewRelicMetricsReporter(nodeNames, NewRelicMetricsReporter.Mode.CONSUMER));
+            metrics.addReporter(new NewRelicMetricsReporter(ClientType.CONSUMER, nodes));
             initialized = true;
         }
     }

instrumentation/kafka-clients-metrics-3.0.0/src/main/java/org/apache/kafka/clients/producer/KafkaProducer_Instrumentation.java

@@ -7,6 +7,7 @@
 
 package org.apache.kafka.clients.producer;
 
+import com.nr.instrumentation.kafka.ClientType;
 import org.apache.kafka.clients.producer.internals.ProducerMetadata;
 import org.apache.kafka.common.metrics.Metrics;
 import org.apache.kafka.common.Node;
@@ -44,11 +45,7 @@ public class KafkaProducer_Instrumentation<K, V> {
     public KafkaProducer_Instrumentation() {
         if (!initialized) {
             List<Node> nodes = metadata.fetch().nodes();
-            Set<String> nodeNames = new HashSet<>(nodes.size());
-            for (Node node : nodes) {
-                nodeNames.add(node.host() + ":" + node.port());
-            }
-            metrics.addReporter(new NewRelicMetricsReporter(nodeNames, NewRelicMetricsReporter.Mode.PRODUCER));
+            metrics.addReporter(new NewRelicMetricsReporter(ClientType.PRODUCER, nodes));
             initialized = true;
         }
     }