Ghetto Asyncio for Python 2

ghetto_asyncio.py

# -*- coding: utf-8 -*-
import os
import sys
import select
import errno
import threading
import Queue
from functools import wraps, partial
from contextlib import contextmanager, closing
import logging
from collections import defaultdict


import SocketServer
import socket


logging.basicConfig()
logger = logging.getLogger(__name__)


def defer(fn):
    def decorated(*args, **kwargs):
        @wraps(fn)
        def deferred():
            return fn(*args, **kwargs)
        return deferred
    return decorated


def catch(fn):
    def decorated(*args, **kwargs):
        try:
            return fn(*args, **kwargs)
        except:
            logger.exception('Unhandled exception in %s', fn.__name__)
    return decorated


def get_fd(int_or_fileobj):
    if isinstance(int_or_fileobj, int):
        return int_or_fileobj
    return int_or_fileobj.fileno()


def execute(poll, coros):
    for coro, deferred, done in coros:
        try:
            fd, deferred = coro.send(deferred())
            fd = get_fd(fd)
            yield coro, deferred, done
        except StopIteration:
            done.set()
        except:
            logger.exception('Exception in deferred')
            done.set()


def deregister(poll, coros, fd):
    poll.unregister(fd)
    for coro, deferred, done in coros.pop(fd):
        done.set()


@contextmanager
def read_reactor():
    poll = select.poll()

    interrupt_r, interrupt_w = os.pipe()

    def interrupt():
        read = partial(os.read, interrupt_r, 1)
        while True:
            data = yield interrupt_r, read
            if not data:
                break
            assert data == '\0'

    queue = Queue.Queue()

    shutdown = threading.Event()

    @catch
    def run():
        coros = defaultdict(list)

        while not shutdown.is_set():
            while not queue.empty():
                coro, done = queue.get()
                # fd = get_fd(fd)

                fd, deferred = coro.next()
                fd = get_fd(fd)

                coros[fd].append((coro, deferred, done))

                poll.register(
                    fd,
                    select.POLLIN | select.POLLPRI | select.POLLERR | select.POLLHUP |
                    select.POLLNVAL
                )

            try:
                ready = poll.poll()
            except select.error as err:
                if err.args[0] != errno.EINTR:
                    raise
                continue

            for fd, flags in ready:
                if flags & select.POLLIN | select.POLLPRI:
                    coros[fd] = list(execute(poll, coros[fd]))
                    if not coros[fd]:
                        deregister(poll, coros, fd)
                else:
                    deregister(poll, coros, fd)


    def put(coro):
        done = threading.Event()
        queue.put((coro, done))
        os.write(interrupt_w, '\0')
        return done.wait

    put(interrupt())

    thread = threading.Thread(target=run)
    thread.daemon = True
    thread.start()

    try:

        yield put
    finally:
        shutdown.set()
        os.close(interrupt_w)
        thread.join()


class SocketReader(object):
    __slots__ = (
        '_connection',
        '_eof',
    )

    def __init__(self, connection):
        self._connection = connection
        self._eof = False

    def read(self, count):
        def deferred():
            data = self._connection.recv(count)
            self._eof = not data
            return data
        return self._connection.fileno(), deferred

    def at_eof(self):
        return self._eof


class SocketWriter(object):
    __slots__ = (
        '_connection',
    )

    def __init__(self, connection):
        self._connection = connection

    def write(self, value):
        self._connection.sendall(value)

    def write_eof(self):
        try:
            self._connection.shutdown(socket.SHUT_WR)
        except socket.error:
            pass


def pipe(reader, writer):
    try:
        while not reader.at_eof():
            data = yield reader.read(16 * 1024)
            if data:
                writer.write(data)
    finally:
        writer.write_eof()


def gather(*waits):
    for wait in waits:
        wait()


class DispatchHandler(SocketServer.BaseRequestHandler, object):
    def __init__(self, loop, *args, **kwargs):
        self.loop = loop
        super(DispatchHandler, self).__init__(*args, **kwargs)

    @catch
    def handle(self):
        with closing(self.request):
            connection = socket.socket()
            connection.connect(('www.zombo.com', 80))

            with closing(connection):
                gather(
                    loop(
                        pipe(
                            SocketReader(self.request),
                            SocketWriter(connection)
                        )
                    ),
                    loop(
                        pipe(
                            SocketReader(connection),
                            SocketWriter(self.request)
                        )
                    )
                )


class DispatchServer(SocketServer.ThreadingTCPServer):
    daemon_threads = True
    allow_reuse_address = True
    request_queue_size = 50


if __name__ == '__main__':
    print('Listening on 9966')
    with read_reactor() as loop:
        DispatchServer(('127.0.0.1', 9966), partial(DispatchHandler, loop)).serve_forever()