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