From 66868119dc3c42c3cc6ea0b41ade81285ef1c9de Mon Sep 17 00:00:00 2001
From: Erik Johnston <erik@matrix.org>
Date: Wed, 20 Jul 2016 15:47:28 +0100
Subject: [PATCH] Add metrics for psutil derived memory usage

---
 synapse/app/homeserver.py      |  3 +++
 synapse/metrics/__init__.py    |  9 +++++++-
 synapse/metrics/metric.py      | 38 ++++++++++++++++++++++++++++++++++
 synapse/python_dependencies.py |  1 +
 4 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/synapse/app/homeserver.py b/synapse/app/homeserver.py
index 9c2dd3295..fe68ceb07 100755
--- a/synapse/app/homeserver.py
+++ b/synapse/app/homeserver.py
@@ -51,6 +51,7 @@ from synapse.api.urls import (
 from synapse.config.homeserver import HomeServerConfig
 from synapse.crypto import context_factory
 from synapse.util.logcontext import LoggingContext
+from synapse.metrics import register_memory_metrics
 from synapse.metrics.resource import MetricsResource, METRICS_PREFIX
 from synapse.replication.resource import ReplicationResource, REPLICATION_PREFIX
 from synapse.federation.transport.server import TransportLayerServer
@@ -335,6 +336,8 @@ def setup(config_options):
         hs.get_datastore().start_doing_background_updates()
         hs.get_replication_layer().start_get_pdu_cache()
 
+        register_memory_metrics(hs)
+
     reactor.callWhenRunning(start)
 
     return hs
diff --git a/synapse/metrics/__init__.py b/synapse/metrics/__init__.py
index bdd7292a3..cce3dba47 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 341043952..d100841a7 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,
+        ]
diff --git a/synapse/python_dependencies.py b/synapse/python_dependencies.py
index e024cec0a..799d35da5 100644
--- a/synapse/python_dependencies.py
+++ b/synapse/python_dependencies.py
@@ -36,6 +36,7 @@ REQUIREMENTS = {
     "blist": ["blist"],
     "pysaml2>=3.0.0,<4.0.0": ["saml2>=3.0.0,<4.0.0"],
     "pymacaroons-pynacl": ["pymacaroons"],
+    "psutil>=2.0.0": ["psutil>=2.0.0"],
 }
 CONDITIONAL_REQUIREMENTS = {
     "web_client": {