Merge pull request prometheus#110 from prometheus/desc

brian-brazil · web-flow · commit 102061849609 · 2016-11-16T11:31:02.000Z
Add describe and limiting of http expostion
diff --git a/README.md b/README.md
@@ -365,6 +365,19 @@ REGISTRY.register(CustomCollector())
 
 `SummaryMetricFamily` and `HistogramMetricFamily` work similarly.
 
+A collector may implement a `describe` method which returns metrics in the same
+format as `collect` (though you don't have to include the samples). This is
+used to predetermine the names of time series a `CollectorRegistry` exposes and
+thus to detect collisions and duplicate registrations.
+
+Usually custom collectors do not have to implement `describe`. If `describe` is
+not implemented and the CollectorRegistry was created with `auto_desribe=True`
+(which is the case for the default registry) then `collect` will be called at
+registration time instead of `describe`. If this could cause problems, either
+implement a proper `describe`, or if that's not practical have `describe`
+return an empty list.
+
+
 ## Multiprocess Mode (Gunicorn)
 
 **Experimental: This feature is new and has rough edges.**
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -38,29 +38,94 @@ class CollectorRegistry(object):
     Metric objects. The returned metrics should be consistent with the Prometheus
     exposition formats.
     '''
-    def __init__(self):
-        self._collectors = set()
+    def __init__(self, auto_describe=False):
+        self._collector_to_names = {}
+        self._names_to_collectors = {}
+        self._auto_describe = auto_describe
         self._lock = Lock()
 
     def register(self, collector):
         '''Add a collector to the registry.'''
         with self._lock:
-            self._collectors.add(collector)
+            names = self._get_names(collector)
+            for name in names:
+                if name in self._names_to_collectors:
+                    raise ValueError('Timeseries already present '
+                            'in CollectorRegistry: ' + name)
+            for name in names:
+                self._names_to_collectors[name] = collector
+            self._collector_to_names[collector] = names
 
     def unregister(self, collector):
         '''Remove a collector from the registry.'''
         with self._lock:
-            self._collectors.remove(collector)
+            for name in self._collector_to_names[collector]:
+                del self._names_to_collectors[name]
+            del self._collector_to_names[collector]
+
+    def _get_names(self, collector):
+        '''Get names of timeseries the collector produces.'''
+        desc_func = None
+        # If there's a describe function, use it.
+        try:
+            desc_func = collector.describe
+        except AttributeError:
+            pass
+        # Otherwise, if auto describe is enabled use the collect function.
+        if not desc_func and self._auto_describe:
+            desc_func = collector.collect
+
+        if not desc_func:
+            return []
+
+        result = []
+        type_suffixes = {
+            'summary': ['', '_sum', '_count'],
+            'histogram': ['_bucket', '_sum', '_count']
+        }
+        for metric in desc_func():
+            for suffix in type_suffixes.get(metric.type, ['']):
+                result.append(metric.name + suffix)
+        return result
 
     def collect(self):
         '''Yields metrics from the collectors in the registry.'''
         collectors = None
         with self._lock:
-            collectors = copy.copy(self._collectors)
+            collectors = copy.copy(self._collector_to_names)
         for collector in collectors:
             for metric in collector.collect():
                 yield metric
 
+    def restricted_registry(self, names):
+        '''Returns object that only collects some metrics.
+
+        Returns an object which upon collect() will return
+        only samples with the given names.
+
+        Intended usage is:
+            generate_latest(REGISTRY.restricted_registry(['a_timeseries']))
+
+        Experimental.'''
+        names = set(names)
+        collectors = set()
+        with self._lock:
+            for name in names:
+                if name in self._names_to_collectors:
+                    collectors.add(self._names_to_collectors[name])
+        metrics = []
+        for collector in collectors:
+            for metric in collector.collect():
+                samples = [s for s in metric.samples if s[0] in names]
+                if samples:
+                    m = Metric(metric.name, metric.documentation, metric.type)
+                    m.samples = samples
+                    metrics.append(m)
+        class RestrictedRegistry(object):
+            def collect(self):
+                return metrics
+        return RestrictedRegistry()
+
     def get_sample_value(self, name, labels=None):
         '''Returns the sample value, or None if not found.
 
@@ -75,7 +140,7 @@ def get_sample_value(self, name, labels=None):
         return None
 
 
-REGISTRY = CollectorRegistry()
+REGISTRY = CollectorRegistry(auto_describe=True)
 '''The default registry.'''
 
 _METRIC_TYPES = ('counter', 'gauge', 'summary', 'histogram', 'untyped')
@@ -476,6 +541,10 @@ def init(name, documentation, labelnames=(), namespace='', subsystem='', registr
         if not _METRIC_NAME_RE.match(full_name):
             raise ValueError('Invalid metric name: ' + full_name)
 
+        def describe():
+            return [Metric(full_name, documentation, cls._type)]
+        collector.describe = describe
+
         def collect():
             metric = Metric(full_name, documentation, cls._type)
             for suffix, labels, value in collector._samples():
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -15,13 +15,14 @@
     from BaseHTTPServer import HTTPServer
     from urllib2 import build_opener, Request, HTTPHandler
     from urllib import quote_plus
+    from urlparse import parse_qs, urlparse
 except ImportError:
     # Python 3
     unicode = str
     from http.server import BaseHTTPRequestHandler
     from http.server import HTTPServer
     from urllib.request import build_opener, Request, HTTPHandler
-    from urllib.parse import quote_plus
+    from urllib.parse import quote_plus, parse_qs, urlparse
 
 
 CONTENT_TYPE_LATEST = str('text/plain; version=0.0.4; charset=utf-8')
@@ -34,7 +35,11 @@ def prometheus_app(environ, start_response):
         status = str('200 OK')
         headers = [(str('Content-type'), CONTENT_TYPE_LATEST)]
         start_response(status, headers)
-        return [generate_latest(registry)]
+        params = parse_qs(environ['QUERY_STRING'])
+        r = registry
+        if 'name[]' in params:
+            r = r.restricted_registry(params['name[]'])
+        return [generate_latest(r)]
     return prometheus_app
 
 
@@ -73,7 +78,11 @@ def do_GET(self):
         self.send_response(200)
         self.send_header('Content-Type', CONTENT_TYPE_LATEST)
         self.end_headers()
-        self.wfile.write(generate_latest(core.REGISTRY))
+        registry = core.REGISTRY
+        params = parse_qs(urlparse(self.path).query)
+        if 'name[]' in params:
+            registry = registry.restricted_registry(params['name[]'])
+        self.wfile.write(generate_latest(registry))
 
     def log_message(self, format, *args):
         return
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -382,6 +382,64 @@ def test_bad_constructors(self):
         self.assertRaises(ValueError, HistogramMetricFamily, 'h', 'help', buckets={}, sum_value=1, labels=['a'])
         self.assertRaises(KeyError, HistogramMetricFamily, 'h', 'help', buckets={}, sum_value=1)
 
+class TestCollectorRegistry(unittest.TestCase):
+    def test_duplicate_metrics_raises(self):
+        registry = CollectorRegistry()
+        Counter('c', 'help', registry=registry)
+        self.assertRaises(ValueError, Counter, 'c', 'help', registry=registry)
+        self.assertRaises(ValueError, Gauge, 'c', 'help', registry=registry)
+
+        Gauge('g', 'help', registry=registry)
+        self.assertRaises(ValueError, Gauge, 'g', 'help', registry=registry)
+        self.assertRaises(ValueError, Counter, 'g', 'help', registry=registry)
+
+        Summary('s', 'help', registry=registry)
+        self.assertRaises(ValueError, Summary, 's', 'help', registry=registry)
+        # We don't currently expose quantiles, but let's prevent future
+        # clashes anyway.
+        self.assertRaises(ValueError, Gauge, 's', 'help', registry=registry)
+
+        Histogram('h', 'help', registry=registry)
+        self.assertRaises(ValueError, Histogram, 'h', 'help', registry=registry)
+        # Clashes aggaint various suffixes.
+        self.assertRaises(ValueError, Summary, 'h', 'help', registry=registry)
+        self.assertRaises(ValueError, Counter, 'h_count', 'help', registry=registry)
+        self.assertRaises(ValueError, Counter, 'h_sum', 'help', registry=registry)
+        self.assertRaises(ValueError, Counter, 'h_bucket', 'help', registry=registry)
+        # The name of the histogram itself isn't taken.
+        Counter('h', 'help', registry=registry)
+
+    def test_unregister_works(self):
+        registry = CollectorRegistry()
+        s = Summary('s', 'help', registry=registry)
+        self.assertRaises(ValueError, Counter, 's_count', 'help', registry=registry)
+        registry.unregister(s)
+        Counter('s_count', 'help', registry=registry)
+
+    def custom_collector(self, metric_family, registry):
+        class CustomCollector(object):
+            def collect(self):
+                return [metric_family]
+        registry.register(CustomCollector())
+
+    def test_autodescribe_disabled_by_default(self):
+        registry = CollectorRegistry()
+        self.custom_collector(CounterMetricFamily('c', 'help', value=1), registry)
+        self.custom_collector(CounterMetricFamily('c', 'help', value=1), registry)
+
+        registry = CollectorRegistry(auto_describe=True)
+        self.custom_collector(CounterMetricFamily('c', 'help', value=1), registry)
+        self.assertRaises(ValueError, self.custom_collector, CounterMetricFamily('c', 'help', value=1), registry)
+
+    def test_restricted_registry(self):
+        registry = CollectorRegistry()
+        Counter('c', 'help', registry=registry)
+        Summary('s', 'help', registry=registry).observe(7)
+
+        m = Metric('s', 'help', 'summary')
+        m.samples = [('s_sum', {}, 7)]
+        self.assertEquals([m], registry.restricted_registry(['s_sum']).collect())
+
 
 if __name__ == '__main__':
     unittest.main()
