Socket接收停止异常：send中assert影响及可靠数据传输方案

为什么assert会影响接收端行为？ #

当发送端的assert 0 == 1触发时，程序直接抛出AssertionError并终止进程。操作系统会自动关闭该进程持有的所有socket连接，此时接收端的socket.recv()会收到空字节串（b''），触发if not chunk:的条件，从而退出循环。

如果移除这个assert，发送端在调用socket.send()后进程不会终止，socket连接保持打开状态。接收端收完所有数据后，会一直阻塞等待新数据——因为它不知道发送端已经发完了，socket连接未关闭，所以永远触发不了break条件。

注意：socket.send()只是把数据放到操作系统发送缓冲区就返回，最终数据都会发出去，但核心问题是接收端没有收到“传输结束”的信号。

如何实现序列化数据的可靠传输？ #

要实现可靠传输，核心是让接收端明确知道“数据已全部发送完成”，常见方案有三种：

方案1：发送完数据后优雅关闭socket写端 #

发送端发送完成后，调用socket.shutdown(socket.SHUT_WR)（仅关闭写方向，不关闭整个socket，还能接收响应），或直接调用socket.close()（无需再通信时）。这样接收端的recv()会收到空字节，判断传输结束。

修改后的发送端代码：

import pickle
import socket

def send_model(socket, state_dict):
    """Sends a PyTorch model state_dict over a network socket."""
    serialized = pickle.dumps(state_dict)
    # 确保所有数据发送完毕
    total_sent = 0
    while total_sent < len(serialized):
        sent = socket.send(serialized[total_sent:])
        if sent == 0:
            raise RuntimeError("Socket connection broken")
        total_sent += sent
    # 关闭写端，告知接收端数据发送完成
    socket.shutdown(socket.SHUT_WR)

方案2：先发送数据长度，再发实际数据 #

先把序列化后的数据长度打包成固定字节的报文（比如4字节大端整数），接收端先接收长度，再按长度接收对应字节数的数据，无需依赖socket关闭判断结束。

修改后的发送端：

import pickle
import socket
import struct

def send_model(socket, state_dict):
    serialized = pickle.dumps(state_dict)
    # 打包数据长度为4字节大端整数
    length = struct.pack('!I', len(serialized))
    socket.sendall(length)
    # 发送实际数据
    socket.sendall(serialized)

接收端代码：

import pickle
import socket
import struct

def recv_model(socket):
    # 先接收4字节的长度信息
    length_data = b''
    while len(length_data) < 4:
        chunk = socket.recv(4 - len(length_data))
        if not chunk:
            raise RuntimeError("Socket connection broken")
        length_data += chunk
    length = struct.unpack('!I', length_data)[0]
    
    # 接收指定长度的数据
    data = b''
    while len(data) < length:
        chunk = socket.recv(min(10000, length - len(data)))
        if not chunk:
            raise RuntimeError("Socket connection broken")
        data += chunk
    
    return pickle.loads(data)

方案3：使用应用层结束标记 #

在数据末尾添加特殊结束标记（比如b'END_OF_DATA'），接收端接收时检查是否收到该标记，收到后停止接收。但要注意如果数据本身包含该标记会出错，仅适合数据内容可控的场景。

补充说明 #

• 实际工程中建议用socket.sendall()代替socket.send()，它会循环调用send()直到所有数据发送完成。
• 使用pickle传输数据要注意安全，不要接收不可信来源的数据，避免反序列化漏洞。

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