2 files changed, 57 insertions, 1 deletions
diff --git a/synapse/metrics/__init__.py b/synapse/metrics/__init__.py
index bdd7292a30..76d5998d75 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,21 @@ class Metrics(object):
         return self._register(CacheMetric, *args, **kwargs)
 
 
+def register_memory_metrics(hs):
+    try:
+        import psutil
+        process = psutil.Process()
+        process.memory_info().rss
+    except (ImportError, AttributeError):
+        logger.warn(
+            "psutil is not installed or incorrect version."
+            " Disabling memory metrics."
+        )
+        return
+    metric = MemoryUsageMetric(hs, psutil)
+    all_metrics.append(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..e81af29895 100644
--- a/synapse/metrics/metric.py
+++ b/synapse/metrics/metric.py
@@ -153,3 +153,43 @@ 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, psutil):
+        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):
+        if not self.memory_snapshots:
+            return []
+
+        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,
+        ]