001package io.prometheus.metrics.model.registry;
002
003import static io.prometheus.metrics.model.snapshots.PrometheusNaming.prometheusName;
004
005import io.prometheus.metrics.model.snapshots.MetricSnapshot;
006import io.prometheus.metrics.model.snapshots.MetricSnapshots;
007import java.util.List;
008import java.util.Set;
009import java.util.concurrent.ConcurrentHashMap;
010import java.util.concurrent.CopyOnWriteArrayList;
011import java.util.function.Predicate;
012
013public class PrometheusRegistry {
014
015  public static final PrometheusRegistry defaultRegistry = new PrometheusRegistry();
016
017  private final Set<String> prometheusNames = ConcurrentHashMap.newKeySet();
018  private final List<Collector> collectors = new CopyOnWriteArrayList<>();
019  private final List<MultiCollector> multiCollectors = new CopyOnWriteArrayList<>();
020
021  public void register(Collector collector) {
022    String prometheusName = collector.getPrometheusName();
023    if (prometheusName != null) {
024      if (!prometheusNames.add(prometheusName)) {
025        throw new IllegalStateException(
026            "Can't register "
027                + prometheusName
028                + " because a metric with that name is already registered.");
029      }
030    }
031    collectors.add(collector);
032  }
033
034  public void register(MultiCollector collector) {
035    for (String prometheusName : collector.getPrometheusNames()) {
036      if (!prometheusNames.add(prometheusName)) {
037        throw new IllegalStateException(
038            "Can't register " + prometheusName + " because that name is already registered.");
039      }
040    }
041    multiCollectors.add(collector);
042  }
043
044  public void unregister(Collector collector) {
045    collectors.remove(collector);
046    String prometheusName = collector.getPrometheusName();
047    if (prometheusName != null) {
048      prometheusNames.remove(collector.getPrometheusName());
049    }
050  }
051
052  public void unregister(MultiCollector collector) {
053    multiCollectors.remove(collector);
054    for (String prometheusName : collector.getPrometheusNames()) {
055      prometheusNames.remove(prometheusName(prometheusName));
056    }
057  }
058
059  public void clear() {
060    collectors.clear();
061    multiCollectors.clear();
062    prometheusNames.clear();
063  }
064
065  public MetricSnapshots scrape() {
066    return scrape((PrometheusScrapeRequest) null);
067  }
068
069  public MetricSnapshots scrape(PrometheusScrapeRequest scrapeRequest) {
070    MetricSnapshots.Builder result = MetricSnapshots.builder();
071    for (Collector collector : collectors) {
072      MetricSnapshot snapshot =
073          scrapeRequest == null ? collector.collect() : collector.collect(scrapeRequest);
074      if (snapshot != null) {
075        if (result.containsMetricName(snapshot.getMetadata().getName())) {
076          throw new IllegalStateException(
077              snapshot.getMetadata().getPrometheusName() + ": duplicate metric name.");
078        }
079        result.metricSnapshot(snapshot);
080      }
081    }
082    for (MultiCollector collector : multiCollectors) {
083      MetricSnapshots snapshots =
084          scrapeRequest == null ? collector.collect() : collector.collect(scrapeRequest);
085      for (MetricSnapshot snapshot : snapshots) {
086        if (result.containsMetricName(snapshot.getMetadata().getName())) {
087          throw new IllegalStateException(
088              snapshot.getMetadata().getPrometheusName() + ": duplicate metric name.");
089        }
090        result.metricSnapshot(snapshot);
091      }
092    }
093    return result.build();
094  }
095
096  public MetricSnapshots scrape(Predicate<String> includedNames) {
097    if (includedNames == null) {
098      return scrape();
099    }
100    return scrape(includedNames, null);
101  }
102
103  public MetricSnapshots scrape(
104      Predicate<String> includedNames, PrometheusScrapeRequest scrapeRequest) {
105    if (includedNames == null) {
106      return scrape(scrapeRequest);
107    }
108    MetricSnapshots.Builder result = MetricSnapshots.builder();
109    for (Collector collector : collectors) {
110      String prometheusName = collector.getPrometheusName();
111      // prometheusName == null means the name is unknown, and we have to scrape to learn the name.
112      // prometheusName != null means we can skip the scrape if the name is excluded.
113      if (prometheusName == null || includedNames.test(prometheusName)) {
114        MetricSnapshot snapshot =
115            scrapeRequest == null
116                ? collector.collect(includedNames)
117                : collector.collect(includedNames, scrapeRequest);
118        if (snapshot != null) {
119          result.metricSnapshot(snapshot);
120        }
121      }
122    }
123    for (MultiCollector collector : multiCollectors) {
124      List<String> prometheusNames = collector.getPrometheusNames();
125      // empty prometheusNames means the names are unknown, and we have to scrape to learn the
126      // names.
127      // non-empty prometheusNames means we can exclude the collector if all names are excluded by
128      // the filter.
129      boolean excluded = !prometheusNames.isEmpty();
130      for (String prometheusName : prometheusNames) {
131        if (includedNames.test(prometheusName)) {
132          excluded = false;
133          break;
134        }
135      }
136      if (!excluded) {
137        MetricSnapshots snapshots =
138            scrapeRequest == null
139                ? collector.collect(includedNames)
140                : collector.collect(includedNames, scrapeRequest);
141        for (MetricSnapshot snapshot : snapshots) {
142          if (snapshot != null) {
143            result.metricSnapshot(snapshot);
144          }
145        }
146      }
147    }
148    return result.build();
149  }
150}