TCP Socket通信出现数据丢失，求代码问题排查方案

问题背景 #

在同机环境下进行TCP Socket开发，服务端向连接的客户端发送随机长度（由"12345"重复10万-20万次生成）的字节数据，模拟可变长度数据传输场景。服务端采用先发送消息长度（4字节小端格式）、再发送实际消息的方案保障完整性，但100次连接测试中多次出现数据丢失，接收的消息长度与约定长度不符。

服务端代码 #

import socket
import struct
import random as rand

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_address = ('localhost', 12345)
sock.bind(server_address)
sock.listen(1)

while True:
    message = "12345" * rand.randint(100000, 200000)
    connection, client_address = sock.accept()
    message_in_bytes = message.encode()
    length = len(message_in_bytes)
    print(client_address, "connected and message length is ", length)
    length_bytes = struct.pack("<I", length)
    connection.send(length_bytes + message_in_bytes)
    connection.close()

客户端代码（原问题版本） #

import socket
import struct

def connect_server():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_address = ('localhost', 12345)
    sock.connect(server_address)
    message_length_bytes = sock.recv(4)
    length_of_message = struct.unpack("<I", message_length_bytes)[0]
    whole_message = sock.recv(length_of_message)
    sock.close()
    if length_of_message == len(whole_message):
        return "success"
    else:
        return "failed"

stats={'success':0,'failed':0}

for _ in range(100):
    stats[connect_server()] +=1

print(stats)

问题根源 #

TCP是流式协议，数据以字节流形式传输，不存在"消息边界"：

• 客户端调用sock.recv(n)时，仅会返回当前缓冲区中可用的字节，不一定等于请求的n字节，尤其是当数据量较大时，数据会被拆分成多个TCP段分批到达。
• 服务端的connection.send()也可能无法一次性发送全部数据，返回值为实际发送的字节数，原代码未处理这种情况，可能导致部分数据未发送就关闭连接。

解决方案 #

1. 客户端：循环接收直到获取完整数据 #

实现通用的recv_all函数，确保接收指定长度的字节：

import socket
import struct

def recv_all(sock, length):
    """循环接收直到获取指定长度的字节"""
    data = b''
    while len(data) < length:
        chunk = sock.recv(length - len(data))
        if not chunk:
            # 连接提前关闭，数据不完整
            return None
        data += chunk
    return data

def connect_server():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_address = ('localhost', 12345)
    sock.connect(server_address)
    
    # 接收长度字段
    message_length_bytes = recv_all(sock, 4)
    if not message_length_bytes:
        sock.close()
        return "failed"
    length_of_message = struct.unpack("<I", message_length_bytes)[0]
    
    # 接收完整消息
    whole_message = recv_all(sock, length_of_message)
    sock.close()
    
    if whole_message is not None and length_of_message == len(whole_message):
        return "success"
    else:
        return "failed"

stats={'success':0,'failed':0}

for _ in range(100):
    stats[connect_server()] +=1

print(stats)

2. 服务端：确保数据完整发送后再关闭连接 #

修改服务端发送逻辑，循环发送直到所有数据都被发送：

import socket
import struct
import random as rand

def send_all(sock, data):
    """循环发送直到所有数据都被发送"""
    sent = 0
    while sent < len(data):
        sent += sock.send(data[sent:])

sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server_address = ('localhost', 12345)
sock.bind(server_address)
sock.listen(1)

while True:
    message = "12345" * rand.randint(100000, 200000)
    connection, client_address = sock.accept()
    message_in_bytes = message.encode()
    length = len(message_in_bytes)
    print(client_address, "connected and message length is ", length)
    length_bytes = struct.pack("<I", length)
    send_all(connection, length_bytes + message_in_bytes)
    connection.close()

验证效果 #

修改后重新运行100次测试，可实现100%的success，解决数据丢失问题。

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