Commit 05cc74d1 authored by David Wilson's avatar David Wilson

core: Support profiling

parent b827ee1b
......@@ -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')
......
......@@ -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
......
......@@ -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:
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment