001package io.prometheus.client;
002
003import java.util.ArrayList;
004import java.util.Arrays;
005import java.util.Collections;
006import java.util.Enumeration;
007import java.util.HashMap;
008import java.util.HashSet;
009import java.util.Iterator;
010import java.util.List;
011import java.util.Map;
012import java.util.NoSuchElementException;
013import java.util.Set;
014
015/**
016 * A registry of Collectors.
017 * <p>
018 * The majority of users should use the {@link #defaultRegistry}, rather than instantiating their own.
019 * <p>
020 * Creating a registry other than the default is primarily useful for unittests, or
021 * pushing a subset of metrics to the <a href="https://github.com/prometheus/pushgateway">Pushgateway</a>
022 * from batch jobs.
023 */
024public class CollectorRegistry {
025  /**
026   * The default registry.
027   */
028  public static final CollectorRegistry defaultRegistry = new CollectorRegistry(true);
029
030  private final Object namesCollectorsLock = new Object();
031  private final Map<Collector, List<String>> collectorsToNames = new HashMap<Collector, List<String>>();
032  private final Map<String, Collector> namesToCollectors = new HashMap<String, Collector>();
033
034  private final boolean autoDescribe;
035
036  public CollectorRegistry() {
037    this(false);
038  }
039
040  public CollectorRegistry(boolean autoDescribe) {
041    this.autoDescribe = autoDescribe;
042  }
043
044  /**
045   * Register a Collector.
046   * <p>
047   * A collector can be registered to multiple CollectorRegistries.
048   */
049  public void register(Collector m) {
050    List<String> names = collectorNames(m);
051    assertNoDuplicateNames(m, names);
052    synchronized (namesCollectorsLock) {
053      for (String name : names) {
054        if (namesToCollectors.containsKey(name)) {
055          throw new IllegalArgumentException("Failed to register Collector of type " + m.getClass().getSimpleName()
056                  + ": " + name + " is already in use by another Collector of type "
057                  + namesToCollectors.get(name).getClass().getSimpleName());
058        }
059      }
060      for (String name : names) {
061        namesToCollectors.put(name, m);
062      }
063      collectorsToNames.put(m, names);
064    }
065  }
066
067  private void assertNoDuplicateNames(Collector m, List<String> names) {
068    Set<String> uniqueNames = new HashSet<String>();
069    for (String name : names) {
070      if (!uniqueNames.add(name)) {
071        throw new IllegalArgumentException("Failed to register Collector of type " + m.getClass().getSimpleName()
072                + ": The Collector exposes the same name multiple times: " + name);
073      }
074    }
075  }
076
077  /**
078   * Unregister a Collector.
079   */
080  public void unregister(Collector m) {
081    synchronized (namesCollectorsLock) {
082      List<String> names = collectorsToNames.remove(m);
083      for (String name : names) {
084        namesToCollectors.remove(name);
085      }
086    }
087  }
088
089  /**
090   * Unregister all Collectors.
091   */
092  public void clear() {
093    synchronized (namesCollectorsLock) {
094      collectorsToNames.clear();
095      namesToCollectors.clear();
096    }
097  }
098
099  /**
100   * A snapshot of the current collectors.
101   */
102  private Set<Collector> collectors() {
103    synchronized (namesCollectorsLock) {
104      return new HashSet<Collector>(collectorsToNames.keySet());
105    }
106  }
107
108  private List<String> collectorNames(Collector m) {
109    List<Collector.MetricFamilySamples> mfs;
110    if (m instanceof Collector.Describable) {
111      mfs = ((Collector.Describable) m).describe();
112    } else if (autoDescribe) {
113      mfs = m.collect();
114    } else {
115      mfs = Collections.emptyList();
116    }
117
118    List<String> names = new ArrayList<String>();
119    for (Collector.MetricFamilySamples family : mfs) {
120      names.addAll(Arrays.asList(family.getNames()));
121    }
122    return names;
123  }
124
125  /**
126   * Enumeration of metrics of all registered collectors.
127   */
128  public Enumeration<Collector.MetricFamilySamples> metricFamilySamples() {
129    return new MetricFamilySamplesEnumeration();
130  }
131
132  /**
133   * Enumeration of metrics matching the specified names.
134   * <p>
135   * Note that the provided set of names will be matched against the time series
136   * name and not the metric name. For instance, to retrieve all samples from a
137   * histogram, you must include the '_count', '_sum' and '_bucket' names.
138   */
139  public Enumeration<Collector.MetricFamilySamples> filteredMetricFamilySamples(Set<String> includedNames) {
140    return new MetricFamilySamplesEnumeration(new SampleNameFilter.Builder().nameMustBeEqualTo(includedNames).build());
141  }
142
143  /**
144   * Enumeration of metrics where {@code sampleNameFilter.test(name)} returns {@code true} for each {@code name} in
145   * {@link Collector.MetricFamilySamples#getNames()}.
146   * @param sampleNameFilter may be {@code null}, indicating that the enumeration should contain all metrics.
147   */
148  public Enumeration<Collector.MetricFamilySamples> filteredMetricFamilySamples(Predicate<String> sampleNameFilter) {
149    return new MetricFamilySamplesEnumeration(sampleNameFilter);
150  }
151
152  class MetricFamilySamplesEnumeration implements Enumeration<Collector.MetricFamilySamples> {
153
154    private final Iterator<Collector> collectorIter;
155    private Iterator<Collector.MetricFamilySamples> metricFamilySamples;
156    private Collector.MetricFamilySamples next;
157    private final Predicate<String> sampleNameFilter;
158
159    MetricFamilySamplesEnumeration(Predicate<String> sampleNameFilter) {
160      this.sampleNameFilter = sampleNameFilter;
161      this.collectorIter = filteredCollectorIterator();
162      findNextElement();
163    }
164
165    private Iterator<Collector> filteredCollectorIterator() {
166      if (sampleNameFilter == null) {
167        return collectors().iterator();
168      } else {
169        HashSet<Collector> collectors = new HashSet<Collector>();
170        synchronized (namesCollectorsLock) {
171          for (Map.Entry<String, Collector> entry : namesToCollectors.entrySet()) {
172            // Note that namesToCollectors contains keys for all combinations of suffixes (_total, _info, etc.).
173            if (sampleNameFilter.test(entry.getKey())) {
174              collectors.add(entry.getValue());
175            }
176          }
177        }
178        return collectors.iterator();
179      }
180    }
181
182    MetricFamilySamplesEnumeration() {
183      this(null);
184    }
185
186    private void findNextElement() {
187      next = null;
188
189      while (metricFamilySamples != null && metricFamilySamples.hasNext()) {
190        next = metricFamilySamples.next().filter(sampleNameFilter);
191        if (next != null) {
192          return;
193        }
194      }
195
196      while (collectorIter.hasNext()) {
197        metricFamilySamples = collectorIter.next().collect(sampleNameFilter).iterator();
198        while (metricFamilySamples.hasNext()) {
199          next = metricFamilySamples.next().filter(sampleNameFilter);
200          if (next != null) {
201            return;
202          }
203        }
204      }
205    }
206
207    public Collector.MetricFamilySamples nextElement() {
208      Collector.MetricFamilySamples current = next;
209      if (current == null) {
210        throw new NoSuchElementException();
211      }
212      findNextElement();
213      return current;
214    }
215
216    public boolean hasMoreElements() {
217      return next != null;
218    }
219  }
220
221  /**
222   * Returns the given value, or null if it doesn't exist.
223   * <p>
224   * This is inefficient, and intended only for use in unittests.
225   */
226  public Double getSampleValue(String name) {
227    return getSampleValue(name, new String[]{}, new String[]{});
228  }
229
230  /**
231   * Returns the given value, or null if it doesn't exist.
232   * <p>
233   * This is inefficient, and intended only for use in unittests.
234   */
235  public Double getSampleValue(String name, String[] labelNames, String[] labelValues) {
236    for (Collector.MetricFamilySamples metricFamilySamples : Collections.list(metricFamilySamples())) {
237      for (Collector.MetricFamilySamples.Sample sample : metricFamilySamples.samples) {
238        if (sample.name.equals(name)
239                && Arrays.equals(sample.labelNames.toArray(), labelNames)
240                && Arrays.equals(sample.labelValues.toArray(), labelValues)) {
241          return sample.value;
242        }
243      }
244    }
245    return null;
246  }
247
248}