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.ArrayList;
015import java.util.Arrays;
016import java.util.Enumeration;
017import java.util.List;
018import java.util.concurrent.atomic.AtomicInteger;
019import java.util.function.Predicate;
020import java.util.zip.GZIPOutputStream;
021
022/** Prometheus scrape endpoint. */
023public class PrometheusScrapeHandler {
024
025  private final PrometheusRegistry registry;
026  private final ExpositionFormats expositionFormats;
027  private final Predicate<String> nameFilter;
028  private final AtomicInteger lastResponseSize = new AtomicInteger(2 << 9); //  0.5 MB
029  private final List<String> supportedFormats;
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    supportedFormats = new ArrayList<>(Arrays.asList("openmetrics", "text"));
048    if (expositionFormats.getPrometheusProtobufWriter().isAvailable()) {
049      supportedFormats.add("prometheus-protobuf");
050    }
051  }
052
053  public void handleRequest(PrometheusHttpExchange exchange) throws IOException {
054    try {
055      PrometheusHttpRequest request = exchange.getRequest();
056      MetricSnapshots snapshots = scrape(request);
057      if (writeDebugResponse(snapshots, exchange)) {
058        return;
059      }
060      ByteArrayOutputStream responseBuffer =
061          new ByteArrayOutputStream(lastResponseSize.get() + 1024);
062      String acceptHeader = request.getHeader("Accept");
063      ExpositionFormatWriter writer = expositionFormats.findWriter(acceptHeader);
064      writer.write(responseBuffer, snapshots);
065      lastResponseSize.set(responseBuffer.size());
066      PrometheusHttpResponse response = exchange.getResponse();
067      response.setHeader("Content-Type", writer.getContentType());
068
069      if (shouldUseCompression(request)) {
070        response.setHeader("Content-Encoding", "gzip");
071        try (GZIPOutputStream gzipOutputStream =
072            new GZIPOutputStream(response.sendHeadersAndGetBody(200, 0))) {
073          responseBuffer.writeTo(gzipOutputStream);
074        }
075      } else {
076        int contentLength = responseBuffer.size();
077        if (contentLength > 0) {
078          response.setHeader("Content-Length", String.valueOf(contentLength));
079        }
080        if (request.getMethod().equals("HEAD")) {
081          // The HTTPServer implementation will throw an Exception if we close the output stream
082          // without sending a response body, so let's not close the output stream in case of a HEAD
083          // response.
084          response.sendHeadersAndGetBody(200, -1);
085        } else {
086          try (OutputStream outputStream = response.sendHeadersAndGetBody(200, contentLength)) {
087            responseBuffer.writeTo(outputStream);
088          }
089        }
090      }
091    } catch (IOException e) {
092      exchange.handleException(e);
093    } catch (RuntimeException e) {
094      exchange.handleException(e);
095    } finally {
096      exchange.close();
097    }
098  }
099
100  private Predicate<String> makeNameFilter(ExporterFilterProperties props) {
101    if (props.getAllowedMetricNames() == null
102        && props.getExcludedMetricNames() == null
103        && props.getAllowedMetricNamePrefixes() == null
104        && props.getExcludedMetricNamePrefixes() == null) {
105      return null;
106    } else {
107      return MetricNameFilter.builder()
108          .nameMustBeEqualTo(props.getAllowedMetricNames())
109          .nameMustNotBeEqualTo(props.getExcludedMetricNames())
110          .nameMustStartWith(props.getAllowedMetricNamePrefixes())
111          .nameMustNotStartWith(props.getExcludedMetricNamePrefixes())
112          .build();
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 MetricSnapshots scrape(PrometheusHttpRequest request) {
130
131    Predicate<String> filter = makeNameFilter(request.getParameterValues("name[]"));
132    if (filter != null) {
133      return registry.scrape(filter, request);
134    } else {
135      return registry.scrape(request);
136    }
137  }
138
139  private boolean writeDebugResponse(MetricSnapshots snapshots, PrometheusHttpExchange exchange)
140      throws IOException {
141    String debugParam = exchange.getRequest().getParameter("debug");
142    PrometheusHttpResponse response = exchange.getResponse();
143    if (debugParam == null) {
144      return false;
145    } else {
146      response.setHeader("Content-Type", "text/plain; charset=utf-8");
147      int responseStatus = supportedFormats.contains(debugParam) ? 200 : 500;
148      OutputStream body = response.sendHeadersAndGetBody(responseStatus, 0);
149      switch (debugParam) {
150        case "openmetrics":
151          expositionFormats.getOpenMetricsTextFormatWriter().write(body, snapshots);
152          break;
153        case "text":
154          expositionFormats.getPrometheusTextFormatWriter().write(body, snapshots);
155          break;
156        case "prometheus-protobuf":
157          String debugString =
158              expositionFormats.getPrometheusProtobufWriter().toDebugString(snapshots);
159          body.write(debugString.getBytes(StandardCharsets.UTF_8));
160          break;
161        default:
162          body.write(
163              ("debug="
164                      + debugParam
165                      + ": Unsupported query parameter. Valid values are 'openmetrics', "
166                      + "'text', and 'prometheus-protobuf'.")
167                  .getBytes(StandardCharsets.UTF_8));
168          break;
169      }
170      return true;
171    }
172  }
173
174  private boolean shouldUseCompression(PrometheusHttpRequest request) {
175    Enumeration<String> encodingHeaders = request.getHeaders("Accept-Encoding");
176    if (encodingHeaders == null) {
177      return false;
178    }
179    while (encodingHeaders.hasMoreElements()) {
180      String encodingHeader = encodingHeaders.nextElement();
181      String[] encodings = encodingHeader.split(",");
182      for (String encoding : encodings) {
183        if (encoding.trim().equalsIgnoreCase("gzip")) {
184          return true;
185        }
186      }
187    }
188    return false;
189  }
190}