coroutine rocks

echoserv.py

import socket
from pyos import Scheduler
from pyos import NewTask
from pyos import Accept, Send, Receive


def handle_client(client, addr):
    print "Connection from",  addr
    while True:
        data = yield Receive(client, 65536)
        if not data:
            break
        yield Send(client, data)
    client.close()
    print "Client closed"
    yield


def server(port):
    print "Server starting"
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
    sock.bind(("", port))
    sock.listen(5)
    while True:
        client, addr = yield Accept(sock)
        yield NewTask(handle_client(client, addr))


def alive():
    while True:
        print "I'm alive!"
        yield

sched = Scheduler()
#sched.new(alive())
sched.new(server(8000))
sched.mainloop()

pyos.py

from Queue import Queue
import select
import types


class Task(object):
    taskid = 0

    def __init__(self, target):
        Task.taskid += 1
        self.stack = []

        self.tid = Task.taskid
        self.target = target
        self.sendval = None

    def run(self):
        while True:
            try:
                result = self.target.send(self.sendval)
                if isinstance(result, SystemCall):
                    return result
                if isinstance(result, types.GeneratorType):
                    self.stack.append(self.target)
                    self.sendval = None
                    self.target = result
                else:
                    if not self.stack: return
                    self.sendval = result
                    self.target = self.stack.pop()
            except StopIteration:
                if not self.stack: raise
                self.sendval = None
                self.target = self.stack.pop()


class SystemCall(object):
    def __init__(self):
        self.task = None
        self.sched = None

    def handle(self):
        pass


class Scheduler(object):

    def __init__(self):
        self.ready = Queue()
        self.taskmap = {}
        self.exit_waiting = {}
        self.read_waiting = {}
        self.write_waiting = {}

    def new(self, target):
        task = Task(target)
        self.taskmap[task.tid] = task
        self.schedule(task)
        return task.tid

    def schedule(self, task):
        self.ready.put(task)

    def exit(self, task):
        print 'Task {} terminated.'.format(task.tid)
        del self.taskmap[task.tid]
        for task in self.exit_waiting.pop(task.tid, []):
            self.schedule(task)

    def waitforexit(self, task, waittid):
        if waittid in self.taskmap:
            self.exit_waiting.setdefault(waittid, []).append(task)
            return True
        else:
            return False

    def waitforread(self, task, fd):
        self.read_waiting[fd] = task

    def waitforwrite(self, task, fd):
        self.write_waiting[fd] = task

    def ioloop(self, timeout):
        if self.read_waiting or self.write_waiting:
            r, w, e = select.select(self.read_waiting.keys(),
                                    self.write_waiting.keys(),
                                    [],
                                    timeout)
            for fd in r:
                self.schedule(self.read_waiting.pop(fd))
            for fd in w:
                self.schedule(self.write_waiting.pop(fd))

    def iotask(self):
        while True:
            if self.ready.empty():
                self.ioloop(None)
            else:
                self.ioloop(0)
            yield

    def mainloop(self):
        self.new(self.iotask())
        while self.taskmap:
            task = self.ready.get()
            try:
                result = task.run()
                if isinstance(result, SystemCall):
                    result.task = task
                    result.sched = self
                    result.handle()
                    continue
            except StopIteration:
                self.exit(task)
                continue
            self.schedule(task)


class GetTid(SystemCall):
    def handle(self):
        self.task.sendval = self.task.tid
        self.sched.schedule(self.task)


class NewTask(SystemCall):
    def __init__(self, target):
        self.target = target

    def handle(self):
        tid = self.sched.new(self.target)
        self.task.sendval = tid
        self.sched.schedule(self.task)


class KillTask(SystemCall):
    def __init__(self, tid):
        self.tid = tid

    def handle(self):
        task = self.sched.taskmap.get(self.tid)
        if task:
            task.target.close()
            self.task.sendval = True
        else:
            self.task.sendval = False
        self.sched.schedule(self.task)


class WaitTask(SystemCall):
    def __init__(self, tid):
        self.tid = tid

    def handle(self):
        result = self.sched.waitforexit(self.task, self.tid)
        self.task.sendval = result
        if not result:
            self.sched.schedule(self.task)


class ReadWait(SystemCall):
    def __init__(self, f):
        self.f = f

    def handle(self):
        self.sched.waitforread(self.task, self.f.fileno())


class WriteWait(SystemCall):
    def __init__(self, f):
        self.f = f

    def handle(self):
        self.sched.waitforwrite(self.task, self.f.fileno())


def Accept(sock):
    yield ReadWait(sock)
    yield sock.accept()


def Send(sock, buf):
    while True:
        yield WriteWait(sock)
        sz = yield sock.send(buf)
        buf = sock[sz:]


def Receive(sock, maxbytes):
    yield ReadWait(sock)
    yield sock.recv(maxbytes)


def foo():
    mytid = yield GetTid()
    for i in xrange(10):
        print "I'm foo", mytid
        yield  # context switch


def main():
    child = yield NewTask(foo())
    print 'waiting for child'
    yield WaitTask(child)
    print 'child done'
    yield WaitTask(child)
    print 'wait failed'


if __name__ == '__main__':
    sched = Scheduler()
    sched.new(main())
    sched.mainloop()