Merge pull request #20487 from vespa-engine/balder/improve-metrics-aggergation-speed

Balder/improve metrics aggergation speed

3 files changed, 52 insertions, 47 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 44eca2f57b4..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()) {
@@ -118,58 +115,50 @@ public class VespaMetrics {
      * In order to include a metric, it must exist in the given map of metric to consumers.
      * Each returned metric will contain a collection of consumers that it should be routed to.
      */
-    private class GetServiceMetricsConsumer implements MetricsParser.Consumer {
+    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 Map<DimensionId, String> extractDimensions(Map<DimensionId, String> dimensions, List<Dimension> configuredDimensions) {
-        if ( ! configuredDimensions.isEmpty()) {
-            Map<DimensionId, String> dims = new HashMap<>(dimensions);
-            configuredDimensions.forEach(d -> dims.put(d.key(), d.value()));
-            dimensions = Collections.unmodifiableMap(dims);
+        private static Metric metricWithConfigProperties(Metric candidate,
+                                                         ConfiguredMetric configuredMetric,
+                                                         Set<ConsumerId> consumers) {
+            Metric metric = candidate.clone();
+            metric.setDimensions(extractDimensions(candidate.getDimensions(), configuredMetric.dimension()));
+            metric.setConsumers(extractConsumers(consumers));
+
+            if (configuredMetric.outputname() != null && !configuredMetric.outputname().id.isEmpty())
+                metric.setName(configuredMetric.outputname());
+            return metric;
         }
-        return dimensions;
-    }
-
-    private Set<ConsumerId> extractConsumers(Set<ConsumerId> configuredConsumers) {
-        Set<ConsumerId> consumers = Collections.emptySet();
-        if (configuredConsumers != null) {
-            consumers = configuredConsumers;
+        private static Map<DimensionId, String> extractDimensions(Map<DimensionId, String> dimensions, List<Dimension> configuredDimensions) {
+            if ( ! configuredDimensions.isEmpty()) {
+                Map<DimensionId, String> dims = new HashMap<>(dimensions);
+                configuredDimensions.forEach(d -> dims.put(d.key(), d.value()));
+                dimensions = Collections.unmodifiableMap(dims);
+            }
+            return dimensions;
         }
-        return consumers;
-    }
 
-    private Metric metricWithConfigProperties(Metric candidate,
-                                              ConfiguredMetric configuredMetric,
-                                              Map<ConfiguredMetric, Set<ConsumerId>> consumersByMetric) {
-        Metric metric = candidate.clone();
-        metric.setDimensions(extractDimensions(candidate.getDimensions(), configuredMetric.dimension()));
-        metric.setConsumers(extractConsumers(consumersByMetric.get(configuredMetric)));
-
-        if (configuredMetric.outputname() != null && !configuredMetric.outputname().id.isEmpty())
-            metric.setName(configuredMetric.outputname());
-        return metric;
-    }
-
-    /**
-     * 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 static Set<ConsumerId> extractConsumers(Set<ConsumerId> configuredConsumers) {
+            Set<ConsumerId> consumers = Collections.emptySet();
+            if (configuredConsumers != null) {
+                consumers = configuredConsumers;
+            }
+            return consumers;
+        }
     }
 
     private Optional<MetricsPacket.Builder> getSystemMetrics(VespaService service) {
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/service/MetricsParser.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/service/MetricsParser.java
index 079633b28a1..8157ecb72fd 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/service/MetricsParser.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/service/MetricsParser.java
@@ -145,6 +145,7 @@ public class MetricsParser {
         }
 
         JsonNode aggregates = metric.get("values");
+        String prefix = name + ".";
         for (Iterator<?> it = aggregates.fieldNames(); it.hasNext(); ) {
             String aggregator = (String) it.next();
             JsonNode aggregatorValue = aggregates.get(aggregator);
@@ -155,7 +156,7 @@ public class MetricsParser {
             if (value == null) {
                 throw new IllegalArgumentException("Value for aggregator '" + aggregator + "' is not a number");
             }
-            String metricName = new StringBuilder().append(name).append(".").append(aggregator).toString();
+            String metricName = prefix + aggregator;
             consumer.consume(new Metric(MetricId.toMetricId(metricName), value, timestamp, dim, description));
         }
     }