Python Socket接收分块数据后无报错退出致客户端挂起问题

核心问题 #

1. 数据接收逻辑错误：
   初始接收的recv_data = conn.recv(data_length)未计入bytes_received，导致后续循环接收的字节数叠加后，总数据长度远超实际需要的data_length，触发pickle.loads()解析失败，且因未捕获异常直接退出函数。
2. 断开消息处理逻辑漏洞：
   收到DISCONNECT_MESSAGE时仅设置connected=False，但后续仍会执行遍历data（此时data是字符串）的代码，触发TypeError（字符串无法按字典取['name']），直接终止handle_client函数，导致客户端因未收到预期响应而挂起。
3. 未处理异常导致隐式退出：
   仅对删除文件的逻辑做了异常捕获，其余核心逻辑（如pickle解析、文件保存、模型发送）无异常处理，一旦出错会直接退出线程，无任何报错提示。

修复后的代码 #

import pickle
import socket
import os
import threading

HEADER = 64
PORT = 5050
SERVER = 'Left this out for obvious reasons'
ADDR = (SERVER, PORT)
FORMAT = 'utf-8'
DISCONNECT_MESSAGE = "..DISCONNECT"

server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
server.bind(ADDR)

def handle_client(conn, addr):
    print(f"[NEW CONNECTION] {addr} connected.")
    connected = True
    try:
        while connected:
            # 接收数据长度头部
            data_length_bytes = conn.recv(HEADER)
            if not data_length_bytes:
                break
            data_length = int(data_length_bytes.decode(FORMAT).strip())
            
            # 正确分块接收完整数据
            recv_data = b''
            bytes_received = 0
            while bytes_received < data_length:
                chunk = conn.recv(min(4096, data_length - bytes_received))
                if not chunk:
                    # 连接中断
                    connected = False
                    break
                recv_data += chunk
                bytes_received += len(chunk)
            
            if not connected:
                break
                
            data = pickle.loads(recv_data)
            
            # 处理断开请求
            if data == DISCONNECT_MESSAGE:
                connected = False
                print(f"[DISCONNECT] {addr} requested disconnect.")
                continue  # 跳过后续处理逻辑
            
            # 清空incoming文件夹
            try:
                for filename in os.listdir('incoming'):
                    file_path = os.path.join('incoming', filename)
                    if os.path.isfile(file_path):
                        os.remove(file_path)
            except FileNotFoundError:
                print('[SERVER_INFO] ./incoming/ folder does not exist, creating it.')
                os.makedirs('incoming')
            except Exception as e:
                print(f'[SERVER_ERROR] Failed to clear incoming folder: {str(e)}')
            
            # 保存图片文件
            try:
                for file in data:
                    save_path = os.path.join('incoming', file['name'])
                    with open(save_path, "wb") as f:
                        f.write(file['data'])
                    print(f'[SAVING] Saved file: {file["name"]}')
            except Exception as e:
                print(f'[SERVER_ERROR] Failed to save files: {str(e)}')
                continue
            
            # 发送模型文件
            try:
                with open('siamesemodel.h5','rb') as f:
                    model_data = f.read()
                pickled_data = pickle.dumps(model_data)
                send_length = str(len(pickled_data)).encode(FORMAT)
                send_length += b' ' * (HEADER - len(send_length))
                conn.send(send_length)
                conn.send(pickled_data)
                print(f'[SEND] Sent model to {addr}')
            except Exception as e:
                print(f'[SERVER_ERROR] Failed to send model: {str(e)}')
    except Exception as e:
        print(f'[CLIENT_ERROR] Connection with {addr} failed: {str(e)}')
    finally:
        conn.close()
        print(f"[CONNECTION CLOSED] {addr} disconnected.")

def start():
    # 确保incoming文件夹存在
    if not os.path.exists('incoming'):
        os.makedirs('incoming')
    server.listen()
    print(f"[LISTENING] Server is listening on {SERVER}")
    while True:
        conn, addr = server.accept()
        thread = threading.Thread(target=handle_client, args=(conn, addr))
        thread.start()
        print(f"[ACTIVE CONNECTIONS] {threading.active_count() - 1}")

print("[STARTING] server is starting...")
start()

关键修复点 #

• 修正数据接收逻辑：初始化recv_data为空字节串，循环接收时每次取剩余需要接收的字节数（min(4096, data_length - bytes_received)），确保接收的数据长度完全匹配data_length。
• 优化断开消息处理：收到断开消息后，使用continue跳过后续业务逻辑，避免无效遍历操作。
• 全局异常捕获：在handle_client外层添加异常捕获，任何错误都会输出日志并正常关闭连接，避免线程隐式退出。
• 完善文件夹处理：检查incoming文件夹是否存在，不存在则自动创建；删除文件时仅删除文件（避免误删子文件夹）。
• 添加详细日志：每个关键步骤添加日志，便于调试和排查问题。

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