Python Socket传输pickle数据偶现UnpicklingError截断问题求解

问题根因 #

• 核心原因是TCP socket的recv()方法不保证一次性返回指定长度的全部数据，仅会返回当前内核缓冲区中可用的最多字节数。当出现网络波动、系统缓冲区数据不足的场景时，调用rscSock.recv(int(resp['arg0'])+1)只会拿到部分序列化数据，就会触发pickle data was truncated报错，因为触发场景是偶发的网络/系统状态变化，所以大部分时候运行正常。
• 分块接收卡住的原因是没有设置正确的循环终止条件：未统计已接收的总字节数和预期长度对比，导致收够数据后仍在等待新的数据包，自然会卡住。
• 额外潜在问题：当前使用的conn.send(graphs)也存在类似问题，send()同样不保证一次性把所有数据发送完成，仅会返回实际发送的字节数，当数据量大、网络状态不佳时，服务端可能只发送了部分数据，客户端自然无法收全完整数据。

解决方案 #

首先实现两个通用的可靠收发工具函数，保证完全接收/发送指定长度的数据：

# 可靠接收指定长度的所有字节
def recv_all(conn, expected_len):
    received = b''
    remaining = expected_len
    while remaining > 0:
        # 每次最多收4096字节，不够就收剩余需要的长度
        chunk = conn.recv(min(4096, remaining))
        if not chunk:
            raise ConnectionError("连接已中断，接收失败")
        received += chunk
        remaining -= len(chunk)
    return received

# 可靠发送所有字节
def send_all(conn, data):
    total_sent = 0
    remaining = len(data)
    while remaining > 0:
        sent = conn.send(data[total_sent:])
        if sent == 0:
            raise ConnectionError("连接已中断，发送失败")
        total_sent += sent
        remaining -= sent
    return total_sent

修改服务端代码 #

将发送pickle数据的部分替换为可靠发送：

elif cmd['cmd'] == 'RSC_VIEW_GRAPHS':
    graphs = pickle.dumps(genGraphs(userSession['uid'], cmd['arg0'], cmd['arg1']))
    send(conn, 'RSC_IMG_DATA', len(graphs))
    send_all(conn, graphs) # 替换原有的conn.send(graphs)
    del graphs

修改客户端代码 #

将接收pickle数据的部分替换为可靠接收：

if resp['resp'] == 'RSC_IMG_DATA':
    expected_len = int(resp['arg0'])
    graphs_bytes = recv_all(rscSock, expected_len)
    graphs = pickle.loads(graphs_bytes)

可选优化 #

现有通用send()、receive()函数里的收发逻辑也可以替换为上述可靠收发函数，避免控制信令（JSON交互）也出现偶发截断问题，进一步降低异常概率。

内容的提问来源于stack exchange，提问作者Burns Milheron