Initial attempt at a scalar callback-based metric to give instantaneous snapshot gauges

3 files changed, 45 insertions, 2 deletions

diff --git a/synapse/metrics/__init__.py b/synapse/metrics/__init__.py
index d5c30bbe41..d7584fc0bc 100644
--- a/synapse/metrics/__init__.py
+++ b/synapse/metrics/__init__.py
@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .metric import CounterMetric, CacheCounterMetric
+from .metric import CounterMetric, CallbackMetric, CacheCounterMetric
 
 
 # We'll keep all the available metrics in a single toplevel dict, one shared
@@ -43,6 +43,15 @@ class Metrics(object):
 
         return metric
 
+    def register_callback(self, name, callback, *args, **kwargs):
+        full_name = "%s.%s" % (self.name_prefix, name)
+
+        metric = CallbackMetric(full_name, *args, callback=callback, **kwargs)
+
+        self._register(metric)
+
+        return metric
+
     def register_cachecounter(self, name, *args, **kwargs):
         full_name = "%s.%s" % (self.name_prefix, name)
 
diff --git a/synapse/metrics/metric.py b/synapse/metrics/metric.py
index 00b149f6f6..8a497fc154 100644
--- a/synapse/metrics/metric.py
+++ b/synapse/metrics/metric.py
@@ -62,6 +62,20 @@ class CounterMetric(BaseMetric):
                 for k in sorted(self.counts.keys())]
 
 
+class CallbackMetric(BaseMetric):
+    """A metric that returns the numeric value returned by a callback whenever
+    it is rendered. Typically this is used to implement gauges that yield the
+    size or other state of some in-memory object by actively querying it."""
+
+    def __init__(self, name, callback, keys=[]):
+        super(CallbackMetric, self).__init__(name, keys=keys)
+
+        self.callback = callback
+
+    def render(self):
+        # TODO(paul): work out something we can do with keys and vectors
+        return ["%s %d" % (self.name, self.callback())]
+
 class CacheCounterMetric(object):
     """A combination of two CounterMetrics, one to count cache hits and one to
     count misses.
diff --git a/tests/metrics/test_metric.py b/tests/metrics/test_metric.py
index 93e8e27e4f..b7bb375ce0 100644
--- a/tests/metrics/test_metric.py
+++ b/tests/metrics/test_metric.py
@@ -15,7 +15,9 @@
 
 from tests import unittest
 
-from synapse.metrics.metric import CounterMetric, CacheCounterMetric
+from synapse.metrics.metric import (
+    CounterMetric, CallbackMetric, CacheCounterMetric
+)
 
 
 class CounterMetricTestCase(unittest.TestCase):
@@ -61,6 +63,24 @@ class CounterMetricTestCase(unittest.TestCase):
         ])
 
 
+class CallbackMetricTestCase(unittest.TestCase):
+
+    def test_callback(self):
+        d = dict()
+
+        metric = CallbackMetric("size", lambda: len(d))
+
+        self.assertEquals(metric.render(), [
+            "size 0",
+        ])
+
+        d["key"] = "value"
+
+        self.assertEquals(metric.render(), [
+            "size 1",
+        ])
+
+
 class CacheCounterMetricTestCase(unittest.TestCase):
 
     def test_cachecounter(self):