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.expositionformats.OpenMetrics2TextFormatWriter;
010import io.prometheus.metrics.model.registry.MetricNameFilter;
011import io.prometheus.metrics.model.registry.PrometheusRegistry;
012import io.prometheus.metrics.model.snapshots.MetricSnapshots;
013import java.io.ByteArrayOutputStream;
014import java.io.IOException;
015import java.io.OutputStream;
016import java.nio.charset.StandardCharsets;
017import java.util.ArrayList;
018import java.util.Arrays;
019import java.util.Enumeration;
020import java.util.List;
021import java.util.concurrent.atomic.AtomicInteger;
022import java.util.function.Predicate;
023import java.util.zip.GZIPOutputStream;
024import javax.annotation.Nullable;
025
026/** Prometheus scrape endpoint. */
027@StableApi
028public class PrometheusScrapeHandler {
029
030  private final PrometheusRegistry registry;
031  private final ExpositionFormats expositionFormats;
032  @Nullable private final Predicate<String> nameFilter;
033  private final AtomicInteger lastResponseSize = new AtomicInteger(2 << 9); //  0.5 MB
034  private final List<String> supportedFormats;
035  private final boolean preferUncompressedResponse;
036
037  public PrometheusScrapeHandler() {
038    this(PrometheusProperties.get(), PrometheusRegistry.defaultRegistry);
039  }
040
041  public PrometheusScrapeHandler(PrometheusRegistry registry) {
042    this(PrometheusProperties.get(), registry);
043  }
044
045  public PrometheusScrapeHandler(PrometheusProperties config) {
046    this(config, PrometheusRegistry.defaultRegistry);
047  }
048
049  public PrometheusScrapeHandler(PrometheusProperties config, PrometheusRegistry registry) {
050    this.expositionFormats = ExpositionFormats.init(config);
051    this.preferUncompressedResponse =
052        config.getExporterHttpServerProperties().isPreferUncompressedResponse();
053    this.registry = registry;
054    this.nameFilter = makeNameFilter(config.getExporterFilterProperties());
055    supportedFormats = new ArrayList<>(Arrays.asList("openmetrics", "text"));
056    if (expositionFormats.getPrometheusProtobufWriter().isAvailable()) {
057      supportedFormats.add("prometheus-protobuf");
058    }
059  }
060
061  public void handleRequest(PrometheusHttpExchange exchange) throws IOException {
062    try {
063      PrometheusHttpRequest request = exchange.getRequest();
064      String[] includedNames = null;
065      String debugParam = null;
066      try {
067        includedNames = request.getParameterValues("name[]");
068        debugParam = request.getParameter("debug");
069      } catch (InvalidQueryParameterException e) {
070        writeInvalidQueryParametersResponse(exchange);
071        return;
072      }
073      MetricSnapshots snapshots = scrape(request, includedNames);
074      String acceptHeader = request.getHeader("Accept");
075      if (writeDebugResponse(
076          snapshots, EscapingScheme.fromAcceptHeader(acceptHeader), debugParam, exchange)) {
077        return;
078      }
079      ExpositionFormatWriter writer = expositionFormats.findWriter(acceptHeader);
080      EscapingScheme escapingScheme =
081          EscapingScheme.fromAcceptHeader(
082              acceptHeader, defaultEscapingScheme(writer, acceptHeader));
083      PrometheusHttpResponse response = exchange.getResponse();
084      response.setHeader("Content-Type", writer.getContentType());
085
086      if (shouldUseCompression(request)) {
087        response.setHeader("Content-Encoding", "gzip");
088        try (GZIPOutputStream gzipOutputStream =
089            new GZIPOutputStream(response.sendHeadersAndGetBody(200, 0))) {
090          writer.write(gzipOutputStream, snapshots, escapingScheme);
091        }
092      } else {
093        ByteArrayOutputStream responseBuffer =
094            new ByteArrayOutputStream(lastResponseSize.get() + 1024);
095        writer.write(responseBuffer, snapshots, escapingScheme);
096        lastResponseSize.set(responseBuffer.size());
097        int contentLength = responseBuffer.size();
098        if (contentLength > 0) {
099          response.setHeader("Content-Length", String.valueOf(contentLength));
100        }
101        if (request.getMethod().equals("HEAD")) {
102          // The HTTPServer implementation will throw an Exception if we close the output stream
103          // without sending a response body, so let's not close the output stream in case of a HEAD
104          // response.
105          response.sendHeadersAndGetBody(200, -1);
106        } else {
107          try (OutputStream outputStream = response.sendHeadersAndGetBody(200, contentLength)) {
108            responseBuffer.writeTo(outputStream);
109          }
110        }
111      }
112    } catch (IOException e) {
113      exchange.handleException(e);
114    } catch (RuntimeException e) {
115      exchange.handleException(e);
116    } finally {
117      exchange.close();
118    }
119  }
120
121  /**
122   * OpenMetrics 2.0 supports UTF-8 names and has no escaping parameter, so names are not escaped
123   * unless the scraper asks for it. With contentNegotiation=false the OM2 writer also serves OM 1.0
124   * and unversioned requests, so 2.0.0 must have been requested explicitly.
125   */
126  private EscapingScheme defaultEscapingScheme(
127      ExpositionFormatWriter writer, @Nullable String acceptHeader) {
128    if (writer instanceof OpenMetrics2TextFormatWriter
129        && "2.0.0".equals(parseOpenMetricsVersion(acceptHeader))) {
130      return EscapingScheme.ALLOW_UTF8;
131    }
132    return EscapingScheme.DEFAULT;
133  }
134
135  @Nullable
136  private static String parseOpenMetricsVersion(@Nullable String acceptHeader) {
137    if (acceptHeader == null) {
138      return null;
139    }
140    for (String mediaType : acceptHeader.split(",")) {
141      if (mediaType.contains("application/openmetrics-text")) {
142        for (String param : mediaType.split(";")) {
143          String[] tokens = param.split("=");
144          if (tokens.length == 2 && tokens[0].trim().equals("version")) {
145            return tokens[1].trim();
146          }
147        }
148        return null;
149      }
150    }
151    return null;
152  }
153
154  @Nullable
155  private Predicate<String> makeNameFilter(ExporterFilterProperties props) {
156    if (props.getAllowedMetricNames() == null
157        && props.getExcludedMetricNames() == null
158        && props.getAllowedMetricNamePrefixes() == null
159        && props.getExcludedMetricNamePrefixes() == null) {
160      return null;
161    } else {
162      return MetricNameFilter.builder()
163          .nameMustBeEqualTo(props.getAllowedMetricNames())
164          .nameMustNotBeEqualTo(props.getExcludedMetricNames())
165          .nameMustStartWith(props.getAllowedMetricNamePrefixes())
166          .nameMustNotStartWith(props.getExcludedMetricNamePrefixes())
167          .build();
168    }
169  }
170
171  @Nullable
172  private Predicate<String> makeNameFilter(@Nullable String[] includedNames) {
173    Predicate<String> result = null;
174    if (includedNames != null && includedNames.length > 0) {
175      result = MetricNameFilter.builder().nameMustBeEqualTo(includedNames).build();
176    }
177    if (result != null && nameFilter != null) {
178      result = result.and(nameFilter);
179    } else if (nameFilter != null) {
180      result = nameFilter;
181    }
182    return result;
183  }
184
185  private MetricSnapshots scrape(PrometheusHttpRequest request, @Nullable String[] includedNames) {
186
187    Predicate<String> filter = makeNameFilter(includedNames);
188    if (filter != null) {
189      return registry.scrape(filter, request);
190    } else {
191      return registry.scrape(request);
192    }
193  }
194
195  private boolean writeDebugResponse(
196      MetricSnapshots snapshots,
197      EscapingScheme escapingScheme,
198      @Nullable String debugParam,
199      PrometheusHttpExchange exchange)
200      throws IOException {
201    PrometheusHttpResponse response = exchange.getResponse();
202    if (debugParam == null) {
203      return false;
204    } else {
205      response.setHeader("Content-Type", "text/plain; charset=utf-8");
206      int responseStatus = supportedFormats.contains(debugParam) ? 200 : 500;
207      OutputStream body = response.sendHeadersAndGetBody(responseStatus, 0);
208      switch (debugParam) {
209        case "openmetrics":
210          expositionFormats.getOpenMetricsTextFormatWriter().write(body, snapshots, escapingScheme);
211          break;
212        case "text":
213          expositionFormats.getPrometheusTextFormatWriter().write(body, snapshots, escapingScheme);
214          break;
215        case "prometheus-protobuf":
216          String debugString =
217              expositionFormats
218                  .getPrometheusProtobufWriter()
219                  .toDebugString(snapshots, escapingScheme);
220          body.write(debugString.getBytes(StandardCharsets.UTF_8));
221          break;
222        default:
223          body.write(
224              ("debug="
225                      + debugParam
226                      + ": Unsupported query parameter. Valid values are 'openmetrics', "
227                      + "'text', and 'prometheus-protobuf'.")
228                  .getBytes(StandardCharsets.UTF_8));
229          break;
230      }
231      return true;
232    }
233  }
234
235  private void writeInvalidQueryParametersResponse(PrometheusHttpExchange exchange)
236      throws IOException {
237    PrometheusHttpResponse response = exchange.getResponse();
238    response.setHeader("Content-Type", "text/plain; charset=utf-8");
239    byte[] message = "Invalid query parameters".getBytes(StandardCharsets.UTF_8);
240    try (OutputStream outputStream = response.sendHeadersAndGetBody(400, message.length)) {
241      outputStream.write(message);
242    }
243  }
244
245  private boolean shouldUseCompression(PrometheusHttpRequest request) {
246    if (preferUncompressedResponse) {
247      return false;
248    }
249
250    Enumeration<String> encodingHeaders = request.getHeaders("Accept-Encoding");
251    if (encodingHeaders == null) {
252      return false;
253    }
254    while (encodingHeaders.hasMoreElements()) {
255      String encodingHeader = encodingHeaders.nextElement();
256      String[] encodings = encodingHeader.split(",");
257      for (String encoding : encodings) {
258        if (encoding.trim().equalsIgnoreCase("gzip")) {
259          return true;
260        }
261      }
262    }
263    return false;
264  }
265}