From 05cc74d142af50f9022f7e770bb467dfe9757399 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Thu, 21 Sep 2017 21:19:24 +0530
Subject: [PATCH] core: Support profiling

---
 mitogen/core.py    | 38 ++++++++++++++++++++++++++++++++------
 mitogen/fakessh.py |  1 +
 mitogen/master.py  | 12 ++++++++++--
 3 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/mitogen/core.py b/mitogen/core.py
index eb6701b1..abce8f78 100644
--- a/mitogen/core.py
+++ b/mitogen/core.py
@@ -154,6 +154,27 @@ def enable_debug_logging():
     root.handlers.insert(0, handler)
 
 
+_profile_hook = lambda name, func, *args: func(*args)
+
+def enable_profiling():
+    global _profile_hook
+    import cProfile, pstats
+    def _profile_hook(name, func, *args):
+        profiler = cProfile.Profile()
+        profiler.enable()
+        try:
+            return func(*args)
+        finally:
+            profiler.create_stats()
+            fp = open('/tmp/mitogen.stats.%d.%s.log' % (os.getpid(), name), 'w')
+            try:
+                stats = pstats.Stats(profiler, stream=fp)
+                stats.sort_stats('cumulative')
+                stats.print_stats()
+            finally:
+                fp.close()
+
+
 class Message(object):
     dst_id = None
     src_id = None
@@ -952,8 +973,11 @@ class Broker(object):
         self._writers = set()
         self._waker = Waker(self)
         self.start_receive(self._waker)
-        self._thread = threading.Thread(target=self._broker_main,
-                                        name='mitogen-broker')
+        self._thread = threading.Thread(
+            target=_profile_hook,
+            args=('broker', self._broker_main),
+            name='mitogen-broker'
+        )
         self._thread.start()
 
     def defer(self, func, *args, **kwargs):
@@ -1075,7 +1099,9 @@ class ExternalContext(object):
     def _on_broker_shutdown(self):
         self.channel.close()
 
-    def _setup_master(self, parent_id, context_id, key, in_fd, out_fd):
+    def _setup_master(self, profiling, parent_id, context_id, key, in_fd, out_fd):
+        if profiling:
+            enable_profiling()
         self.broker = Broker()
         self.router = Router(self.broker)
         self.master = Context(self.router, 0, 'master')
@@ -1172,9 +1198,9 @@ class ExternalContext(object):
                     Message.pickled(e, dst_id=msg.src_id, handle=msg.reply_to)
                 )
 
-    def main(self, parent_id, context_id, key, debug, log_level,
+    def main(self, parent_id, context_id, key, debug, profiling, log_level,
              in_fd=100, out_fd=1, core_src_fd=101, setup_stdio=True):
-        self._setup_master(parent_id, context_id, key, in_fd, out_fd)
+        self._setup_master(profiling, parent_id, context_id, key, in_fd, out_fd)
         try:
             try:
                 self._setup_logging(debug, log_level)
@@ -1190,7 +1216,7 @@ class ExternalContext(object):
                           self.parent, context_id, os.getpid())
                 LOG.debug('Recovered sys.executable: %r', sys.executable)
 
-                self._dispatch_calls()
+                _profile_hook('main', self._dispatch_calls)
                 LOG.debug('ExternalContext.main() normal exit')
             except BaseException:
                 LOG.exception('ExternalContext.main() crashed')
diff --git a/mitogen/fakessh.py b/mitogen/fakessh.py
index ed53667c..785f8030 100644
--- a/mitogen/fakessh.py
+++ b/mitogen/fakessh.py
@@ -382,6 +382,7 @@ def run(dest, router, args, deadline=None, econtext=None):
                 context_id,                     # context_id
                 fakessh.key,                    # key
                 router.debug,                   # debug
+                router.profiling,               # profiling
                 logging.getLogger().level,      # log_level
                 sock2.fileno(),                 # in_fd
                 sock2.fileno(),                 # out_fd
diff --git a/mitogen/master.py b/mitogen/master.py
index 1b7b25d7..46cfa455 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -530,7 +530,11 @@ class Stream(mitogen.core.Stream):
     #: True to cause context to write verbose /tmp/mitogen.<pid>.log.
     debug = False
 
-    def construct(self, remote_name=None, python_path=None, debug=False, **kwargs):
+    #: True to cause context to write /tmp/mitogen.stats.<pid>.<thread>.log.
+    profiling = False
+
+    def construct(self, remote_name=None, python_path=None, debug=False,
+                  profiling=False, **kwargs):
         """Get the named context running on the local machine, creating it if
         it does not exist."""
         super(Stream, self).construct(**kwargs)
@@ -542,6 +546,7 @@ class Stream(mitogen.core.Stream):
             remote_name %= (getpass.getuser(), socket.gethostname(), os.getpid())
         self.remote_name = remote_name
         self.debug = debug
+        self.profiling = profiling
 
     def on_shutdown(self, broker):
         """Request the slave gracefully shut itself down."""
@@ -589,10 +594,11 @@ class Stream(mitogen.core.Stream):
     def get_preamble(self):
         source = inspect.getsource(mitogen.core)
         source += '\nExternalContext().main%r\n' % ((
-            mitogen.context_id,       # parent_id
+            mitogen.context_id,        # parent_id
             self.remote_id,            # context_id
             self.key,
             self.debug,
+            self.profiling,
             LOG.level or logging.getLogger().level or logging.INFO,
         ),)
 
@@ -748,6 +754,7 @@ class ChildIdAllocator(object):
 
 class Router(mitogen.core.Router):
     debug = False
+    profiling = False
 
     def __init__(self, *args, **kwargs):
         super(Router, self).__init__(*args, **kwargs)
@@ -804,6 +811,7 @@ class Router(mitogen.core.Router):
     def connect(self, method_name, name=None, **kwargs):
         klass = METHOD_NAMES[method_name]()
         kwargs.setdefault('debug', self.debug)
+        kwargs.setdefault('profiling', self.profiling)
 
         via = kwargs.pop('via', None)
         if via is not None: