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