Add support for Enum/StateSet

brian-brazil · brian-brazil · commit 1aa8405786ed · 2018-09-07T17:17:03.000+01:00
Signed-off-by: Brian Brazil &lt;brian.brazil@robustperception.io&gt;
diff --git a/README.md b/README.md
@@ -169,7 +169,6 @@ with h.time():
   pass
 ```
 
-
 ### Info
 
 Info tracks key-value information, usually about a whole target.
@@ -180,6 +179,17 @@ i = Info('my_build_version', 'Description of info')
 i.info({'version': '1.2.3', 'buildhost': 'foo@bar'})
 ```
 
+### Enum
+
+Enum tracks which of a set of states something is currently in.
+
+```python
+from prometheus_client import Enum
+e = Enum('my_task_state', 'Description of enum',
+        states=['starting', 'running', 'stopped'])
+e.state('running')
+```
+
 ### Labels
 
 All metrics can have labels, allowing grouping of related time series.
@@ -425,7 +435,7 @@ This comes with a number of limitations:
 
 - Registries can not be used as normal, all instantiated metrics are exported
 - Custom collectors do not work (e.g. cpu and memory metrics)
-- Info metrics do not work
+- Info and Enum metrics do not work
 - The pushgateway cannot be used
 - Gauges cannot use the `pid` label
 
diff --git a/prometheus_client/__init__.py b/prometheus_client/__init__.py
@@ -5,7 +5,7 @@
 from . import process_collector
 from . import platform_collector
 
-__all__ = ['Counter', 'Gauge', 'Summary', 'Histogram', 'Info']
+__all__ = ['Counter', 'Gauge', 'Summary', 'Histogram', 'Info', 'Enum']
 
 CollectorRegistry = core.CollectorRegistry
 REGISTRY = core.REGISTRY
@@ -15,6 +15,7 @@
 Summary = core.Summary
 Histogram = core.Histogram
 Info = core.Info
+Enum = core.Enum
 
 CONTENT_TYPE_LATEST = exposition.CONTENT_TYPE_LATEST
 generate_latest = exposition.generate_latest
diff --git a/prometheus_client/core.py b/prometheus_client/core.py
@@ -151,7 +151,8 @@ def get_sample_value(self, name, labels=None):
 REGISTRY = CollectorRegistry(auto_describe=True)
 '''The default registry.'''
 
-_METRIC_TYPES = ('counter', 'gauge', 'summary', 'histogram', 'untyped', 'info')
+_METRIC_TYPES = ('counter', 'gauge', 'summary', 'histogram', 
+        'untyped', 'info', 'stateset')
 
 
 class Metric(object):
@@ -354,6 +355,35 @@ def add_metric(self, labels, value):
             dict(dict(zip(self._labelnames, labels)), **value), 1))
 
 
+class StateSetMetricFamily(Metric):
+    '''A single stateset and its samples.
+
+    For use by custom collectors.
+    '''
+    def __init__(self, name, documentation, value=None, labels=None):
+        Metric.__init__(self, name, documentation, 'stateset')
+        if labels is not None and value is not None:
+            raise ValueError('Can only specify at most one of value and labels.')
+        if labels is None:
+            labels = []
+        self._labelnames = tuple(labels)
+        if value is not None:
+            self.add_metric([], value)
+
+    def add_metric(self, labels, value):
+        '''Add a metric to the metric family.
+
+        Args:
+          labels: A list of label values
+          value: A dict of string state names to booleans
+        '''
+        labels = tuple(labels)
+        for state, enabled in value.items():
+            v = (1 if enabled else 0)
+            self.samples.append((self.name,
+                dict(zip(self._labelnames + (self.name,), labels + (state,))), v))
+
+
 class _MutexValue(object):
     '''A float protected by a mutex.'''
 
@@ -983,19 +1013,18 @@ class Info(object):
      Example usage:
         from prometheus_client import Info
 
-        g = Info('my_build_info', 'Description of info')
-        g.info({'version': '1.2.3', 'buildhost': 'foo@bar'})
+        i = Info('my_build', 'Description of info')
+        i.info({'version': '1.2.3', 'buildhost': 'foo@bar'})
 
      Info metrics do not work in multiprocess mode.
     '''
     _type = 'info'
     _reserved_labelnames = []
-    _labelnames = []
-    _value = {}
 
     def __init__(self, name, labelnames, labelvalues):
         self._labelnames = set(labelnames)
         self._lock = Lock()
+        self._value = {}
 
     def info(self, val):
         '''Set info metric.'''
@@ -1011,6 +1040,44 @@ def _samples(self):
             return (('_info', self._value, 1.0,), )
 
 
+@_MetricWrapper
+class Enum(object):
+    '''Enum metric, which of a set of states is true.
+
+     Example usage:
+        from prometheus_client import Enum
+
+        e = Enum('task_state', 'Description of enum', 
+          states=['starting', 'running', 'stopped'])
+        e.state('running')
+
+     The first listed state will be the default.
+     Enum metrics do not work in multiprocess mode.
+    '''
+    _type = 'stateset'
+    _reserved_labelnames = []
+
+    def __init__(self, name, labelnames, labelvalues, states=None):
+        if name in labelnames:
+            raise ValueError('Overlapping labels for Enum metric: %s' % (name,))
+        if not states:
+            raise ValueError('No states provided for Enum metric: %s' % (name,))
+        self._name = name
+        self._states = states
+        self._value = 0
+        self._lock = Lock()
+
+    def state(self, state):
+        '''Set enum metric state.'''
+        with self._lock:
+            self._value = self._states.index(state)
+
+    def _samples(self):
+        with self._lock:
+            return [('', {self._name: s}, 1 if i == self._value else 0,)
+                       for i, s in enumerate(self._states)]
+
+
 class _ExceptionCounter(object):
     def __init__(self, counter, exception):
         self._counter = counter
diff --git a/prometheus_client/exposition.py b/prometheus_client/exposition.py
@@ -76,6 +76,8 @@ def generate_latest(registry=core.REGISTRY):
         elif mtype == 'info':
             mname = mname + '_info'
             mtype = 'gauge'
+        elif mtype == 'stateset':
+            mtype = 'gauge'
         output.append('# HELP {0} {1}'.format(
             mname, metric.documentation.replace('\\', r'\\').replace('\n', r'\n')))
         output.append('\n# TYPE {0} {1}\n'.format(mname, mtype))
diff --git a/tests/test_core.py b/tests/test_core.py
@@ -20,6 +20,8 @@
     HistogramMetricFamily,
     Info,
     InfoMetricFamily,
+    Enum,
+    StateSetMetricFamily,
     Metric,
     Summary,
     SummaryMetricFamily,
@@ -289,7 +291,7 @@ def test_setting_buckets(self):
         self.assertRaises(ValueError, Histogram, 'h', 'help', registry=None, buckets=[3, 1])
 
     def test_labels(self):
-        self.assertRaises(ValueError, Histogram, 'h2', 'help', registry=None, labelnames=['le'])
+        self.assertRaises(ValueError, Histogram, 'h', 'help', registry=None, labelnames=['le'])
 
         self.labels.labels('a').observe(2)
         self.assertEqual(0, self.registry.get_sample_value('hl_bucket', {'le': '1.0', 'l': 'a'}))
@@ -363,6 +365,35 @@ def test_labels(self):
         self.assertEqual(1, self.registry.get_sample_value('il_info', {'l': 'a', 'foo': 'bar'}))
 
 
+class TestEnum(unittest.TestCase):
+    def setUp(self):
+        self.registry = CollectorRegistry()
+        self.enum = Enum('e', 'help', states=['a', 'b', 'c'], registry=self.registry)
+        self.labels = Enum('el', 'help', ['l'], states=['a', 'b', 'c'], registry=self.registry)
+
+    def test_enum(self):
+        self.assertEqual(1, self.registry.get_sample_value('e', {'e': 'a'}))
+        self.assertEqual(0, self.registry.get_sample_value('e', {'e': 'b'}))
+        self.assertEqual(0, self.registry.get_sample_value('e', {'e': 'c'}))
+
+        self.enum.state('b')
+        self.assertEqual(0, self.registry.get_sample_value('e', {'e': 'a'}))
+        self.assertEqual(1, self.registry.get_sample_value('e', {'e': 'b'}))
+        self.assertEqual(0, self.registry.get_sample_value('e', {'e': 'c'}))
+
+        self.assertRaises(ValueError, self.enum.state, 'd')
+        self.assertRaises(ValueError, Enum, 'e', 'help', registry=None)
+
+    def test_labels(self):
+        self.labels.labels('a').state('c')
+        self.assertEqual(0, self.registry.get_sample_value('el', {'l': 'a', 'el': 'a'}))
+        self.assertEqual(0, self.registry.get_sample_value('el', {'l': 'a', 'el': 'b'}))
+        self.assertEqual(1, self.registry.get_sample_value('el', {'l': 'a', 'el': 'c'}))
+
+        e = Enum('e', 'help', registry=None, labelnames=['e'])
+        self.assertRaises(ValueError, e.labels, '')
+
+
 class TestMetricWrapper(unittest.TestCase):
     def setUp(self):
         self.registry = CollectorRegistry()
@@ -524,6 +555,18 @@ def test_info_labels(self):
         self.custom_collector(cmf)
         self.assertEqual(1, self.registry.get_sample_value('i_info', {'a': 'b', 'c': 'd'}))
 
+    def test_stateset(self):
+        self.custom_collector(StateSetMetricFamily('s', 'help', value={'a': True, 'b': True,}))
+        self.assertEqual(1, self.registry.get_sample_value('s', {'s': 'a'}))
+        self.assertEqual(1, self.registry.get_sample_value('s', {'s': 'b'}))
+
+    def test_stateset_labels(self):
+        cmf = StateSetMetricFamily('s', 'help', labels=['foo'])
+        cmf.add_metric(['bar'], {'a': False, 'b': False,})
+        self.custom_collector(cmf)
+        self.assertEqual(0, self.registry.get_sample_value('s', {'foo': 'bar', 's': 'a'}))
+        self.assertEqual(0, self.registry.get_sample_value('s', {'foo': 'bar', 's': 'b'}))
+
     def test_bad_constructors(self):
         self.assertRaises(ValueError, UntypedMetricFamily, 'u', 'help', value=1, labels=[])
         self.assertRaises(ValueError, UntypedMetricFamily, 'u', 'help', value=1, labels=['a'])
@@ -550,6 +593,9 @@ def test_bad_constructors(self):
         self.assertRaises(ValueError, InfoMetricFamily, 'i', 'help', value={}, labels=[])
         self.assertRaises(ValueError, InfoMetricFamily, 'i', 'help', value={}, labels=['a'])
 
+        self.assertRaises(ValueError, StateSetMetricFamily, 's', 'help', value={'a': True}, labels=[])
+        self.assertRaises(ValueError, StateSetMetricFamily, 's', 'help', value={'a': True}, labels=['a'])
+
     def test_labelnames(self):
         cmf = UntypedMetricFamily('u', 'help', labels=iter(['a']))
         self.assertEqual(('a',), cmf._labelnames)
diff --git a/tests/test_exposition.py b/tests/test_exposition.py
@@ -9,7 +9,7 @@
 else:
     import unittest
 
-from prometheus_client import Gauge, Counter, Summary, Histogram, Info, Metric
+from prometheus_client import Gauge, Counter, Summary, Histogram, Info, Enum, Metric
 from prometheus_client import CollectorRegistry, generate_latest
 from prometheus_client import push_to_gateway, pushadd_to_gateway, delete_from_gateway
 from prometheus_client import CONTENT_TYPE_LATEST, instance_ip_grouping_key
@@ -77,6 +77,11 @@ def test_info(self):
         i.labels('c', 'd').info({'foo': 'bar'})
         self.assertEqual(b'# HELP ii_info A info\n# TYPE ii_info gauge\nii_info{a="c",b="d",foo="bar"} 1.0\n', generate_latest(self.registry))
 
+    def test_enum(self):
+        i = Enum('ee', 'An enum', ['a', 'b'], registry=self.registry, states=['foo', 'bar'])
+        i.labels('c', 'd').state('bar')
+        self.assertEqual(b'# HELP ee An enum\n# TYPE ee gauge\nee{a="c",b="d",ee="foo"} 0.0\nee{a="c",b="d",ee="bar"} 1.0\n', generate_latest(self.registry))
+
     def test_unicode(self):
         c = Counter('cc', '\u4500', ['l'], registry=self.registry)
         c.labels('\u4500').inc()
