Tornado stack_context live call data example

example-output.json

{
    "_debug": [
        {
            "kwargs": {},
            "time": 1.8904321193695068,
            "name": "get_impl:6cf3e692-2e98-45ae-bfd4-597c47851b0b",
            "calls": [
                {
                    "kwargs": {},
                    "time": 0.5965301990509033,
                    "name": "sleep_for:8eaeb1ff-c600-493c-8cb9-b5dfea9b9f52",
                    "calls": [],
                    "args": [
                        "0.5948681633672656"
                    ],
                    "end": 1428269617.930542,
                    "start": 1428269617.3340118
                },
                {
                    "kwargs": {},
                    "time": 1.8900480270385742,
                    "name": "nested_sleep_for_level_1:28f3f663-9eb1-4cd6-bc31-2f09fccf7ea9",
                    "calls": [
                        {
                            "kwargs": {},
                            "time": 0.9447321891784668,
                            "name": "nested_sleep_for_level_2:33934624-378c-4f4f-af3e-e7b6fa981ff7",
                            "calls": [],
                            "args": [
                                "0.9427484355324981"
                            ],
                            "end": 1428269619.2226226,
                            "start": 1428269618.2778904
                        },
                        {
                            "kwargs": {},
                            "time": 0.9445490837097168,
                            "name": "nested_sleep_for_level_2:c540412c-86fb-4849-b753-05ad2b8d0a3c",
                            "calls": [],
                            "args": [
                                "0.9427484355324981"
                            ],
                            "end": 1428269619.2234523,
                            "start": 1428269618.2789032
                        },
                        {
                            "kwargs": {},
                            "time": 0.944650411605835,
                            "name": "nested_sleep_for_level_2:1b0047d3-bc4e-407f-b56c-63a4f29811f8",
                            "calls": [],
                            "args": [
                                "0.9427484355324981"
                            ],
                            "end": 1428269619.2228742,
                            "start": 1428269618.2782238
                        }

                    ],
                    "args": [
                        "0.9427484355324981"
                    ],
                    "end": 1428269619.223697,
                    "start": 1428269617.333649
                }

            ],
            "args": [
                "<deliciouscake.handler.Handler object at 0x7f7550c3e438>",
                "baz"
            ],
            "end": 1428269619.2238536,
            "start": 1428269617.3334215
        }
    ],
    "foo": "baz"
}

handler.py

import functools

from random import random
from tornado import web, gen

from .wrapper import profiled_func, profiling_context


@profiled_func
@gen.coroutine
def nested_sleep_for_level_2(seconds):
    yield gen.sleep(seconds)
    return


@profiled_func
@gen.coroutine
def nested_sleep_for_level_1(seconds):
    yield gen.sleep(seconds)
    yield [nested_sleep_for_level_2(seconds), nested_sleep_for_level_2(seconds), nested_sleep_for_level_2(seconds)]
    return


@profiled_func
@gen.coroutine
def sleep_for(seconds):
    yield gen.sleep(seconds)
    return


class Handler(web.RequestHandler):
    @gen.coroutine
    def get(self, path):
        if path == 'profile':
            with profiling_context(functools.partial(self.get_impl, 'baz')) as ctx:
                result = yield ctx.result
            result['_debug'] = ctx.profiling_info
        else:
            result = yield self.get_impl('bar')

        self.finish(result)

    @profiled_func
    @gen.coroutine
    def get_impl(self, arg):
        yield [nested_sleep_for_level_1(random()), sleep_for(random())]
        return {'foo': arg}

main.py

from tornado import web, ioloop

from .handler import Handler


def main():
    app = web.Application([
        (r'/(.*)', Handler)
    ])
    app.listen(8000)
    ioloop.IOLoop.current().start()


if __name__ == '__main__':
    main()

wrapper.py

import contextlib
import threading
from uuid import uuid4
from time import time

import addict
import functools
from tornado.stack_context import StackContext


_STATE = threading.local()

_STATE.request_id = None

_STATE.stack = []
_STATE.profiling = addict.Dict()


def profiled_func(_impl):
    @functools.wraps(_impl)
    def _wrapped(*args, **kwargs):
        request_id = kwargs.pop('request_id', _STATE.request_id)

        if request_id:
            invocation_id = str(uuid4())
            with StackContext(functools.partial(profiling_gathering_context, request_id, invocation_id,
                                                _impl.__name__, *args, **kwargs)):
                return _impl(*args, **kwargs)
        else:
            return _impl(*args, **kwargs)

    return _wrapped


@contextlib.contextmanager
def profiling_gathering_context(request_id, invocation_id, func_name, *func_args, **func_kwargs):
    prev_request_id = _STATE.request_id
    _STATE.request_id = request_id
    _STATE.stack.append((request_id, invocation_id, func_name))

    curr_profile = _STATE.profiling[request_id]
    for k in _STATE.stack:
        curr_profile = curr_profile[k]

    start = time()
    curr_profile.start = min(curr_profile.get('start', start + 1), start)
    curr_profile.args = func_args
    curr_profile.kwargs = func_kwargs
    try:
        yield
    finally:
        end = time()
        curr_profile.end = max(curr_profile.get('end', end - 1), end)
        _STATE.request_id = prev_request_id
        _STATE.stack.pop()


def _format_profiling_info(profiling_info_dict):
    calls = []
    for k, v in profiling_info_dict.items():
        if not isinstance(k, str):
            req_id, invoc_id, func_name = k
            this_start = v.start
            this_end = v.end
            this_args = v.args
            this_kwargs = v.kwargs
            this_calls = _format_profiling_info(v)
            calls.append({
                'name': '{}:{}'.format(func_name, invoc_id),
                'start': this_start,
                'end': this_end,
                'time': this_end - this_start,
                'calls': this_calls,
                'args': [str(arg) for arg in this_args],
                'kwargs': {str(k): str(v) for k, v in this_kwargs},
            })
    return calls


class _ProfilingInfo(object):
    def __init__(self, result, profiling_info):
        self.result = result
        self._profiling_info = profiling_info

    @property
    def profiling_info(self):
        return _format_profiling_info(self._profiling_info)


@contextlib.contextmanager
def profiling_context(func):
    request_id = str(uuid4())

    try:
        result = func(request_id=request_id)
        profiling_info_ref = _STATE.profiling[request_id]
        yield _ProfilingInfo(result, profiling_info_ref)
    finally:
        del _STATE.profiling[request_id]