python3 IO多路复用之selector

        最近在写工具遇到IO多路复用，具体原理就不细说了，百度很多。大概意思就是说系统可以一次性管理N个socket文件描述符，然后不断遍历，每次有可读/可写的对象时就返回，应用层就可以直接使用。没有返回时，应用层不用阻塞等待，可以节省大量资源。适合大批量IO操作，比如后台服务。

这次练习使用selector，是对select的高级封装，当然也包含pool、epoll等模块，取决于系统支持哪个。

server端：

import selectors
import socket


# 自动选择最佳IO多路机制
sel = selectors.DefaultSelector()


def write(conn, mask):
    """
    发送数据
    :param conn: 连接
    :param mask: event类型
    :return:
    """
    messages = message.get(conn)
    if messages:
        # 获取消息
        this_message = messages.pop(0)
        # 发送消息
        conn.send(this_message)
        # 切换为可读状态
        sel.modify(conn, selectors.EVENT_READ, read)
    else:
        # 如果没有数据则注销监听
        sel.unregister(conn)


def read(conn, mask):
    """
    读取消息
    :param conn: 连接
    :param mask: event类型
    :return:
    """
    # 读取数据
    data = conn.recv(1024)
    if not data:
        # 没有数据则表示断开连接
        message.pop(conn)

    if conn not in message:
        message[conn] = []
    message[conn].append(data)

    # 修改当前连接为可写状态
    sel.modify(conn, selectors.EVENT_WRITE, write)


def accept(sock, mask):
    """
    server连接读取
    :param sock:
    :param mask:
    :return:
    """
    # 等待客户端建立连接
    client, addr = sock.accept()
    # 设置非阻塞
    client.setblocking(False)
    # 注册连接，加入监听列表
    sel.register(client, selectors.EVENT_READ, read)


# 初始化socket
server = socket.socket()
# 绑定地址和端口
server.bind(('0.0.0.0', 8889))
# 监听
server.listen(100)

# 记录接收到的消息，后面要发出去
message = {}

# 设置非阻塞状态，很重要
server.setblocking(False)

# 注册server的IO事件， accept是收到消息后回调用的
sel.register(server, selectors.EVENT_READ, accept)

while 1:
    # 检测有无活动对象（读、写事件），没有就阻塞
    events = sel.select()
    for key, mask in events:
        callback = key.data         # data就是注册时写的accept
        callback(key.fileobj, mask)

client端：

import socket

s = socket.socket()

ss = s.connect(('127.0.0.1', 8889))

while 1:
    message = input('>>>输入待发送消息: ')
    if message == 'exit':
        s.send(b'')
        break
    s.send(message.encode())
    data = s.recv(1024)
    print(f'client收到消息: {data.decode("utf-8")}')


s.close()

官方文档：selectors --- 高级 I/O 复用库 — Python 3.11.0 文档