2 files changed, 46 insertions, 1 deletions
diff --git a/synapse/metrics/__init__.py b/synapse/metrics/__init__.py
index bdd7292a30..cce3dba47c 100644
--- a/synapse/metrics/__init__.py
+++ b/synapse/metrics/__init__.py
@@ -27,7 +27,8 @@ import gc
 from twisted.internet import reactor
 
 from .metric import (
-    CounterMetric, CallbackMetric, DistributionMetric, CacheMetric
+    CounterMetric, CallbackMetric, DistributionMetric, CacheMetric,
+    MemoryUsageMetric,
 )
 
 
@@ -66,6 +67,12 @@ class Metrics(object):
         return self._register(CacheMetric, *args, **kwargs)
 
 
+def register_memory_metrics(hs):
+    metric = MemoryUsageMetric(hs)
+    all_metrics.append(metric)
+    return metric
+
+
 def get_metrics_for(pkg_name):
     """ Returns a Metrics instance for conveniently creating metrics
     namespaced with the given name prefix. """
diff --git a/synapse/metrics/metric.py b/synapse/metrics/metric.py
index 341043952a..d100841a7f 100644
--- a/synapse/metrics/metric.py
+++ b/synapse/metrics/metric.py
@@ -16,6 +16,8 @@
 
 from itertools import chain
 
+import psutil
+
 
 # TODO(paul): I can't believe Python doesn't have one of these
 def map_concat(func, items):
@@ -153,3 +155,39 @@ class CacheMetric(object):
             """%s:total{name="%s"} %d""" % (self.name, self.cache_name, total),
             """%s:size{name="%s"} %d""" % (self.name, self.cache_name, size),
         ]
+
+
+class MemoryUsageMetric(object):
+    """Keeps track of the current memory usage, using psutil.
+
+    The class will keep the current min/max/sum/counts of rss over the last
+    WINDOW_SIZE_SEC, by polling UPDATE_HZ times per second
+    """
+
+    UPDATE_HZ = 2  # number of times to get memory per second
+    WINDOW_SIZE_SEC = 30  # the size of the window in seconds
+
+    def __init__(self, hs):
+        clock = hs.get_clock()
+        self.memory_snapshots = []
+        self.process = psutil.Process()
+
+        clock.looping_call(self._update_curr_values, 1000 / self.UPDATE_HZ)
+
+    def _update_curr_values(self):
+        max_size = self.UPDATE_HZ * self.WINDOW_SIZE_SEC
+        self.memory_snapshots.append(self.process.memory_info().rss)
+        self.memory_snapshots[:] = self.memory_snapshots[-max_size:]
+
+    def render(self):
+        max_rss = max(self.memory_snapshots)
+        min_rss = min(self.memory_snapshots)
+        sum_rss = sum(self.memory_snapshots)
+        len_rss = len(self.memory_snapshots)
+
+        return [
+            "process_psutil_rss:max %d" % max_rss,
+            "process_psutil_rss:min %d" % min_rss,
+            "process_psutil_rss:total %d" % sum_rss,
+            "process_psutil_rss:count %d" % len_rss,
+        ]