Precompute mapping from MetricId to Map<ConfiguredMetric, Set<ConsumerId>> to avoid costly streaming and filtering on every metric.

2 files changed, 27 insertions, 20 deletions

diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/MetricsConsumers.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/MetricsConsumers.java
index 7c6cae660a7..457e27a5896 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/MetricsConsumers.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/MetricsConsumers.java
@@ -4,9 +4,11 @@ package ai.vespa.metricsproxy.core;
 
 import ai.vespa.metricsproxy.core.ConsumersConfig.Consumer;
 import ai.vespa.metricsproxy.metric.model.ConsumerId;
+import ai.vespa.metricsproxy.metric.model.MetricId;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -33,11 +35,20 @@ public class MetricsConsumers {
     // All consumers for each metric (more useful than the opposite map).
     private final Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric;
 
+    // All consumers for each metric, by metric id
+    private final Map<MetricId, Map<ConfiguredMetric, Set<ConsumerId>>> consumersByMetricByMetricId;
+
     public MetricsConsumers(ConsumersConfig config) {
         consumerMetrics = config.consumer().stream().collect(
                 toUnmodifiableLinkedMap(consumer -> ConsumerId.toConsumerId(consumer.name()), consumer -> convert(consumer.metric())));
 
         consumersByMetric = createConsumersByMetric(consumerMetrics);
+        consumersByMetricByMetricId = new HashMap<>();
+        consumersByMetric.forEach((configuredMetric, consumers) -> {
+            var consumersByMetric = consumersByMetricByMetricId.computeIfAbsent(configuredMetric.id(), id -> new HashMap<>());
+            var consumerSet = consumersByMetric.computeIfAbsent(configuredMetric, id -> new HashSet<>());
+            consumerSet.addAll(consumers);
+        });
     }
 
     /**
@@ -52,6 +63,10 @@ public class MetricsConsumers {
         return consumersByMetric;
     }
 
+    public Map<ConfiguredMetric, Set<ConsumerId>> getConsumersByMetric(MetricId id) {
+        return consumersByMetricByMetricId.get(id);
+    }
+
     public Set<ConsumerId> getAllConsumers() {
         return unmodifiableSet(consumerMetrics.keySet());
     }
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/VespaMetrics.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/VespaMetrics.java
index 9af0d0f5171..3629e81582a 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/VespaMetrics.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/core/VespaMetrics.java
@@ -23,7 +23,6 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.stream.Collectors;
 
 import static ai.vespa.metricsproxy.metric.dimensions.PublicDimensions.INTERNAL_SERVICE_ID;
 import static ai.vespa.metricsproxy.metric.model.DimensionId.toDimensionId;
@@ -69,15 +68,13 @@ public class VespaMetrics {
     public List<MetricsPacket.Builder> getMetrics(List<VespaService> services) {
         List<MetricsPacket.Builder> metricsPackets = new ArrayList<>();
 
-        Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric = metricsConsumers.getConsumersByMetric();
-
         for (VespaService service : services) {
             // One metrics packet for system metrics
             Optional<MetricsPacket.Builder> systemCheck = getSystemMetrics(service);
             systemCheck.ifPresent(metricsPackets::add);
 
             MetricAggregator aggregator = new MetricAggregator(service.getDimensions());
-            GetServiceMetricsConsumer metricsConsumer = new GetServiceMetricsConsumer(consumersByMetric, aggregator);
+            GetServiceMetricsConsumer metricsConsumer = new GetServiceMetricsConsumer(metricsConsumers, aggregator);
             service.consumeMetrics(metricsConsumer);
 
             if (! aggregator.getAggregated().isEmpty()) {
@@ -120,24 +117,27 @@ public class VespaMetrics {
      */
     private static class GetServiceMetricsConsumer implements MetricsParser.Consumer {
         private final MetricAggregator aggregator;
-        private final Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric;
-        GetServiceMetricsConsumer(Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric, MetricAggregator aggregator) {
-            this.consumersByMetric = consumersByMetric;
+        private final MetricsConsumers metricsConsumers;
+        GetServiceMetricsConsumer(MetricsConsumers metricsConsumers, MetricAggregator aggregator) {
+            this.metricsConsumers = metricsConsumers;
             this.aggregator = aggregator;
         }
 
         @Override
         public void consume(Metric candidate) {
-            getConfiguredMetrics(candidate.getName(), consumersByMetric.keySet()).forEach(
-                    configuredMetric -> aggregator.aggregate(
-                            metricWithConfigProperties(candidate, configuredMetric, consumersByMetric)));
+            Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric = metricsConsumers.getConsumersByMetric(candidate.getName());
+            if (consumersByMetric != null) {
+                consumersByMetric.keySet().forEach(
+                        configuredMetric -> aggregator.aggregate(
+                                metricWithConfigProperties(candidate, configuredMetric, consumersByMetric.get(configuredMetric))));
+            }
         }
         private static Metric metricWithConfigProperties(Metric candidate,
                                                          ConfiguredMetric configuredMetric,
-                                                         Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric) {
+                                                         Set<ConsumerId> consumers) {
             Metric metric = candidate.clone();
             metric.setDimensions(extractDimensions(candidate.getDimensions(), configuredMetric.dimension()));
-            metric.setConsumers(extractConsumers(consumersByMetric.get(configuredMetric)));
+            metric.setConsumers(extractConsumers(consumers));
 
             if (configuredMetric.outputname() != null && !configuredMetric.outputname().id.isEmpty())
                 metric.setName(configuredMetric.outputname());
@@ -159,14 +159,6 @@ public class VespaMetrics {
             }
             return consumers;
         }
-        /**
-         * Returns all configured metrics (for any consumer) that have the given id as 'name'.
-         */
-        private static Set<ConfiguredMetric> getConfiguredMetrics(MetricId id, Set<ConfiguredMetric> configuredMetrics) {
-            return configuredMetrics.stream()
-                    .filter(m -> m.id().equals(id))
-                    .collect(Collectors.toSet());
-        }
     }
 
     private Optional<MetricsPacket.Builder> getSystemMetrics(VespaService service) {