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