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
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    synchronized (collectorsToNames) {
052      for (String name : names) {
053        if (namesToCollectors.containsKey(name)) {
054          throw new IllegalArgumentException("Collector already registered that provides name: " + name);
055        }
056      }
057      for (String name : names) {
058        namesToCollectors.put(name, m);
059      }
060      collectorsToNames.put(m, names);
061    }
062  }
063
064  /**
065   * Unregister a Collector.
066   */
067  public void unregister(Collector m) {
068    synchronized (collectorsToNames) {
069      List<String> names = collectorsToNames.remove(m);
070      for (String name : names) {
071        namesToCollectors.remove(name);
072      }
073    }
074  }
075
076  /**
077   * Unregister all Collectors.
078   */
079  public void clear() {
080    synchronized (collectorsToNames) {
081      collectorsToNames.clear();
082      namesToCollectors.clear();
083    }
084  }
085
086  /**
087   * A snapshot of the current collectors.
088   */
089  private Set<Collector> collectors() {
090    synchronized (collectorsToNames) {
091      return new HashSet<Collector>(collectorsToNames.keySet());
092    }
093  }
094
095  private List<String> collectorNames(Collector m) {
096    List<Collector.MetricFamilySamples> mfs;
097    if (m instanceof Collector.Describable) {
098      mfs = ((Collector.Describable) m).describe();
099    } else if (autoDescribe) {
100      mfs = m.collect();
101    } else {
102      mfs = Collections.emptyList();
103    }
104
105    List<String> names = new ArrayList<String>();
106    for (Collector.MetricFamilySamples family : mfs) {
107      switch (family.type) {
108        case SUMMARY:
109          names.add(family.name + "_count");
110          names.add(family.name + "_sum");
111          names.add(family.name);
112          break;
113        case HISTOGRAM:
114          names.add(family.name + "_count");
115          names.add(family.name + "_sum");
116          names.add(family.name + "_bucket");
117          names.add(family.name);
118          break;
119        default:
120          names.add(family.name);
121      }
122    }
123    return names;
124  }
125
126  /**
127   * Enumeration of metrics of all registered collectors.
128   */
129  public Enumeration<Collector.MetricFamilySamples> metricFamilySamples() {
130    return new MetricFamilySamplesEnumeration();
131  }
132
133  public Enumeration<Collector.MetricFamilySamples> filteredMetricFamilySamples(Set<String> includedNames) {
134    return new MetricFamilySamplesEnumeration(includedNames);
135  }
136
137  class MetricFamilySamplesEnumeration implements Enumeration<Collector.MetricFamilySamples> {
138
139    private final Iterator<Collector> collectorIter;
140    private Iterator<Collector.MetricFamilySamples> metricFamilySamples;
141    private Collector.MetricFamilySamples next;
142    private Set<String> includedNames;
143
144    MetricFamilySamplesEnumeration(Set<String> includedNames) {
145      this.includedNames = includedNames;
146      collectorIter = includedCollectorIterator(includedNames);
147      findNextElement();
148    }
149
150    private Iterator<Collector> includedCollectorIterator(Set<String> includedNames) {
151      if (includedNames.isEmpty()) {
152        return collectors().iterator();
153      } else {
154        HashSet<Collector> collectors = new HashSet<Collector>();
155        synchronized (namesToCollectors) {
156          for (Map.Entry<String, Collector> entry : namesToCollectors.entrySet()) {
157            if (includedNames.contains(entry.getKey())) {
158              collectors.add(entry.getValue());
159            }
160          }
161        }
162
163        return collectors.iterator();
164      }
165    }
166
167    MetricFamilySamplesEnumeration() {
168      this(Collections.<String>emptySet());
169    }
170
171    private void findNextElement() {
172      next = null;
173
174      while (metricFamilySamples != null && metricFamilySamples.hasNext()) {
175        next = filter(metricFamilySamples.next());
176        if (next != null) {
177          return;
178        }
179      }
180
181      if (next == null) {
182        while (collectorIter.hasNext()) {
183          metricFamilySamples = collectorIter.next().collect().iterator();
184          while (metricFamilySamples.hasNext()) {
185            next = filter(metricFamilySamples.next());
186            if (next != null) {
187              return;
188            }
189          }
190        }
191      }
192    }
193
194    private Collector.MetricFamilySamples filter(Collector.MetricFamilySamples next) {
195      if (includedNames.isEmpty()) {
196        return next;
197      } else {
198        Iterator<Collector.MetricFamilySamples.Sample> it = next.samples.iterator();
199        while (it.hasNext()) {
200            if (!includedNames.contains(it.next().name)) {
201                it.remove();
202            }
203        }
204        if (next.samples.size() == 0) {
205          return null;
206        }
207        return next;
208      }
209    }
210
211    public Collector.MetricFamilySamples nextElement() {
212      Collector.MetricFamilySamples current = next;
213      if (current == null) {
214        throw new NoSuchElementException();
215      }
216      findNextElement();
217      return current;
218    }
219
220    public boolean hasMoreElements() {
221      return next != null;
222    }
223  }
224
225  /**
226   * Returns the given value, or null if it doesn't exist.
227   * <p>
228   * This is inefficient, and intended only for use in unittests.
229   */
230  public Double getSampleValue(String name) {
231    return getSampleValue(name, new String[]{}, new String[]{});
232  }
233
234  /**
235   * Returns the given value, or null if it doesn't exist.
236   * <p>
237   * This is inefficient, and intended only for use in unittests.
238   */
239  public Double getSampleValue(String name, String[] labelNames, String[] labelValues) {
240    for (Collector.MetricFamilySamples metricFamilySamples : Collections.list(metricFamilySamples())) {
241      for (Collector.MetricFamilySamples.Sample sample : metricFamilySamples.samples) {
242        if (sample.name.equals(name)
243                && Arrays.equals(sample.labelNames.toArray(), labelNames)
244                && Arrays.equals(sample.labelValues.toArray(), labelValues)) {
245          return sample.value;
246        }
247      }
248    }
249    return null;
250  }
251
252}