如何将Python多线程Socket客户端/服务端转为asyncio并单脚本测试？

我正在编写一个工具对基于asyncio的服务进行端到端测试。最初打算在一个终端启动服务，另一个终端运行测试，但后来意识到应该可以在单个脚本中实现。我可以用concurrent.futures.ThreadPoolExecutor实现，但在转换为await/async def逻辑时遇到困难。

可行的ThreadPoolExecutor示例 #

import argparse
import socket
import concurrent.futures
import threading
import socketserver


class TCPHandler(socketserver.BaseRequestHandler):
    def handle(self):
        print(f'Got data: {self.request.recv(1024).strip().decode()}')

def started_server(*, server):
    print('starting server')
    server.serve_forever()
    print('server thread closing')


def run_client(*, host, port, server):
    print('client started, attempting connection')
    with socket.create_connection((host, port), timeout=0.5) as conn:
        print('connected')
        conn.send(b'hello werld')
    print('closing server')
    server.shutdown()
    print('cancelled')

def test_the_server(*, host, port):
    ex = concurrent.futures.ThreadPoolExecutor(max_workers=3)
    print('server a')
    quitter = threading.Event()
    server = socketserver.TCPServer((host, port), TCPHandler)
    a = ex.submit(started_server, server=server)
    b = ex.submit(run_client, host=host, port=port, server=server)
    print(a.result(), b.result())
    print('server b')


def do_it():  # Shia LeBeouf!
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("-p", "--port", type=int, default=60025)
    args = parser.parse_args()

    exit(test_the_server(host=args.host, port=args.port))


if __name__ == "__main__":
    do_it()

失败的asyncio尝试 #

我尝试转换为asyncio循环，但实现会阻塞，且调用server.close()无法触发asyncio.Server.wait_closed()完成服务关闭。以下是失败代码：

import asyncio
import argparse
import socket
import concurrent.futures
import threading
import socketserver


class EchoHandler(asyncio.Protocol):
    def data_received(self, data):
        print(f"Got this data: {data.decode()}")

async def run_server(*, server):
    print('starting server')
    server = await server
    async with server:
        print('start serving')
        await server.start_serving()
        print('waiting on close')
        await server.wait_closed()
    print('server coro closing')

def started_server(*, server):
    print('server thread started')
    asyncio.run(run_server(server=server))
    print('server thread finished')

def run_client(*, host, port, server):
    print('client started, attempting connection')
    with socket.create_connection((host, port), timeout=0.5) as conn:
        print('connected')
        conn.send(b'hello werld')
    print('closing server')
    server.close()
    print('cancelled')

async def fnord(reader, writer):
    data = await reader.read(100)
    message = data.decode()
    print('got', message)

def test_the_server(*, host, port):
    ex = concurrent.futures.ThreadPoolExecutor(max_workers=3)
    print('server a')
    quitter = threading.Event()
    #server = socketserver.TCPServer((host, port), TCPHandler)
    server = asyncio.start_server(fnord, host, port)
    a = ex.submit(started_server, server=server)
    b = ex.submit(run_client, host=host, port=port, server=server)
    print(a.result(), b.result())
    print('server b')


def do_it():  # Shia LeBeouf!
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("-p", "--port", type=int, default=60025)
    args = parser.parse_args()

    exit(test_the_server(host=args.host, port=args.port))


if __name__ == "__main__":
    do_it()

解决方案 #

方案1：线程+asyncio（兼容同步客户端） #

核心问题是你跨线程操作了未解析的协程对象，且没有用正确的线程间信号控制服务关闭。修正后的代码如下：

import asyncio
import argparse
import socket
import threading
import time

async def fnord(reader, writer):
    data = await reader.read(100)
    message = data.decode()
    print('got', message)
    # 可选：回复客户端完成完整交互
    writer.write(b'received: ' + data)
    await writer.drain()
    writer.close()
    await writer.wait_closed()

async def run_server(host, port, stop_event):
    print('starting server')
    server = await asyncio.start_server(fnord, host, port)
    print('start serving')
    await server.start_serving()
    
    # 等待线程间停止信号
    await stop_event.wait()
    
    print('closing server')
    server.close()
    await server.wait_closed()
    print('server coro closing')

def server_thread(host, port, stop_event):
    print('server thread started')
    # 在线程内创建独立的事件循环
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    loop.run_until_complete(run_server(host, port, stop_event))
    print('server thread finished')

def run_client(host, port):
    print('client started, attempting connection')
    with socket.create_connection((host, port), timeout=0.5) as conn:
        print('connected')
        conn.send(b'hello werld')
        # 可选：接收服务端回复
        response = conn.recv(1024)
        print(f'received response: {response.decode()}')
    print('client finished')

def test_the_server(host, port):
    print('server a')
    # 创建线程间停止信号
    stop_event = asyncio.Event()
    
    # 启动服务线程
    server_t = threading.Thread(target=server_thread, args=(host, port, stop_event))
    server_t.start()
    
    # 简单延时等待服务启动（实际可改用端口检测逻辑）
    time.sleep(0.1)
    
    # 运行同步客户端
    run_client(host, port)
    
    # 触发服务停止
    stop_event.set()
    server_t.join()
    
    print('server b')

def do_it():
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("-p", "--port", type=int, default=60025)
    args = parser.parse_args()
    test_the_server(host=args.host, port=args.port)

if __name__ == "__main__":
    do_it()

方案2：全asyncio实现（更简洁） #

如果客户端也可以改为异步逻辑，完全不需要线程，直接在同一个事件循环中完成：

import asyncio
import argparse

async def fnord(reader, writer):
    data = await reader.read(100)
    message = data.decode()
    print('got', message)
    writer.write(b'received: ' + data)
    await writer.drain()
    writer.close()
    await writer.wait_closed()

async def run_client(host, port):
    print('client started, attempting connection')
    reader, writer = await asyncio.open_connection(host, port)
    print('connected')
    writer.write(b'hello werld')
    await writer.drain()
    
    response = await reader.read(100)
    print(f'received response: {response.decode()}')
    
    writer.close()
    await writer.wait_closed()
    print('client finished')

async def test_the_server(host, port):
    print('server a')
    # 启动服务
    server = await asyncio.start_server(fnord, host, port)
    print('start serving')
    
    # 运行客户端任务
    await run_client(host, port)
    
    # 关闭服务
    print('closing server')
    server.close()
    await server.wait_closed()
    print('server coro closing')
    print('server b')

def do_it():
    parser = argparse.ArgumentParser(usage=__doc__)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("-p", "--port", type=int, default=60025)
    args = parser.parse_args()
    asyncio.run(test_the_server(host=args.host, port=args.port))

if __name__ == "__main__":
    do_it()

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