用 Python 从零构建异步回显服务器

简介

让我们从 0 开始，搭建一个异步服务输出服务器。

套接字

套接字（socket），是不同计算机中实现通信的一种方式，你可以理解成一个接口，它会在客户端和服务端建立连接，一台发送数据，一台接收数据，靠的就是套接字。

阻塞套接字服务器

import socket# socket.AF_INET: IPv4 主机号 + 端口号
# socket.SOCK_STREAM: TCP 连接
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)# socket.SOL_SOCKET: 套接字选项
# socket.SO_REUSEADDR: 允许重用本地地址,避免端口被占用
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('localhost', 8000)
server_socket.bind(server_address) # 绑定地址和端口号
# socket.listen(): 监听连接请求
server_socket.listen()connection, client_address = server_socket.accept()
print(f'我获取了一个连接地址: {client_address}')

启动后，尝试连接服务器

# 启动服务器python server.py# 连接服务器
nc localhost 8000# 输出
我获取了一个连接地址: ('127.0.0.1', 60525)

从套接字读取和写入数据

前面的示例只能接收一次连接，而且无法读取和写入数据，让我们修改一下

import socket# socket.AF_INET: IPv4 主机号 + 端口号
# socket.SOCK_STREAM: TCP 连接
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)# socket.SOL_SOCKET: 套接字选项
# socket.SO_REUSEADDR: 允许重用本地地址,避免端口被占用
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('localhost', 8000)
server_socket.bind(server_address) # 绑定地址和端口号
# socket.listen(): 监听连接请求
server_socket.listen()try:connection, client_address = server_socket.accept()print(f'我获取了一个连接地址: {client_address}')buffer = b''while buffer[-5:] != b'quit\n':data = connection.recv(1024)if not data:breakelse:buffer += dataprint(f'接收数据: {data}')print(f'接收到的所有数据: {buffer}')connection.sendall(buffer)finally:server_socket.close()

这里最重要的就是 recv 方法，它可以从套接字中接收数据，并且写入服务端，然后使用 sendall 方法将接收到的数据，最后全部发送给客户端。

允许多个连接的服务器

如果你尝试使用多个客户端连接上面的服务器，你会发现只有第一个会生效，让我们改写一下，让它可以支持连接多个客户端。

import socket# socket.AF_INET: IPv4 主机号 + 端口号
# socket.SOCK_STREAM: TCP 连接
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)# socket.SOL_SOCKET: 套接字选项
# socket.SO_REUSEADDR: 允许重用本地地址,避免端口被占用
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('localhost', 8000)
server_socket.bind(server_address) # 绑定地址和端口号
# socket.listen(): 监听连接请求
server_socket.listen()connections = []try:while True:connection, client_address = server_socket.accept()print(f'我获取了一个连接地址: {client_address}')buffer = b''while buffer[-5:] != b'quit\n':data = connection.recv(1024)if not data:breakelse:buffer += dataprint(f'接收数据: {data}')print(f'接收到的所有数据: {buffer}')connection.sendall(buffer)finally:server_socket.close()

对比一下代码，你会发现，只不过是把每一个连接放入一个列表中，然后循环遍历接收数据，但是这段示例依然不够优秀，当运行时你会发现，每次第二次连接都会阻塞，都会要第一次连接断开后，第二次连接才会生效。

原因是因为我们使用的是阻塞套接字，想要允许多个客户端连接，我们需要使用非阻塞套接字。

使用非阻塞套接字

python 套接字中设置非阻塞套接字比较简单，一行代码就搞定了setblocking(False)

import socket# socket.AF_INET: IPv4 主机号 + 端口号
# socket.SOCK_STREAM: TCP 连接
server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)# socket.SOL_SOCKET: 套接字选项
# socket.SO_REUSEADDR: 允许重用本地地址,避免端口被占用
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('localhost', 8000)
server_socket.bind(server_address) # 绑定地址和端口号
# socket.listen(): 监听连接请求
server_socket.listen()
server_socket.setblocking(False) # 设置为非阻塞模式connections = []try:while True:try:connection, client_address = server_socket.accept()connection.setblocking(False)print(f'我获取了一个连接地址: {client_address}')connections.append(connection)print(f'当前连接数: {len(connections)}')except BlockingIOError:passfor connection in connections:try:buffer = b''while buffer[-5:] != b'quit\n':data = connection.recv(1024)if not data:breakelse:buffer += dataprint(f'接收数据: {data}')print(f'接收到的所有数据: {buffer}')connection.send(buffer)except BlockingIOError:passfinally:server_socket.close()

尝试运行，会发现 cpu 飙到 100% ，我们需要更好的方法。

使用 selectors 模块优化

Python 的选择器模块是连接操作系统的低层接口，cpu 消耗很少。

import selectors
import socketselector = selectors.DefaultSelector()server_socket = socket.socket()
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('127.0.0.1', 8000)
server_socket.setblocking(False)
server_socket.bind(server_address)
server_socket.listen()selector.register(server_socket, selectors.EVENT_READ)while True:events = selector.select(timeout=1)if len(events) == 0:print('没有事件，等待一会吧')for event, _ in events:event_socket = event.fileobjif event_socket == server_socket:client_socket, client_address = server_socket.accept()print(f'连接来自 {client_address}')client_socket.setblocking(False)selector.register(client_socket, selectors.EVENT_READ)else:data = event_socket.recv(1024)if data:print(f'我获取了数据 {data}')event_socket.send(b'You Got ' + data)else:print('客户端断开连接')selector.unregister(event_socket)event_socket.close()

代码解释

# 创建选择器（自动选择适合当前操作系统的机制）
selector = selectors.DefaultSelector()# 创建服务器 socket 并设置参数
server_socket = socket.socket()
server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)# 注册服务器 socket，监听“可读事件” (EVENT_READ)
# 意思是：当有客户端连接进来，我们就会收到事件通知。
selector.register(server_socket, selectors.EVENT_READ)# 返回所有“发生了事件的 socket 对象”，等待最多 1 秒；
events = selector.select(timeout=1)

因为 selectors 模块是操作系统级别的高效事件通知系统，所以使用这个程序它的 CPU 使用率很低。

Python 自带的 asyncio 默认事件循环基于 selectors 模块

使用 asyncio 构建

回显服务器

import asyncio
import socketasync def echo(connection, loop):while data := await loop.sock_recv(connection, 1024):await loop.sock_sendall(connection, data)async def listen_for_connection(server_socket, loop):while True:connection, address = await loop.sock_accept(server_socket)connection.setblocking(False)print(f'获取到连接: {address}')asyncio.create_task(echo(connection, loop))async def main():server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('localhost', 8000)server_socket.setblocking(False)server_socket.bind(server_address)server_socket.listen()await listen_for_connection(server_socket, asyncio.get_event_loop())asyncio.run(main())

:=是 python3.8 新增的海象运算符

while data := await ...相当于：data = await ...
while data:

解决服务器的错误任务

import asyncio
import socket
import loggingasync def echo(connection, loop):try:while data := await loop.sock_recv(connection, 1024):if data == b'boom\n':raise Exception('Boom!')await loop.sock_sendall(connection, data)except Exception as e:logging.exception(f'服务器报错啦: {e}')finally:connection.close()print('关闭连接')async def listen_for_connection(server_socket, loop):while True:connection, address = await loop.sock_accept(server_socket)connection.setblocking(False)print(f'获取到连接: {address}')asyncio.create_task(echo(connection, loop))async def main():server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)server_address = ('localhost', 8000)server_socket.setblocking(False)server_socket.bind(server_address)server_socket.listen()await listen_for_connection(server_socket, asyncio.get_event_loop())asyncio.run(main())