Avoid buffering entire prometheus model serialisation

4 files changed, 43 insertions, 6 deletions

diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/PrometheusResponse.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/PrometheusResponse.java
new file mode 100644
index 00000000000..c9e60510d56
--- /dev/null
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/PrometheusResponse.java
@@ -0,0 +1,30 @@
+package ai.vespa.metricsproxy.http;
+
+import ai.vespa.metricsproxy.metric.model.prometheus.PrometheusModel;
+import com.yahoo.container.jdisc.HttpResponse;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.io.OutputStreamWriter;
+import java.io.Writer;
+
+/**
+ * @author jonmv
+ */
+public class PrometheusResponse extends HttpResponse {
+
+    private final PrometheusModel model;
+
+    public PrometheusResponse(int status, PrometheusModel model) {
+        super(status);
+        this.model = model;
+    }
+
+    @Override
+    public void render(OutputStream outputStream) throws IOException {
+        Writer writer = new OutputStreamWriter(outputStream);
+        model.serialize(writer);
+        writer.flush();
+    }
+
+}
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 58b51020bb9..0d6228150ed 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
@@ -3,7 +3,7 @@
 package ai.vespa.metricsproxy.http.application;
 
 import ai.vespa.metricsproxy.core.MetricsConsumers;
-import ai.vespa.metricsproxy.http.TextResponse;
+import ai.vespa.metricsproxy.http.PrometheusResponse;
 import ai.vespa.metricsproxy.metric.model.ConsumerId;
 import ai.vespa.metricsproxy.metric.model.DimensionId;
 import ai.vespa.metricsproxy.metric.model.MetricsPacket;
@@ -75,7 +75,7 @@ public class ApplicationMetricsHandler extends HttpHandlerBase {
         }
     }
 
-    private TextResponse applicationPrometheusResponse(String requestedConsumer) {
+    private HttpResponse applicationPrometheusResponse(String requestedConsumer) {
         ConsumerId consumer = getConsumerOrDefault(requestedConsumer, metricsConsumers);
         var metricsByNode = metricsRetriever.getMetrics(consumer);
 
@@ -87,7 +87,7 @@ public class ApplicationMetricsHandler extends HttpHandlerBase {
                         .map(builder -> builder.putDimension(DimensionId.toDimensionId("hostname"), element.hostname))
                         .map(MetricsPacket.Builder::build))
                 .toList();
-        return new TextResponse(200, toPrometheusModel(metricsForAllNodes).serialize());
+        return new PrometheusResponse(200, toPrometheusModel(metricsForAllNodes));
     }
 
 }
diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/prometheus/PrometheusHandler.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/prometheus/PrometheusHandler.java
index d73561b5eff..e609b54b916 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/prometheus/PrometheusHandler.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/prometheus/PrometheusHandler.java
@@ -3,6 +3,7 @@ package ai.vespa.metricsproxy.http.prometheus;
 
 import ai.vespa.metricsproxy.core.MetricsConsumers;
 import ai.vespa.metricsproxy.core.MetricsManager;
+import ai.vespa.metricsproxy.http.PrometheusResponse;
 import ai.vespa.metricsproxy.http.TextResponse;
 import ai.vespa.metricsproxy.http.ValuesFetcher;
 import ai.vespa.metricsproxy.metric.dimensions.ApplicationDimensions;
@@ -56,11 +57,11 @@ public class PrometheusHandler extends HttpHandlerBase {
         return Optional.empty();
     }
 
-    private TextResponse valuesResponse(String consumer) {
+    private HttpResponse valuesResponse(String consumer) {
         try {
             List<MetricsPacket> metrics =  new ArrayList<>(valuesFetcher.fetch(consumer));
             metrics.addAll(nodeMetricGatherer.gatherMetrics());
-            return new TextResponse(OK, toPrometheusModel(metrics).serialize());
+            return new PrometheusResponse(OK, toPrometheusModel(metrics));
         } catch (Exception e) {
             log.log(Level.WARNING, "Got exception when rendering metrics:", e);
             return new TextResponse(INTERNAL_SERVER_ERROR, e.getMessage());
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 fd00ad67ec0..0f3878821f3 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
@@ -11,6 +11,7 @@ import io.prometheus.client.Collector.MetricFamilySamples.Sample;
 import io.prometheus.client.exporter.common.TextFormat;
 
 import java.io.StringWriter;
+import java.io.Writer;
 import java.util.ArrayList;
 import java.util.Enumeration;
 import java.util.Iterator;
@@ -50,12 +51,16 @@ public class PrometheusModel implements Enumeration<MetricFamilySamples> {
 
     public String serialize() {
         var writer = new StringWriter();
+        serialize(writer);
+        return writer.toString();
+    }
+
+    public void serialize(Writer writer) {
         try {
             TextFormat.write004(writer, this);
         } catch (Exception e) {
             throw new PrometheusRenderingException("Could not render metrics. Check the log for details.", e);
         }
-        return writer.toString();
     }
 
     private MetricFamilySamples createMetricFamily(MetricId metricId) {
@@ -70,6 +75,7 @@ public class PrometheusModel implements Enumeration<MetricFamilySamples> {
         }));
         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)
     {