Python selectors场景下如何完整接收TCP分片消息

核心认知纠正 #

TCP是流传输协议，本身不存在「单次send操作边界」的概念：

• 不管是阻塞还是非阻塞模式，recv()返回0仅代表对端已经关闭了连接，绝非「单次发送的所有数据接收完成」
• 你之前阻塞socket场景下的判断逻辑仅适用于「请求发送完成后对端立即关闭连接」的短连接场景，长连接下不管阻塞还是非阻塞都需要应用层自行定义报文边界

非阻塞模式下，recv()在系统内核接收缓冲区没有可读数据时，不会阻塞等待数据到达，会直接抛出BlockingIOError异常，你照搬阻塞模式的死循环读取逻辑，一定会在缓冲区读完时触发报错。

解决方案 #

1. 正确实现非阻塞模式下的持续接收

你需要为每个连接维护独立的持久化接收缓冲区，不要在read回调内定义局部缓冲区（每次回调触发局部变量都会重置），同时在循环读取时捕获BlockingIOError判断当前缓冲区已读完。你可以把连接对应的状态（缓冲区、其他自定义参数）和回调绑定，注册到selectors中，修改后的示例代码如下：

import selectors
import socket
from collections import namedtuple

# 定义连接状态结构：存储接收缓冲区、处理回调
ConnState = namedtuple('ConnState', ['buffer', 'callback'])

sel = selectors.DefaultSelector()

def accept(sock, mask):
    conn, addr = sock.accept()
    print('Connected by', addr)
    conn.setblocking(False)
    # 注册连接时绑定初始状态：空缓冲区+read处理函数
    state = ConnState(buffer=bytearray(), callback=read)
    sel.register(conn, selectors.EVENT_READ, state)

def read(conn, mask):
    state = sel.get_key(conn).data
    while True:
        try:
            data = conn.recv(1000)
        except BlockingIOError:
            # 当前缓冲区所有数据已读完，退出循环等待下一次EVENT_READ
            break
        if not data:
            # 对端关闭连接，清理资源
            print('closing', conn)
            sel.unregister(conn)
            conn.close()
            return
        # 数据追加到对应连接的持久化缓冲区
        state.buffer.extend(data)
    
    # 到这里已经读完当前所有到达的数据，接下来可以判断是否收到完整请求
    # 这里以长度前缀协议举例子：前4字节是大端模式的报文长度
    while len(state.buffer) >=4:
        # 读取前4字节得到报文总长度
        msg_len = int.from_bytes(state.buffer[:4], byteorder='big')
        if len(state.buffer) < 4 + msg_len:
            # 报文还没收全，等待下一次数据到达
            break
        # 切出完整报文
        msg = state.buffer[4:4+msg_len]
        # 剩下的数据留在缓冲区等待下一次处理
        state.buffer = state.buffer[4+msg_len:]
        
        # 这里处理你的业务逻辑，比如回显
        print('received complete msg:', repr(msg))
        # 非阻塞send注意也可能触发BlockingIOError，生产环境需要维护发送缓冲区监听EVENT_WRITE，这里简化处理
        try:
            conn.sendall(msg)
        except BlockingIOError:
            # 生产环境这里需要把数据放到发送缓冲区，注册EVENT_WRITE事件等待可写时再发
            pass

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
sock.bind(('localhost', 1234))
sock.listen(100)
sock.setblocking(False)
sel.register(sock, selectors.EVENT_READ, None)

while True:
    events = sel.select()
    for key, mask in events:
        if key.data is None:
            accept(key.fileobj, mask)
        else:
            state = key.data
            state.callback(key.fileobj, mask)

2. 应用层定义报文边界的常用方案

要区分不同的请求，必须在应用层自行约定协议规则，常用方案有三种：

• 固定长度报文：约定所有请求的长度固定，缓冲区长度达到约定值就切出一个完整请求，适合报文长度固定的场景
• 特殊分隔符：约定特殊字符/字符串作为请求结束标记（比如HTTP的空行、文本协议常用的\r\n），缓冲区中检索到分隔符就切分到分隔符前的内容作为完整请求，注意要处理分隔符被拆分到两次recv结果中的情况
• 长度前缀：每个报文开头用固定字节数存储整个报文的payload长度，如上方案例，先读够长度字段，再按长度读取对应长度的payload，是通用性最强、出错概率最低的方案

内容的提问来源于stack exchange，提问作者Michael B. Ortiz