stack_context.py

import functools
import threading

from tornado import gen, httpserver, ioloop, stack_context, web


class ThreadRequestContext(object):
    """
    A context manager that saves some per-thread state globally.
    Intended for use with Tornado's StackContext.
    """
    
    _state = threading.local()
    _state.data = {}

    class __metaclass__(type):
        # property() doesn't work on classmethods,
        #  see http://stackoverflow.com/q/128573/1231454
        @property
        def data(cls):
            if not hasattr(cls._state, 'data'):
               return {}
            return cls._state.data

    def __init__(self, **data):
        self._data = data

    def __enter__(self):
        self._prev_data = self.__class__.data
        self.__class__._state.data = self._data

    def __exit__(self, *exc):
        self.__class__._state.data = self._prev_data
        del self._prev_data
        return False


class RequestContextHandler(web.RequestHandler):

    def _execute(self, transforms, *args, **kwargs):
        """
        This is an undocumented method on ``web.RequestHandler`` that is called
        on every request.

        The use of functools can be thought of just instantiating a
        ThreadRequestContext instance. StackContext requires that a callable
        that returns a ContextManager be passed to it, not an instance of a
        contenxt manger. (https://github.com/tornadoweb/tornado/blob/master/tornado/stack_context.py#L95)
        The result is an instance of ThreadRequestContext being added to
        tornado.stack_context._state (instance of tornado.stack_context.State).
        Since the ThreadRequestContext contains a thread local, accessing
        the ``data`` property on the class returns whatever is passed to it
        in ``context_data`` below.

        Honestly, since Tornado is single threaded, I don't totally understand
        how this works. I would expect ``_state.data`` to get clobbered on
        each request. I assume that the StackContext is somehow capturing any
        thread locals inside of it and storing outside of the thread. This
        also might be tapping into the underlying epoll or kqueue
        implementation.
        """
        context_data = {'headers': self.request.headers}

        with stack_context.StackContext(
            functools.partial(ThreadRequestContext, **context_data)):
            return super(RequestContextHandler, self)._execute(transforms, *args, **kwargs)


class SimpleHandler(RequestContextHandler):
    """
    Super simple request handler.
    """
    @gen.coroutine
    def get(self, *args, **kwargs):
        yield extract_thread_context()
        #self.finish(str(extract_thread_context()))
        self.finish('asdf')


@gen.coroutine
def extract_thread_context():
    """
    Extract all request header data from the StackContext.
    """
    print ThreadRequestContext.data
    yield gen.sleep(5)


if __name__ == '__main__':
    application = web.Application([
        (r'/', SimpleHandler),
    ])
    http_server = httpserver.HTTPServer(application)
    http_server.listen(8080)
    ioloop.IOLoop.current().start()