Merge pull request #2854 from lbryio/tracemalloc_api

Adds tracemalloc api for memory troubleshooting

lbry/extras/daemon/daemon.py

@@ -1,3 +1,4 @@
+import linecache
 import os
 import re
 import asyncio
@@ -8,6 +9,7 @@ import inspect
 import typing
 import random
 import hashlib
+import tracemalloc
 from urllib.parse import urlencode, quote
 from typing import Callable, Optional, List
 from binascii import hexlify, unhexlify
@@ -4531,6 +4533,85 @@ class Daemon(metaclass=JSONRPCServerType):
         result['node_id'] = hexlify(self.dht_node.protocol.node_id).decode()
         return result
 
+    TRACEMALLOC_DOC = """
+    Controls and queries tracemalloc memory tracing tools for troubleshooting.
+    """
+
+    def jsonrpc_tracemalloc_enable(self):  # pylint: disable=no-self-use
+        """
+        Enable tracemalloc memory tracing
+
+        Usage:
+            jsonrpc_tracemalloc_enable
+
+        Options:
+            None
+
+        Returns:
+            (bool) is it tracing?
+        """
+        tracemalloc.start()
+        return tracemalloc.is_tracing()
+
+    def jsonrpc_tracemalloc_disable(self):  # pylint: disable=no-self-use
+        """
+        Disable tracemalloc memory tracing
+
+        Usage:
+            jsonrpc_tracemalloc_disable
+
+        Options:
+            None
+
+        Returns:
+            (bool) is it tracing?
+        """
+        tracemalloc.stop()
+        return tracemalloc.is_tracing()
+
+    def jsonrpc_tracemalloc_top(self, items: int = 10):  # pylint: disable=no-self-use
+        """
+        Show most common objects, the place that created them and their size.
+
+        Usage:
+            jsonrpc_tracemalloc_top [(<items> | --items=<items>)]
+
+        Options:
+            --items=<items>               : (int) maximum items to return, from the most common
+
+        Returns:
+            (dict) dictionary containing most common objects in memory
+            {
+                "line": (str) filename and line number where it was created,
+                "code": (str) code that created it,
+                "size": (int) size in bytes, for each "memory block",
+                "count" (int) number of memory blocks
+            }
+        """
+        if not tracemalloc.is_tracing():
+            raise Exception("Enable tracemalloc first! See 'tracemalloc set' command.")
+        stats = tracemalloc.take_snapshot().filter_traces((
+            tracemalloc.Filter(False, "<frozen importlib._bootstrap>"),
+            tracemalloc.Filter(False, "<unknown>"),
+            # tracemalloc and linecache here use some memory, but thats not relevant
+            tracemalloc.Filter(False, tracemalloc.__file__),
+            tracemalloc.Filter(False, linecache.__file__),
+        )).statistics('lineno', True)
+        results = []
+        for stat in stats:
+            frame = stat.traceback[0]
+            filename = os.sep.join(frame.filename.split(os.sep)[-2:])
+            line = linecache.getline(frame.filename, frame.lineno).strip()
+            results.append({
+                "line": f"{filename}:{frame.lineno}",
+                "code": line,
+                "size": stat.size,
+                "count": stat.count
+            })
+            if len(results) == items:
+                break
+        return results
+
     COMMENT_DOC = """
     View, create and abandon comments.
     """

tests/integration/other/test_other_commands.py

@@ -38,3 +38,20 @@ class SettingsManagement(CommandTestCase):
         self.assertTrue(self.daemon.analytics_manager.enabled)
         self.assertTrue(loggly.enabled)
         self.daemon.jsonrpc_settings_set('share_usage_data', False)
+
+
+class TroubleshootingCommands(CommandTestCase):
+    async def test_tracemalloc_commands(self):
+        self.addCleanup(self.daemon.jsonrpc_tracemalloc_disable)
+        self.assertFalse(self.daemon.jsonrpc_tracemalloc_disable())
+        self.assertTrue(self.daemon.jsonrpc_tracemalloc_enable())
+
+        class WeirdObject():
+            pass
+        hold_em = [WeirdObject() for _ in range(500)]
+        top = self.daemon.jsonrpc_tracemalloc_top(1)
+        self.assertEqual(1, len(top))
+        self.assertEqual('hold_em = [WeirdObject() for _ in range(500)]', top[0]['code'])
+        self.assertTrue(top[0]['line'].startswith('other/test_other_commands.py:'))
+        self.assertGreaterEqual(top[0]['count'], 500)
+        self.assertGreater(top[0]['size'], 0)  # just matters that its a positive integer