001package io.prometheus.metrics.exporter.common;
002
003import io.prometheus.metrics.config.ExporterFilterProperties;
004import io.prometheus.metrics.config.PrometheusProperties;
005import io.prometheus.metrics.expositionformats.ExpositionFormatWriter;
006import io.prometheus.metrics.expositionformats.ExpositionFormats;
007import io.prometheus.metrics.model.registry.MetricNameFilter;
008import io.prometheus.metrics.model.registry.PrometheusRegistry;
009import io.prometheus.metrics.model.snapshots.MetricSnapshots;
010
011import java.io.ByteArrayOutputStream;
012import java.io.IOException;
013import java.io.OutputStream;
014import java.nio.charset.StandardCharsets;
015import java.util.Arrays;
016import java.util.Enumeration;
017import java.util.concurrent.atomic.AtomicInteger;
018import java.util.function.Predicate;
019import java.util.zip.GZIPOutputStream;
020
021/**
022 * Prometheus scrape endpoint.
023 */
024public class PrometheusScrapeHandler {
025
026    private final PrometheusRegistry registry;
027    private final ExpositionFormats expositionFormats;
028    private final Predicate<String> nameFilter;
029    private AtomicInteger lastResponseSize = new AtomicInteger(2 << 9); //  0.5 MB
030
031    public PrometheusScrapeHandler() {
032        this(PrometheusProperties.get(), PrometheusRegistry.defaultRegistry);
033    }
034
035    public PrometheusScrapeHandler(PrometheusRegistry registry) {
036        this(PrometheusProperties.get(), registry);
037    }
038
039    public PrometheusScrapeHandler(PrometheusProperties config) {
040        this(config, PrometheusRegistry.defaultRegistry);
041    }
042
043    public PrometheusScrapeHandler(PrometheusProperties config, PrometheusRegistry registry) {
044        this.expositionFormats = ExpositionFormats.init(config.getExporterProperties());
045        this.registry = registry;
046        this.nameFilter = makeNameFilter(config.getExporterFilterProperties());
047    }
048
049    public void handleRequest(PrometheusHttpExchange exchange) throws IOException {
050        try {
051            PrometheusHttpRequest request = exchange.getRequest();
052            PrometheusHttpResponse response = exchange.getResponse();
053            MetricSnapshots snapshots = scrape(request);
054            if (writeDebugResponse(snapshots, exchange)) {
055                return;
056            }
057            ByteArrayOutputStream responseBuffer = new ByteArrayOutputStream(lastResponseSize.get() + 1024);
058            String acceptHeader = request.getHeader("Accept");
059            ExpositionFormatWriter writer = expositionFormats.findWriter(acceptHeader);
060            writer.write(responseBuffer, snapshots);
061            lastResponseSize.set(responseBuffer.size());
062            response.setHeader("Content-Type", writer.getContentType());
063
064            if (shouldUseCompression(request)) {
065                response.setHeader("Content-Encoding", "gzip");
066                try (GZIPOutputStream gzipOutputStream = new GZIPOutputStream(response.sendHeadersAndGetBody(200, 0))) {
067                    responseBuffer.writeTo(gzipOutputStream);
068                }
069            } else {
070                int contentLength = responseBuffer.size();
071                if (contentLength > 0) {
072                    response.setHeader("Content-Length", String.valueOf(contentLength));
073                }
074                if (request.getMethod().equals("HEAD")) {
075                    // The HTTPServer implementation will throw an Exception if we close the output stream
076                    // without sending a response body, so let's not close the output stream in case of a HEAD response.
077                    response.sendHeadersAndGetBody(200, -1);
078                } else {
079                    try (OutputStream outputStream = response.sendHeadersAndGetBody(200, contentLength)) {
080                        responseBuffer.writeTo(outputStream);
081                    }
082                }
083            }
084        } catch (IOException e) {
085            exchange.handleException(e);
086        } catch (RuntimeException e) {
087            exchange.handleException(e);
088        } finally {
089            exchange.close();
090        }
091    }
092
093    private Predicate<String> makeNameFilter(ExporterFilterProperties props) {
094        if (props.getAllowedMetricNames() == null && props.getExcludedMetricNames() == null && props.getAllowedMetricNamePrefixes() == null && props.getExcludedMetricNamePrefixes() == null) {
095            return null;
096        } else {
097            return MetricNameFilter.builder()
098                    .nameMustBeEqualTo(props.getAllowedMetricNames())
099                    .nameMustNotBeEqualTo(props.getExcludedMetricNames())
100                    .nameMustStartWith(props.getAllowedMetricNamePrefixes())
101                    .nameMustNotStartWith(props.getExcludedMetricNamePrefixes())
102                    .build();
103        }
104    }
105
106    private MetricSnapshots scrape(PrometheusHttpRequest request) {
107
108        Predicate<String> filter = makeNameFilter(request.getParameterValues("name[]"));
109        if (filter != null) {
110            return registry.scrape(filter, request);
111        } else {
112            return registry.scrape(request);
113        }
114    }
115
116    private Predicate<String> makeNameFilter(String[] includedNames) {
117        Predicate<String> result = null;
118        if (includedNames != null && includedNames.length > 0) {
119            result = MetricNameFilter.builder().nameMustBeEqualTo(includedNames).build();
120        }
121        if (result != null && nameFilter != null) {
122            result = result.and(nameFilter);
123        } else if (nameFilter != null) {
124            result = nameFilter;
125        }
126        return result;
127    }
128
129    private boolean writeDebugResponse(MetricSnapshots snapshots, PrometheusHttpExchange exchange) throws IOException {
130        String debugParam = exchange.getRequest().getParameter("debug");
131        PrometheusHttpResponse response = exchange.getResponse();
132        if (debugParam == null) {
133            return false;
134        } else {
135            response.setHeader("Content-Type", "text/plain; charset=utf-8");
136            boolean supportedFormat = Arrays.asList("openmetrics", "text", "prometheus-protobuf").contains(debugParam);
137            int responseStatus = supportedFormat ? 200 : 500;
138            OutputStream body = response.sendHeadersAndGetBody(responseStatus, 0);
139            switch (debugParam) {
140                case "openmetrics":
141                    expositionFormats.getOpenMetricsTextFormatWriter().write(body, snapshots);
142                    break;
143                case "text":
144                    expositionFormats.getPrometheusTextFormatWriter().write(body, snapshots);
145                    break;
146                case "prometheus-protobuf":
147                    String debugString = expositionFormats.getPrometheusProtobufWriter().toDebugString(snapshots);
148                    body.write(debugString.getBytes(StandardCharsets.UTF_8));
149                    break;
150                default:
151                    body.write(("debug=" + debugParam + ": Unsupported query parameter. Valid values are 'openmetrics', 'text', and 'prometheus-protobuf'.").getBytes(StandardCharsets.UTF_8));
152                    break;
153            }
154            return true;
155        }
156    }
157
158    private boolean shouldUseCompression(PrometheusHttpRequest request) {
159        Enumeration<String> encodingHeaders = request.getHeaders("Accept-Encoding");
160        if (encodingHeaders == null) {
161            return false;
162        }
163        while (encodingHeaders.hasMoreElements()) {
164            String encodingHeader = encodingHeaders.nextElement();
165            String[] encodings = encodingHeader.split(",");
166            for (String encoding : encodings) {
167                if (encoding.trim().equalsIgnoreCase("gzip")) {
168                    return true;
169                }
170            }
171        }
172        return false;
173    }
174}