001package io.prometheus.metrics.model.registry;
002
003import io.prometheus.metrics.model.snapshots.MetricSnapshot;
004import io.prometheus.metrics.model.snapshots.MetricSnapshots;
005
006import java.util.List;
007import java.util.Set;
008import java.util.concurrent.ConcurrentHashMap;
009import java.util.concurrent.CopyOnWriteArrayList;
010import java.util.function.Predicate;
011
012import static io.prometheus.metrics.model.snapshots.PrometheusNaming.prometheusName;
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        prometheusNames.remove(collector.getPrometheusName());
044    }
045
046    public void unregister(MultiCollector collector) {
047        multiCollectors.remove(collector);
048        for (String prometheusName : collector.getPrometheusNames()) {
049            prometheusNames.remove(prometheusName(prometheusName));
050        }
051    }
052
053    public MetricSnapshots scrape() {
054        MetricSnapshots.Builder result = MetricSnapshots.builder();
055        for (Collector collector : collectors) {
056            MetricSnapshot snapshot = collector.collect();
057            if (snapshot != null) {
058                if (result.containsMetricName(snapshot.getMetadata().getName())) {
059                    throw new IllegalStateException(snapshot.getMetadata().getPrometheusName() + ": duplicate metric name.");
060                }
061                result.metricSnapshot(snapshot);
062            }
063        }
064        for (MultiCollector collector : multiCollectors) {
065            for (MetricSnapshot snapshot : collector.collect()) {
066                if (result.containsMetricName(snapshot.getMetadata().getName())) {
067                    throw new IllegalStateException(snapshot.getMetadata().getPrometheusName() + ": duplicate metric name.");
068                }
069                result.metricSnapshot(snapshot);
070            }
071        }
072        return result.build();
073    }
074
075    public MetricSnapshots scrape(Predicate<String> includedNames) {
076        if (includedNames == null) {
077            return scrape();
078        }
079        MetricSnapshots.Builder result = MetricSnapshots.builder();
080        for (Collector collector : collectors) {
081            String prometheusName = collector.getPrometheusName();
082            if (prometheusName == null || includedNames.test(prometheusName)) {
083                MetricSnapshot snapshot = collector.collect(includedNames);
084                if (snapshot != null) {
085                    result.metricSnapshot(snapshot);
086                }
087            }
088        }
089        for (MultiCollector collector : multiCollectors) {
090            List<String> prometheusNames = collector.getPrometheusNames();
091            boolean excluded = true; // the multi-collector is excluded unless at least one name matches
092            for (String prometheusName : prometheusNames) {
093                if (includedNames.test(prometheusName)) {
094                    excluded = false;
095                    break;
096                }
097            }
098            if (!excluded) {
099                for (MetricSnapshot snapshot : collector.collect(includedNames)) {
100                    if (snapshot != null) {
101                        result.metricSnapshot(snapshot);
102                    }
103                }
104            }
105        }
106        return result.build();
107    }
108
109}