From 44007e42bf97c3e46665e94c4eb453064205b31c Mon Sep 17 00:00:00 2001
From: Philip Jenvey <pjenvey@underboss.org>
Date: Thu, 30 Nov 2017 16:37:41 -0800
Subject: [PATCH] feat: make memusage's objgraph/dump_rpy_heap optional

Closes #1091
---
 autopush/memusage.py               | 15 +++++++++------
 autopush/tests/test_integration.py | 17 ++++++++++++++++-
 autopush/web/health.py             | 13 ++++++++++---
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/autopush/memusage.py b/autopush/memusage.py
index 00da5818..bc6fd7b4 100644
--- a/autopush/memusage.py
+++ b/autopush/memusage.py
@@ -7,13 +7,14 @@
 import tempfile
 import zlib
 from StringIO import StringIO
+from typing import Optional  # noqa
 
 from autopush.gcdump import Stat
 
 
-def memusage():
+def memusage(do_dump_rpy_heap=True, do_objgraph=True):
     """Returning a str of memory usage stats"""
-    # type() -> str
+    # type: (Optional[bool], Optional[bool]) -> str
     def trap_err(func, *args, **kwargs):
         try:
             return func(*args, **kwargs)
@@ -25,12 +26,14 @@ def trap_err(func, *args, **kwargs):
     rusage = trap_err(resource.getrusage, resource.RUSAGE_SELF)
     buf.writelines([repr(rusage), '\n\n'])
     trap_err(pmap_extended, buf)
-    # dump rpython's heap before objgraph potentially pollutes the
-    # heap with its heavy workload
-    trap_err(dump_rpy_heap, buf)
+    if do_dump_rpy_heap:
+        # dump rpython's heap before objgraph potentially pollutes the
+        # heap with its heavy workload
+        trap_err(dump_rpy_heap, buf)
     trap_err(get_stats_asmmemmgr, buf)
     buf.write('\n\n')
-    trap_err(objgraph.show_most_common_types, limit=0, file=buf)
+    if do_objgraph:
+        trap_err(objgraph.show_most_common_types, limit=0, file=buf)
     return buf.getvalue()
 
 
diff --git a/autopush/tests/test_integration.py b/autopush/tests/test_integration.py
index b819be2e..69dd0581 100644
--- a/autopush/tests/test_integration.py
+++ b/autopush/tests/test_integration.py
@@ -2109,7 +2109,7 @@ def test_memusage(self):
             "http://localhost:{}/_memusage".format(port),
         )
         assert response.code == 200
-        assert 'rusage' in body
+        assert 'ru_maxrss=' in body
         assert 'Logger' in body
         if find_executable('pmap'):
             assert 'RSS' in body or 'Rss' in body  # pmap -x or -XX/X output
@@ -2118,6 +2118,21 @@ def test_memusage(self):
             assert 'rpy_unicode' in body
             assert 'get_stats_asmmemmgr: (' in body
 
+    @inlineCallbacks
+    def test_memusage_options(self):
+        port = self.ep.conf.memusage_port
+        url = ("http://localhost:{}/_memusage?objgraph=false&"
+               "dump_rpy_heap=false").format(port)
+        response, body = yield _agent('GET', url)
+        assert response.code == 200
+        assert 'ru_maxrss=' in body
+        assert 'Logger' not in body
+        if find_executable('pmap'):
+            assert 'RSS' in body or 'Rss' in body  # pmap -x or -XX/X output
+        if hasattr(sys, 'pypy_version_info'):  # pragma: nocover
+            assert 'size: ' not in body
+            assert 'rpy_unicode' not in body
+
 
 @inlineCallbacks
 def _agent(method, url, contextFactory=None, headers=None, body=None):
diff --git a/autopush/web/health.py b/autopush/web/health.py
index d28265b2..dd120668 100644
--- a/autopush/web/health.py
+++ b/autopush/web/health.py
@@ -97,7 +97,7 @@ def get(self):
 
 
 class MemUsageHandler(BaseWebHandler):
-    """Spits out a tarball of some memory stats.
+    """Spits out some memory stats.
 
     Should be ran on its own port, not accessible externally.
 
@@ -110,11 +110,18 @@ def authenticate_peer_cert(self):
     def get(self):
         """HTTP Get
 
-        Returns that this node is alive, and the version.
+        Returns the memory stats.
 
         """
         from autopush.memusage import memusage
-        d = deferToThread(memusage)
+
+        def enabled(name):
+            return self.get_argument(name, u'true').lower() != u'false'
+        d = deferToThread(
+            memusage,
+            do_dump_rpy_heap=enabled('dump_rpy_heap'),
+            do_objgraph=enabled('objgraph')
+        )
         d.addCallback(self.write)
         d.addCallback(self.finish)
         d.addErrback(self._response_err)