reorx
11/30/2011 - 1:46 AM

threaded decorator for tornado

threaded decorator for tornado

from functools import partial, wraps
from threading import Thread

from tornado import stack_context, ioloop


def inline_threaded(func):

    @wraps(func)
    def wrapper(*args, **kwargs):

        io_loop = ioloop.IOLoop.instance()

        # callback is expected to be in kwargs or to be the last argument
        if 'callback' in kwargs:
            orig_cb = kwargs.pop('callback')
        else:
            args, orig_cb = args[:-1], args[-1]
            assert callable(orig_cb)

        @stack_context.wrap
        def callback(result):
            @wraps(orig_cb)
            def thread_stopper():
                thread.join()
                orig_cb(result)
            io_loop.add_callback(thread_stopper)

        @stack_context.wrap
        def raise_exc(e):
            thread.join()
            raise e

        @wraps(func)
        def exception_handler():
            try:
                callback(func(*args, **kwargs))
            except Exception as e:
                io_loop.add_callback(partial(raise_exc, e))

        thread = Thread(target=exception_handler)
        thread.start()

    return wrapper


def threaded(func):

    @wraps(func)
    def wrapper(*args, **kwargs):

        io_loop = ioloop.IOLoop.instance()

        # callback is expected to be in kwargs or to be the last argument
        if 'callback' in kwargs:
            orig_cb = kwargs['callback']
        else:
            args, orig_cb = args[:-1], args[-1]
            assert callable(orig_cb)

        @stack_context.wrap
        def callback(*args, **kwargs):
            @wraps(orig_cb)
            def thread_stopper():
                thread.join()
                orig_cb(*args, **kwargs)
            io_loop.add_callback(thread_stopper)

        kwargs['callback'] = callback

        @stack_context.wrap
        def raise_exc(e):
            thread.join()
            raise e

        @wraps(func)
        def exception_handler():
            try:
                func(*args, **kwargs)
            except Exception as e:
                io_loop.add_callback(partial(raise_exc, e))

        thread = Thread(target=exception_handler)
        thread.start()

    return wrapper
from time import sleep

from tornado import gen
import tornado.ioloop
import tornado.web

from tornado_threaded import threaded, inline_threaded

@threaded
def slow_func(callback):
    sleep(5)
    callback("Hello, world!")

@inline_threaded
def slow_func2():
    sleep(5)
    return "Hello, world!"

class MainHandler(tornado.web.RequestHandler):
    @tornado.web.asynchronous
    @gen.engine
    def get(self):
        result = yield gen.Task(slow_func)
        self.write(result)
        self.finish()

application = tornado.web.Application([
    (r"/", MainHandler),
])

if __name__ == "__main__":
    application.listen(8888)
    tornado.ioloop.IOLoop.instance().start()