Add new response class, and limit to max 1MB pending output

2 files changed, 36 insertions, 0 deletions

diff --git a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/MetricsJsonResponse.java b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/MetricsJsonResponse.java
new file mode 100644
index 00000000000..b927db790b2
--- /dev/null
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/MetricsJsonResponse.java
@@ -0,0 +1,31 @@
+package ai.vespa.metricsproxy.http;
+
+import com.yahoo.container.jdisc.HttpResponse;
+
+import java.io.IOException;
+import java.io.OutputStream;
+import java.util.function.Consumer;
+
+/**
+ * @author jonmv
+ */
+public class MetricsJsonResponse extends HttpResponse {
+
+    private final Consumer<OutputStream> modelWriter;
+
+    public MetricsJsonResponse(int status, Consumer<OutputStream> modelWriter) {
+        super(status);
+        this.modelWriter = modelWriter;
+    }
+
+    @Override
+    public void render(OutputStream outputStream) throws IOException {
+        modelWriter.accept(outputStream);
+    }
+
+    @Override
+    public long maxPendingBytes() {
+        return 1 << 20;
+    }
+
+}
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
index c9e60510d56..e0c74671c9c 100644
--- a/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/PrometheusResponse.java
+++ b/metrics-proxy/src/main/java/ai/vespa/metricsproxy/http/PrometheusResponse.java
@@ -27,4 +27,9 @@ public class PrometheusResponse extends HttpResponse {
         writer.flush();
     }
 
+    @Override
+    public long maxPendingBytes() {
+        return 1 << 20;
+    }
+
 }