Spring Boot Health Check as Prometheus format

HeathPrometheusEndpoint.java

package ru.formatko;

import com.fasterxml.jackson.databind.ObjectMapper;
import io.prometheus.client.Collector;
import io.prometheus.client.exporter.common.TextFormat;
import java.io.IOException;
import java.io.StringWriter;
import java.io.Writer;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import lombok.Data;
import lombok.SneakyThrows;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.actuate.endpoint.annotation.Endpoint;
import org.springframework.boot.actuate.endpoint.annotation.ReadOperation;
import org.springframework.boot.actuate.endpoint.web.WebEndpointResponse;
import org.springframework.boot.actuate.health.HealthEndpoint;
import org.springframework.boot.actuate.health.HttpCodeStatusMapper;
import org.springframework.boot.actuate.health.Status;
import org.springframework.stereotype.Component;

/**
 * Example:
 * # HELP health_status HealthCheck result in prometheus's response format
 * # TYPE health_status gauge
 * health_status{application="java-service",type="main",} 1.0
 * health_status{application="java-service",type="db",database="PostgreSQL",validationQuery="isValid()",} 1.0
 * health_status{application="java-service",type="diskSpace",total="506332180480",exists="true",threshold="10485760",free="412188921856",} 1.0
 * health_status{application="java-service",type="ping",} 1.0
*/
@Component
@Endpoint(id = "health-check")
public class HeathPrometheusEndpoint {

    private static final String APPLICATION = "application";
    private static final String TYPE = "type";
    public static final String SAMPLE_HEALTH_STATUS = "health_status";

    private final HealthEndpoint healthEndpoint;
    private final String appName;
    private final ObjectMapper mapper;
    private final HttpCodeStatusMapper httpCodeStatusMapper;

    public HeathPrometheusEndpoint(HealthEndpoint healthEndpoint,
                                   ObjectMapper mapper,
                                   @Value("${spring.application.name:}") String appName,
                                   HttpCodeStatusMapper httpCodeStatusMapper) {
        this.healthEndpoint = healthEndpoint;
        this.mapper = mapper;
        this.appName = appName;
        this.httpCodeStatusMapper = httpCodeStatusMapper;
    }

    @ReadOperation(produces = TextFormat.CONTENT_TYPE_004)
    public WebEndpointResponse<String> healthPrometheus() {
        StatusDto status = createStatusDto();
        List<Collector.MetricFamilySamples.Sample> samples = new ArrayList<>();
        samples.add(createMainSample(status));
        samples.addAll(createComponentSamples(status));
        return createStringWebEndpointResponse(status, createMetricFamily(samples));
    }

    @SneakyThrows
    private StatusDto createStatusDto() {
        return mapper.readValue(mapper.writeValueAsString(healthEndpoint.health()), StatusDto.class);
    }


    private Collector.MetricFamilySamples.Sample createMainSample(StatusDto status) {
        Labels labels = new Labels();
        labels.add(APPLICATION, appName);
        labels.add(TYPE, "main");

        return createSample(SAMPLE_HEALTH_STATUS, labels, status.getStatus());
    }

    private List<Collector.MetricFamilySamples.Sample> createComponentSamples(StatusDto status) {
        List<Collector.MetricFamilySamples.Sample> list = new ArrayList<>();
        for (Map.Entry<String, StatusDto> entry : status.components.entrySet()) {
            Labels labels = new Labels();
            labels.add(APPLICATION, appName);
            labels.add(TYPE, entry.getKey());

            StatusDto statusDto = entry.getValue();
            Map<String, Object> details = statusDto.getDetails();
            if (details != null && !details.isEmpty()) {
                details.forEach((k, v) -> labels.add(k, String.valueOf(v)));
            }
            list.add(createSample(SAMPLE_HEALTH_STATUS, labels, statusDto.getStatus()));
        }
        return list;
    }

    private Collector.MetricFamilySamples.Sample createSample(String name, Labels labels, Status status) {
        double v = Status.UP.equals(status) ? 1 : 0;
        return new Collector.MetricFamilySamples.Sample(name, labels.getLabels(), labels.getValues(), v);
    }

    private Collector.MetricFamilySamples createMetricFamily(List<Collector.MetricFamilySamples.Sample> s) {
        return new Collector.MetricFamilySamples(
            "health_status", Collector.Type.GAUGE,
            "HealthCheck result in prometheus's response format", s);
    }

    private WebEndpointResponse<String> createStringWebEndpointResponse(
        StatusDto status, Collector.MetricFamilySamples metricFamilySamples
    ) {
        try {
            Writer writer = new StringWriter();
            TextFormat.write004(writer,
                Collections.enumeration(Collections.singletonList(metricFamilySamples)));

            return wrapResponse(writer.toString(), status);
        } catch (IOException ex) {
            // This actually never happens since StringWriter::write() doesn't throw any
            // IOException
            throw new RuntimeException("Writing metrics failed", ex);
        }
    }

    private WebEndpointResponse<String> wrapResponse(String body, StatusDto status) {
        if (body == null || body.isEmpty()) {
            return new WebEndpointResponse<>("", 500);
        } else {
            int statusCode = httpCodeStatusMapper.getStatusCode(status.getStatus());
            return new WebEndpointResponse<>(body, statusCode);
        }
    }

    public static class Labels {
        private final Map<String, String> map = new HashMap<>();

        public void add(String label, String value) {
            if (value != null && !value.isEmpty()) {
                map.put(label, value);
            }
        }

        public List<String> getLabels() {
            return new ArrayList<>(map.keySet());
        }

        public List<String> getValues() {
            return new ArrayList<>(map.values());
        }
    }

    @Data
    public static class StatusDto {
        private Status status;
        private Map<String, StatusDto> components;
        private Map<String, Object> details;
    }
}