Tornado monitoring wrappers (proof of concept)

demo.py

from tornado import gen
from tornado.httpclient import AsyncHTTPClient
from tornado.httpserver import HTTPServer
from tornado.ioloop import IOLoop
from tornado.options import define, options, parse_command_line
from tornado.web import RequestHandler, Application

define('port', default=8888)
define('trace', default=False)

class DelayHandler(RequestHandler):
    async def get(self):
        await gen.sleep(int(self.get_argument('ms')) / 1000.0)
        self.write('ok')


class MyHandler(RequestHandler):
    async def get(self):
        client = AsyncHTTPClient()
        await client.fetch('http://localhost:%d/delay?ms=100' % options.port)
        await gen.multi([
            client.fetch('http://localhost:%d/delay?ms=50' % options.port),
            client.fetch('http://localhost:%d/delay?ms=20' % options.port),
            client.fetch('http://localhost:%d/delay?ms=30' % options.port),
            ])
        self.write('all done')

def main():
    parse_command_line()
    app = Application([
        ('/', MyHandler),
        ('/delay', DelayHandler),
        ], debug=True)
    if options.trace:
        import monitor
        monitor.install()
        app = monitor.wrap_app(app)
    server = HTTPServer(app)
    server.listen(options.port)
    IOLoop.current().start()

if __name__ == '__main__':
    main()

monitor.py

import contextlib
import functools
import itertools
import threading
import tornado.httpclient
import tornado.httputil
import tornado.simple_httpclient
import tornado.stack_context

id_generator = itertools.count()

request_context = threading.local()

@contextlib.contextmanager
def set_request_id(request_id):
    request_context.id = request_id
    try:
        yield
    finally:
        request_context.id = None

def request_id_context(request_id):
    return tornado.stack_context.StackContext(functools.partial(set_request_id, request_id))

class WrapServerConnectionDelegate(tornado.httputil.HTTPServerConnectionDelegate):
    def __init__(self, app):
        self.app = app

    def start_request(self, server_conn, request_conn):
        request_id = next(id_generator)
        with request_id_context(request_id):
            request_conn = WrapConnection(request_id, request_conn)
            real_delegate = self.app.start_request(server_conn, request_conn)
            return WrapMessageDelegate(request_id, real_delegate)

    def on_close(self, server_conn):
        return self.app.on_close(server_conn)

class WrapMessageDelegate(tornado.httputil.HTTPMessageDelegate):
    def __init__(self, request_id, real_delegate):
        self.request_id = request_id
        self.real_delegate = real_delegate

    def headers_received(self, start_line, headers):
        with request_id_context(self.request_id):
            print('request %s received request headers' % request_context.id)
            return self.real_delegate.headers_received(start_line, headers)

    def data_received(self, chunk):
        with request_id_context(self.request_id):
            return self.real_delegate.headers_received(chunk)

    def finish(self):
        with request_id_context(self.request_id):
            print('request %s received entire request' % request_context.id)
            return self.real_delegate.finish()

    def on_connection_close(self):
        with request_id_context(self.request_id):
            return self.real_delegate.on_connection_close()

class WrapConnection(tornado.httputil.HTTPConnection):
    def __init__(self, request_id, real_connection):
        self.request_id = request_id
        self.real_connection = real_connection

    def write_headers(self, start_line, headers, chunk=None, callback=None):
        with request_id_context(self.request_id):
            print('request %s writing response headers with status code %s' %
                  (request_context.id, start_line.code))
            return self.real_connection.write_headers(start_line, headers, chunk, callback)

    def write(self, chunk, callback=None):
        with request_id_context(self.request_id):
            return self.real_connection.write(chunk, callback)

    def finish(self):
        with request_id_context(self.request_id):
            print('request %s done writing response' % request_context.id)
            return self.real_connection.finish()

    def set_close_callback(self, callback):
        with request_id_context(self.request_id):
            return self.real_connection.set_close_callback(callback)

class WrapAsyncHTTPClient(tornado.httpclient.AsyncHTTPClient):
    def initialize(self, *args, **kw):
        self.real_client = tornado.simple_httpclient.SimpleAsyncHTTPClient()
        super().initialize(*args, **kw)

    def fetch(self, req, **kw):
        print('request %s fetching %s' % (request_context.id, req))
        return self.real_client.fetch(req, **kw)


def wrap_app(app):
    return WrapServerConnectionDelegate(app)

def install():
    tornado.httpclient.AsyncHTTPClient.configure(WrapAsyncHTTPClient)