issues/1506: ensure Prometheus /metrics endpoint streams correctly with clean EOF

Author: chris-wood-mo

api/src/main/java/io/kafbat/ui/controller/PrometheusExposeController.java

@@ -4,10 +4,12 @@
 import io.kafbat.ui.model.KafkaCluster;
 import io.kafbat.ui.service.StatisticsCache;
 import io.kafbat.ui.service.metrics.prometheus.PrometheusMetricsExposer;
+import java.nio.charset.StandardCharsets;
 import java.util.Map;
 import java.util.Optional;
 import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
+import org.springframework.http.HttpStatus;
 import org.springframework.http.ResponseEntity;
 import org.springframework.web.bind.annotation.RestController;
 import org.springframework.web.server.ServerWebExchange;
@@ -20,28 +22,48 @@ public class PrometheusExposeController extends AbstractController implements Pr
   private final StatisticsCache statisticsCache;
 
   @Override
-  public Mono<ResponseEntity<String>> exposeAllMetrics(ServerWebExchange exchange) {
-    return Mono.just(
-        PrometheusMetricsExposer.exposeAllMetrics(
-            clustersStorage.getKafkaClusters()
-                .stream()
-                .filter(KafkaCluster::isExposeMetricsViaPrometheusEndpoint)
-                .collect(Collectors.toMap(KafkaCluster::getName, c -> statisticsCache.get(c).getMetrics()))
-        )
-    );
+  public Mono<ResponseEntity<Void>> exposeAllMetrics(ServerWebExchange exchange) {
+    String metrics = PrometheusMetricsExposer.exposeAllMetrics(
+        clustersStorage.getKafkaClusters()
+            .stream()
+            .filter(KafkaCluster::isExposeMetricsViaPrometheusEndpoint)
+            .collect(Collectors.toMap(
+                KafkaCluster::getName,
+                c -> statisticsCache.get(c).getMetrics()))
+    ).getBody();
+
+    var response = exchange.getResponse();
+    response.getHeaders().add("Content-Type", "text/plain; version=0.0.4");
+
+    byte[] bytes = metrics.getBytes(StandardCharsets.UTF_8);
+    var buffer = response.bufferFactory().wrap(bytes);
+
+    return response.writeWith(Mono.just(buffer))
+        .doOnTerminate(response::setComplete)
+        .thenReturn(ResponseEntity.ok().build());
   }
 
   @Override
-  public Mono<ResponseEntity<String>> exposeClusterMetrics(String clusterName,
-                                                           ServerWebExchange exchange) {
+  public Mono<ResponseEntity<Void>> exposeClusterMetrics(String clusterName, ServerWebExchange exchange) {
     Optional<KafkaCluster> cluster = clustersStorage.getClusterByName(clusterName);
-    if (cluster.isPresent() && cluster.get().isExposeMetricsViaPrometheusEndpoint()) {
-      return Mono.just(PrometheusMetricsExposer.exposeAllMetrics(
-          Map.of(clusterName, statisticsCache.get(cluster.get()).getMetrics())
-      ));
-    } else {
-      return Mono.just(ResponseEntity.notFound().build());
+    if (cluster.isEmpty() || !cluster.get().isExposeMetricsViaPrometheusEndpoint()) {
+      exchange.getResponse().setStatusCode(HttpStatus.NOT_FOUND);
+      return exchange.getResponse().setComplete()
+          .thenReturn(ResponseEntity.notFound().build());
     }
-  }
 
+    String metrics = PrometheusMetricsExposer.exposeAllMetrics(
+        Map.of(clusterName, statisticsCache.get(cluster.get()).getMetrics())
+    ).getBody();
+
+    var response = exchange.getResponse();
+    response.getHeaders().add("Content-Type", "text/plain; version=0.0.4");
+
+    byte[] bytes = metrics.getBytes(StandardCharsets.UTF_8);
+    var buffer = response.bufferFactory().wrap(bytes);
+
+    return response.writeWith(Mono.just(buffer))
+        .doOnTerminate(response::setComplete)
+        .thenReturn(ResponseEntity.ok().build());
+  }
 }

contract-typespec/api/prometheus.tsp

@@ -13,11 +13,13 @@ interface PrometheusExposeApi {
     @get
     @summary("exposeAllMetrics")
     @operationId("exposeAllMetrics")
-    exposeAllMetrics(): string;
+    // Return nothing (the controller handles the stream manually)
+    exposeAllMetrics(): void;
 
     @get
     @route("/{clusterName}")
     @summary("exposeClusterMetrics")
     @operationId("exposeClusterMetrics")
-    exposeClusterMetrics(@path clusterName: string): string;
+    // Also return nothing
+    exposeClusterMetrics(@path clusterName: string): void;
 }