Merge pull request #30936 from vespa-engine/balder/stream-prometheus-metricsv8.332.5

To avoid keeping all Prometheus Metrics in memory while aggregating, …

6 files changed, 75 insertions, 55 deletions

diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsHandler.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsHandler.java
index 0277bda078f..58b51020bb9 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsHandler.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/application/ApplicationMetricsHandler.java
@@ -20,7 +20,6 @@ import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.Executor;
 import java.util.logging.Level;
-import java.util.stream.Collectors;
 
 import static ai.vespa.metricsproxy.http.ValuesFetcher.getConsumerOrDefault;
 import static ai.vespa.metricsproxy.metric.model.json.GenericJsonUtil.toGenericApplicationModel;
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/DimensionId.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/DimensionId.java
index 43cc8fda3c9..fe35822c4af 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/DimensionId.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/DimensionId.java
@@ -1,8 +1,8 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.metricsproxy.metric.model;
 
+import ai.vespa.metricsproxy.metric.model.prometheus.PrometheusUtil;
 import com.yahoo.concurrent.CopyOnWriteHashMap;
-import io.prometheus.client.Collector;
 
 import java.util.Map;
 import java.util.Objects;
@@ -17,7 +17,7 @@ public final class DimensionId {
     private final String idForPrometheus;
     private DimensionId(String id) {
         this.id = id;
-        idForPrometheus = Collector.sanitizeMetricName(id);
+        idForPrometheus = PrometheusUtil.sanitize(id);
     }
 
     public static DimensionId toDimensionId(String id) {
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricId.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricId.java
index 829eb06101f..0080ec86272 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricId.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/MetricId.java
@@ -1,8 +1,8 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.metricsproxy.metric.model;
 
+import ai.vespa.metricsproxy.metric.model.prometheus.PrometheusUtil;
 import com.yahoo.concurrent.CopyOnWriteHashMap;
-import io.prometheus.client.Collector;
 
 import java.util.Map;
 import java.util.Objects;
@@ -18,7 +18,7 @@ public class MetricId {
     private final String idForPrometheus;
     private MetricId(String id) {
         this.id = id;
-        idForPrometheus = Collector.sanitizeMetricName(id);
+        idForPrometheus = PrometheusUtil.sanitize(id);
     }
 
     public static MetricId toMetricId(String id) {
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/ServiceId.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/ServiceId.java
index 28c64b012c1..20641545a7d 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/ServiceId.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/ServiceId.java
@@ -1,7 +1,7 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.metricsproxy.metric.model;
 
-import io.prometheus.client.Collector;
+import ai.vespa.metricsproxy.metric.model.prometheus.PrometheusUtil;
 
 import java.util.Objects;
 
@@ -14,7 +14,7 @@ public class ServiceId {
     private final String idForPrometheus;
     private ServiceId(String id) {
         this.id = id;
-        idForPrometheus = Collector.sanitizeMetricName(id);
+        idForPrometheus = PrometheusUtil.sanitize(id);
     }
 
     public static ServiceId toServiceId(String id) { return new ServiceId(id); }
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusModel.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusModel.java
index 77db2389271..1f55e2f8679 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusModel.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusModel.java
@@ -1,38 +1,51 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.metricsproxy.metric.model.prometheus;
 
+import ai.vespa.metricsproxy.metric.model.DimensionId;
+import ai.vespa.metricsproxy.metric.model.MetricId;
+import ai.vespa.metricsproxy.metric.model.MetricsPacket;
+import ai.vespa.metricsproxy.metric.model.ServiceId;
 import io.prometheus.client.Collector;
+import io.prometheus.client.Collector.MetricFamilySamples;
+import io.prometheus.client.Collector.MetricFamilySamples.Sample;
 import io.prometheus.client.exporter.common.TextFormat;
 
-import java.io.IOException;
 import java.io.StringWriter;
+import java.util.ArrayList;
 import java.util.Enumeration;
 import java.util.Iterator;
 import java.util.List;
-import java.util.logging.Level;
-import java.util.logging.Logger;
+import java.util.Map;
+import java.util.Set;
 
 /**
  * @author yj-jtakagi
  * @author gjoranv
  */
-public class PrometheusModel implements Enumeration<Collector.MetricFamilySamples> {
-    private static Logger log = Logger.getLogger(PrometheusModel.class.getName());
+public class PrometheusModel implements Enumeration<MetricFamilySamples> {
+    private final Map<ServiceId, List<MetricsPacket>> packetsByServiceId;
+    private final Iterator<MetricId> metricIterator;
+    private final Iterator<MetricFamilySamples> statusMetrics;
 
-    private final Iterator<Collector.MetricFamilySamples> metricFamilySamplesIterator;
-
-    PrometheusModel(List<Collector.MetricFamilySamples> metricFamilySamples) {
-        this.metricFamilySamplesIterator = metricFamilySamples.iterator();
+    PrometheusModel(Set<MetricId> metricNames, Map<ServiceId,
+                    List<MetricsPacket>> packetsByServiceId,
+                    List<MetricFamilySamples> statusMetrics)
+    {
+        metricIterator = metricNames.iterator();
+        this.packetsByServiceId = packetsByServiceId;
+        this.statusMetrics = statusMetrics.iterator();
     }
 
     @Override
     public boolean hasMoreElements() {
-        return metricFamilySamplesIterator.hasNext();
+        return metricIterator.hasNext() || statusMetrics.hasNext();
     }
 
     @Override
-    public Collector.MetricFamilySamples nextElement() {
-        return metricFamilySamplesIterator.next();
+    public MetricFamilySamples nextElement() {
+        return metricIterator.hasNext()
+                ? createMetricFamily(metricIterator.next())
+                : statusMetrics.next();
     }
 
     public String serialize() {
@@ -45,4 +58,31 @@ public class PrometheusModel implements Enumeration<Collector.MetricFamilySample
         return writer.toString();
     }
 
+    private MetricFamilySamples createMetricFamily(MetricId metricId) {
+        List<MetricFamilySamples.Sample> sampleList = new ArrayList<>();
+        packetsByServiceId.forEach(((serviceId, packets) -> {
+            for (var packet : packets) {
+                Number metric = packet.metrics().get(metricId);
+                if (metric != null) {
+                    sampleList.add(createSample(serviceId, metricId, metric, packet.timestamp, packet.dimensions()));
+                }
+            }
+        }));
+        return new MetricFamilySamples(metricId.getIdForPrometheus(), Collector.Type.UNKNOWN, "", sampleList);
+    }
+    private static Sample createSample(ServiceId serviceId, MetricId metricId, Number metric,
+                                       Long timeStamp, Map<DimensionId, String> dimensions)
+    {
+        List<String> labels = new ArrayList<>(dimensions.size());
+        List<String> labelValues = new ArrayList<>(dimensions.size());
+        for (var entry : dimensions.entrySet()) {
+            var labelName = entry.getKey().getIdForPrometheus();
+            labels.add(labelName);
+            labelValues.add(entry.getValue());
+        }
+        labels.add("vespa_service");
+        labelValues.add(serviceId.getIdForPrometheus());
+        return new Sample(metricId.getIdForPrometheus(), labels, labelValues, metric.doubleValue(), timeStamp);
+    }
+
 }
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusUtil.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusUtil.java
index 2b0db5381bc..47b97fa0902 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusUtil.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/metric/model/prometheus/PrometheusUtil.java
@@ -1,16 +1,17 @@
 // Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package ai.vespa.metricsproxy.metric.model.prometheus;
 
+import ai.vespa.metricsproxy.metric.model.MetricId;
 import ai.vespa.metricsproxy.metric.model.MetricsPacket;
 import ai.vespa.metricsproxy.metric.model.ServiceId;
 import io.prometheus.client.Collector;
 import io.prometheus.client.Collector.MetricFamilySamples;
-import io.prometheus.client.Collector.MetricFamilySamples.Sample;
 
 import java.util.ArrayList;
-import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 /**
@@ -19,54 +20,34 @@ import java.util.stream.Collectors;
  */
 public class PrometheusUtil {
 
+    public static String sanitize(String name) {
+        String sanitized = Collector.sanitizeMetricName(name);
+        return name.equals(sanitized) ? name : sanitized;
+    }
+
     public static PrometheusModel toPrometheusModel(List<MetricsPacket> metricsPackets) {
+        Set<MetricId> metricNames = new HashSet<>();
+        for (MetricsPacket metricsPacket : metricsPackets) {
+            metricNames.addAll(metricsPacket.metrics().keySet());
+        }
+
         Map<ServiceId, List<MetricsPacket>> packetsByService = metricsPackets.stream()
                 .collect(Collectors.groupingBy(packet -> packet.service));
 
-        List<MetricFamilySamples> metricFamilySamples = new ArrayList<>(packetsByService.size());
-
-        Map<String, List<Sample>> samples = new HashMap<>();
+        List<MetricFamilySamples> statusMetrics = new ArrayList<>(packetsByService.size());
         packetsByService.forEach(((serviceId, packets) -> {
-
             var serviceName = serviceId.getIdForPrometheus();
-            for (var packet : packets) {
-                Long timeStamp = packet.timestamp * 1000;
-                var dimensions = packet.dimensions();
-                List<String> labels = new ArrayList<>(dimensions.size());
-                List<String> labelValues = new ArrayList<>(dimensions.size());
-                for (var entry : dimensions.entrySet()) {
-                    var labelName = entry.getKey().getIdForPrometheus();
-                    labels.add(labelName);
-                    labelValues.add(entry.getValue());
-                }
-                labels.add("vespa_service");
-                labelValues.add(serviceName);
-
-                for (var metric : packet.metrics().entrySet()) {
-                    var metricName = metric.getKey().getIdForPrometheus();
-                    List<Sample> sampleList;
-                    if (samples.containsKey(metricName)) {
-                        sampleList = samples.get(metricName);
-                    } else {
-                        sampleList = new ArrayList<>();
-                        samples.put(metricName, sampleList);
-                        metricFamilySamples.add(new MetricFamilySamples(metricName, Collector.Type.UNKNOWN, "", sampleList));
-                    }
-                    sampleList.add(new Sample(metricName, labels, labelValues, metric.getValue().doubleValue(), timeStamp));
-                }
-            }
             if (!packets.isEmpty()) {
                 var firstPacket = packets.get(0);
                 var statusMetricName = serviceName + "_status";
                 // MetricsPacket status 0 means OK, but it's the opposite in Prometheus.
                 var statusMetricValue = (firstPacket.statusCode == 0) ? 1 : 0;
-                var sampleList = List.of(new Sample(statusMetricName, List.of(), List.of(),
+                var sampleList = List.of(new Collector.MetricFamilySamples.Sample(statusMetricName, List.of(), List.of(),
                         statusMetricValue, firstPacket.timestamp * 1000));
-                metricFamilySamples.add(new MetricFamilySamples(statusMetricName, Collector.Type.UNKNOWN, "status of service", sampleList));
+                statusMetrics.add(new Collector.MetricFamilySamples(statusMetricName, Collector.Type.UNKNOWN, "status of service", sampleList));
             }
         }));
-
-        return new PrometheusModel(metricFamilySamples);
+        return new PrometheusModel(metricNames, packetsByService, statusMetrics);
     }
 
 }