100行代码实现gevent调度模型

Posted 2022-12-03 yueguanghaidao

昨天心血来潮看了https://github.com/locustio/locust/的源码，经常用ab或者httpload的童鞋可以看下locust的代码，支持分布式运行，通过gevent可以单机开N个协程模仿并发用户，分布式rpc用了zmq的PUSH,PULL模式，不得不说zeromq的确简洁。２年前研究过gevent，那时就想用纯Python模拟一下调度模型，后来也就不了了之了。原理其实和gevent一样，也是通过多个greenlet互相switch实现的。

# coding=utf8

import errno
import socket
import select
import greenlet as rawgreenlet
from greenlet import greenlet


class Sock(object):
    def __init__(self, socket_=None):
        if socket_ is None:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        else:
            sock = socket_
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.setblocking(0)
        self.sock = sock

    def listen(self, backlog):
        self.sock.listen(backlog)

    def bind(self, host, port):
        self.sock.bind((host, port))

    def accept(self):
        while True:
            try:
                print("accept")
                client_sock, address = self.sock.accept()
                break
            except socket.error as e:
                print(e)
                if e.args[0] != errno.EAGAIN:
                    raise
            switch = rawgreenlet.getcurrent().switch
            ioloop.wait(self.sock, switch, ioloop.READ)

        return Sock(client_sock), address

    def recv(self, *args):
        while True:
            try:
                print("recv")
                return self.sock.recv(*args)
            except socket.error as e:
                print(e)
                if e.args[0] != errno.EAGAIN:
                    raise
            switch = rawgreenlet.getcurrent().switch
            ioloop.wait(self.sock, switch, ioloop.READ)


def spawn(f):
    g = greenlet(f, parent=ioloop)
    ioloop.add_callback(g.switch)
    return g


class IOLoop(greenlet):
    """"main greenlet"""
    READ = select.EPOLLIN
    WRITE = select.EPOLLOUT
    ERROR = select.EPOLLERR | select.EPOLLHUP

    def __init__(self):
        self.poller = select.epoll()
        self.handler = 
        self.callbacks = []

    def wait(self, sock, callback, event):
        """wait until waiter avaliable"""
        self.add_handler(sock, callback, event)
        self.switch()

    def add_handler(self, sock, callback, event):
        """
        fd: file description
        callback: when event avaliable, invoke callback
        evnet: poll evnet
        """
        fd = sock.fileno()
        if fd in self.handler:
            self.poller.unregister(fd)
        self.poller.register(fd, event)
        self.handler[fd] = callback

    def add_callback(self, callback):
        self.callbacks.append(callback)

    def start(self):
        self.switch()

    def run(self):
        print("ioloop run")
        while True:
            # invoke callback
            while self.callbacks:
                callback = self.callbacks.pop()
                callback()

            # poller
            events = self.poller.poll(1)
            print("poller:", events)
            for fd, event in events:
                if event & self.READ:
                    handler = self.handler[fd]
                    handler()

ioloop = IOLoop()


def f():
    sock = Sock()
    sock.bind("localhost", 8088)
    sock.listen(5)
    client_sock, address = sock.accept()
    print("connection from:", address)
    while 1:
        print(client_sock.recv(10))

    # return this greenlet will dead

spawn(f)
ioloop.start()

github地址https://github.com/Skycrab/code/blob/master/Python/mygevent.py,
ioloop就是主协程，一直在检测是否有回调和socket事件，通过spawn创建一个新的协程，当socket需要accept时，注册可读事件，switch到主协程，当socket可读时再切换到之前运行的协程处。如果熟悉golang的童鞋，模型都是一样的，ioloop是golang底层的主gorutine，spawn函数相当于go关键字，剩下的就都是靠ioloop去切换了。

上面的代码只是个原型，有兴趣的童鞋可以继续完善，当然也可以关注https://github.com/eventlet/eventlet。Python3用户可以关注https://github.com/python/asyncio，不过满屏的async和await看起来真心不舒服，还是觉得gevent代码写的最舒心。

以前写过几篇gevnet的文章，感兴趣的童鞋可以看看，对你理解应该很有帮助。