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