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