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..log. debug = False - def construct(self, remote_name=None, python_path=None, debug=False, **kwargs): + #: True to cause context to write /tmp/mitogen.stats...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: