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