Merge pull request #24013 from vespa-engine/olaa/prometheus-metric-packet-handler

MetricsPacketsHandler supports Prometheus format

3 files changed, 112 insertions, 0 deletions

diff --git a/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java b/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java
index 8fa658bf7fc..480ee96393a 100644
--- a/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java
+++ b/container-core/src/main/java/com/yahoo/container/jdisc/state/MetricsPacketsHandler.java
@@ -18,6 +18,7 @@ import com.yahoo.jdisc.handler.ResponseDispatch;
 import com.yahoo.jdisc.handler.ResponseHandler;
 import com.yahoo.jdisc.http.HttpHeaders;
 
+import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
@@ -100,10 +101,16 @@ public class MetricsPacketsHandler extends AbstractRequestHandler {
             if (query != null && query.equals("array-formatted")) {
                 return getMetricsArray();
             }
+            if ("format=prometheus".equals(query)) {
+                return buildPrometheusOutput();
+            }
+
             String output = jsonToString(getStatusPacket()) + getAllMetricsPackets() + "\n";
             return output.getBytes(StandardCharsets.UTF_8);
         } catch (JsonProcessingException e) {
             throw new RuntimeException("Bad JSON construction.", e);
+        } catch (IOException e) {
+            throw new RuntimeException("Unexcpected IOException.", e);
         }
     }
 
@@ -120,6 +127,13 @@ public class MetricsPacketsHandler extends AbstractRequestHandler {
     }
 
     /**
+     * Returns metrics in Prometheus format
+     */
+    private byte[] buildPrometheusOutput() throws IOException {
+        return PrometheusHelper.buildPrometheusOutput(getSnapshot(), applicationName, timer.currentTimeMillis());
+    }
+
+    /**
      * Exactly one status packet is added to the response.
      */
     private JsonNode getStatusPacket() {
diff --git a/container-core/src/main/java/com/yahoo/container/jdisc/state/PrometheusHelper.java b/container-core/src/main/java/com/yahoo/container/jdisc/state/PrometheusHelper.java
new file mode 100644
index 00000000000..ca12e8161a9
--- /dev/null
+++ b/container-core/src/main/java/com/yahoo/container/jdisc/state/PrometheusHelper.java
@@ -0,0 +1,73 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.container.jdisc.state;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.Map;
+
+import static com.yahoo.container.jdisc.state.JsonUtil.sanitizeDouble;
+
+/**
+ * @author olaa
+ */
+public class PrometheusHelper {
+
+    private static final String HELP_LINE = "# HELP %s\n# TYPE %s untyped\n";
+    private static final String METRIC_LINE = "%s{%s} %s %d\n";
+
+    protected static byte[] buildPrometheusOutput(MetricSnapshot metricSnapshot, String application, long timestamp) throws IOException {
+        var outputStream = new ByteArrayOutputStream();
+
+        for (Map.Entry<MetricDimensions, MetricSet> snapshotEntry : metricSnapshot) {
+            var metricDimensions = snapshotEntry.getKey();
+            var metricSet = snapshotEntry.getValue();
+
+            var dimensionBuilder = new StringBuilder();
+            for (var dimension : metricDimensions) {
+                dimensionBuilder.append(dimension.getKey()).append("=\"").append(dimension.getValue()).append("\",");
+            }
+            dimensionBuilder.append("vespa_service=\"").append(application).append("\",");
+            var dimensions = dimensionBuilder.toString();
+
+            for (var metric : metricSet) {
+                var metricName = metric.getKey();
+                var metricValue = metric.getValue();
+
+                if (metricValue instanceof CountMetric) {
+                    var sanitizedMetricName = getSanitizedMetricName(metricName, "count");
+                    var value = ((CountMetric) metricValue).getCount();
+                    outputStream.write(getMetricLines(sanitizedMetricName, dimensions, value, timestamp));
+                } else if (metricValue instanceof GaugeMetric) {
+                    var gauge = (GaugeMetric) metricValue;
+                    writeGaugeMetrics(outputStream, metricName, gauge, dimensions, timestamp);
+                }
+            }
+        }
+        return outputStream.toByteArray();
+    }
+
+    private static void writeGaugeMetrics(OutputStream outputStream, String metricName, GaugeMetric gaugeMetric, String dimensions, long timestamp) throws IOException {
+        var sanitizedMetricName = getSanitizedMetricName(metricName, "last");
+        var value = sanitizeDouble(gaugeMetric.getLast());
+        outputStream.write(getMetricLines(sanitizedMetricName, dimensions, value, timestamp));
+
+        sanitizedMetricName = getSanitizedMetricName(metricName, "average");
+        value = sanitizeDouble(gaugeMetric.getAverage());
+        outputStream.write(getMetricLines(sanitizedMetricName, dimensions, value, timestamp));
+
+        sanitizedMetricName = getSanitizedMetricName(metricName, "max");
+        value = sanitizeDouble(gaugeMetric.getMax());
+        outputStream.write(getMetricLines(sanitizedMetricName, dimensions, value, timestamp));
+    }
+
+    private static byte[] getMetricLines(String metricName, String dimensions, Number value, long timestamp) {
+        return (String.format(HELP_LINE, metricName, metricName) +
+                String.format(METRIC_LINE, metricName, dimensions, value, timestamp)).getBytes();
+    }
+
+    private static String getSanitizedMetricName(String metricName, String suffix) {
+        return metricName.replaceAll("([-.])", "_") + "_" + suffix;
+    }
+
+}
diff --git a/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java b/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java
index 019fc1e1e78..56efa396297 100644
--- a/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java
+++ b/container-core/src/test/java/com/yahoo/container/jdisc/state/MetricsPacketsHandlerTest.java
@@ -159,6 +159,31 @@ public class MetricsPacketsHandlerTest extends StateHandlerTestBase {
             if (dimensions.has("dim2")) assertDimension(packet, "host", HOST_DIMENSION);
         });
     }
+
+    @Test
+    public void prometheus_metrics() {
+        var context = StateMetricContext.newInstance(Map.of("dim1", "value1"));
+        var snapshot = new MetricSnapshot();
+        snapshot.set(context, "gauge.metric", 0.2);
+        snapshot.add(context, "counter.metric", 5);
+        snapshotProvider.setSnapshot(snapshot);
+        var response = requestAsString("http://localhost/metrics-packets?format=prometheus");
+        var expectedResponse = """
+                # HELP gauge_metric_last
+                # TYPE gauge_metric_last untyped
+                gauge_metric_last{dim1="value1",vespa_service="state-handler-test-base",} 0.2 0
+                # HELP gauge_metric_average
+                # TYPE gauge_metric_average untyped
+                gauge_metric_average{dim1="value1",vespa_service="state-handler-test-base",} 0.2 0
+                # HELP gauge_metric_max
+                # TYPE gauge_metric_max untyped
+                gauge_metric_max{dim1="value1",vespa_service="state-handler-test-base",} 0.2 0
+                # HELP counter_metric_count
+                # TYPE counter_metric_count untyped
+                counter_metric_count{dim1="value1",vespa_service="state-handler-test-base",} 5 0
+                """;
+        assertEquals(expectedResponse, response);
+    }
     
     private List<JsonNode> incrementTimeAndGetJsonPackets() throws Exception {
         advanceToNextSnapshot();