asyncio 简单的 websocket 服务器
pip install websockets~=9.1
一个代理服务器，将所有客户端发来的信息，转发到其他服务器。
import asyncio
import json
import logging
import websockets
import websockets.legacy.client
MODULE_LOGGER = logging.getLogger(__name__)
# 需要连接的其他服务器地址
SERVER_TYPES = {
'audio_server': {'addr': 'ws://localhost:9999', 'msgs': None},
'touch_server': {'addr': 'ws://localhost:8888', 'msgs': None}
}
# 转发服务器，用来接收用户请求
class Dispatcher:
"""A server, would wait for clients to connect, and put the client's message into different Queue.
"""
def __init__(self, ip="localhost", port=9000):
self._host = ip
self._port = port
self.clients = []
self._loop = asyncio.get_event_loop()
self._loop.create_task(self.run())
async def run(self):
"""create a server
"""
self.server = await websockets.serve(self.dispatch, self._host, self._port)
async def register_client(self, ws):
"""register connected client
:param ws: websocket client
"""
if ws not in self.clients:
self.clients.append(ws)
async def unregister_client(self, ws):
"""unregister a disconnected client.
:param ws: websocket client.
:return:
"""
if ws in self.clients:
self.clients.remove(ws)
async def dispatch(self, websocket, path):
"""According to client's request's server type, put the request message into specific queue.
:param websocket:
:param path:
:return:
"""
await self.register_client(websocket)
try:
async for message in websocket: # recv message from a websocket client.
msg_dict = self.parse_message(message)
server_type = msg_dict.get('server')
if server_type in SERVER_TYPES:
await SERVER_TYPES[server_type]['msgs'].put((websocket, message)) # put message info into queue
MODULE_LOGGER.info('Add one message to queue.')
except Exception as exc:
MODULE_LOGGER.error(exc)
finally:
await self.unregister_client(websocket)
def parse_message(self, message):
"""parse the str message into dict type.
message should be like this:
msg = {
'server': 'audio_server', # touch_server, audio_server,...
'method': 'play_audio',
...
}
:param message:
:return:
"""
try:
message = json.loads(message)
return message
except Exception as exe:
MODULE_LOGGER.info(exe)
return {}
# 客户端，用来连接其他的服务器，然后将用户的请求，发送到这个服务器
class Client:
"""local client that connect to a specific server, and send/recv message to/from that server."""
def __init__(self, server, addr, msgs):
"""
:param server: a Dispatcher instance.
:param addr: the server's uri address.
:param msgs: the message queue belong to the server type.
"""
self.server = server
self._remote_addr = addr
self._msgs = msgs
self._loop = asyncio.get_event_loop()
self._loop.create_task(self.run())
async def run(self):
await self._connect()
self._loop.create_task(self.send_message())
self._loop.create_task(self.recv_message())
async def _connect(self):
try:
self.client = await websockets.legacy.client.connect(self._remote_addr)
except:
MODULE_LOGGER.error(f"Cannot connect to: {self._remote_addr}")
def get_user_client(self, address):
"""According to the argument: address, get the user client object
:param list address: client's address
:return:
"""
for client in self.server.clients:
if list(client.remote_address) == address:
return client
return None
async def recv_message(self):
"""recv feedback from server, then send the feedback to user client."""
try:
while True:
msg = await self.client.recv()
msg_json = json.loads(msg)
to_client = msg_json['from'] # the server should send the 'from' keyword.
MODULE_LOGGER.info(f'recv msg :{msg}')
client = self.get_user_client(to_client)
if client:
await client.send(msg)
except:
await self._connect() # reconnect server
self._loop.create_task(self.send_message())
async def send_message(self):
"""Send message to the server(self._remote_addr)
:param addr:
:param msg_queue:
:return:
"""
while True:
client_info = await self._msgs.get() # tuple(client_obj, message)
msg = self.add_msg_source(*client_info)
try:
await self.client.send(msg)
MODULE_LOGGER.info(f"Send message to {self._remote_addr}")
except Exception as exc:
MODULE_LOGGER.error(f"Connect to server: {self._remote_addr} error!")
MODULE_LOGGER.error(exc)
# local client can't connect to server, then feedback to user client.
user_client = client_info[0]
if user_client.open:
await user_client.send("Error send message ...") # Todo. The sent error data should be re-considered
await self._connect() # try to reconnect the server.
self._loop.create_task(self.send_message())
break
def add_msg_source(self, client, msg):
"""add the message a from source, so that we can recognize the message from which client.
:param client:
:param msg:
:return:
"""
msg_json = json.loads(msg)
msg_json['from'] = client.remote_address
return json.dumps(msg_json)
def init_message_queue(maxsize=5):
""" initialize SERVER_TYPES, create asyncio.Queue() for each type of server to store message.
:return:
"""
global SERVER_TYPES
for server in SERVER_TYPES:
SERVER_TYPES[server]['msgs'] = asyncio.Queue(maxsize=maxsize)
if __name__ == "__main__":
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
init_message_queue()
server = Dispatcher()
for key in SERVER_TYPES:
Client(server, SERVER_TYPES[key]['addr'], SERVER_TYPES[key]['msgs'])
loop.run_forever()
创建其他服务器：
import asyncio
import websockets
async def hello(websocket, path):
try:
print("Client connected...")
async for message in websocket:
print(message)
await websocket.send(message)
print(f"sent {message}")
except:
print('Client disconnected...')
async def main(port):
await websockets.serve(hello, 'localhost', port)
asyncio.get_event_loop().create_task(main(9999))
asyncio.get_event_loop().create_task(main(8888))
asyncio.get_event_loop().run_forever()
模拟用户连接：
import os
import json
import asyncio
import websockets
async def recv_from(ws):
while True:
msg = await ws.recv()
print(f'recv from {msg}')
msg = {
'server': 'audio_server', # touch_server, audio_server
'method': 'play',
'num': 0,
}
msg2 = {
'server': 'touch_server', # touch_server, audio_server
'method': 'play',
'num': 0,
}
async def hello(msg):
async with websockets.connect('ws://localhost:9000') as websocket:
asyncio.get_event_loop().create_task(recv_from(websocket))
while True:
await asyncio.sleep(1)
msg['num'] += 1
str_msg = json.dumps(msg)
await websocket.send(str_msg)
print("send to", msg)
asyncio.get_event_loop().create_task(hello(msg))
# asyncio.get_event_loop().create_task(hello(msg2))
asyncio.get_event_loop().run_forever()
上面共有三个脚本，执行顺序是：2，1，3
