tornado 实现websocket
server 端
from tornado.websocket import WebSocketClosedError import tornado.ioloop import tornado.web import tornado.websocket import redis import threading import json import time import consts # 全局变量 clients = {} rooms = {} redis_client = redis.StrictRedis(host='localhost', port=6379, db=0) pubsub = redis_client.pubsub() USER_CONNECTION_PREFIX = consts.RedisKeyPrefix + "_user_connection:" # 设置 Redis 键的过期时间(秒) REDIS_KEY_EXPIRATION = 3600 class MainHandler(tornado.web.RequestHandler): def get(self): self.write("Hello, this is a WebSocket server!") class EchoWebSocket(tornado.websocket.WebSocketHandler): def open(self): self.user_id = self.get_argument("user_id") print(f"WebSocket opened for user_id: {self.user_id}") self.room_id = None # 检查用户是否已经连接 if not self.check_and_set_user_connection(): self.close() return clients[self.user_id] = self self.send_msg(data={"msg": "welcome"}) def on_message(self, message): # print(f"Received message from {self.user_id}: {message}") data = json.loads(message) if not isinstance(data, dict): return try: action = data.get("action") message = data.get("message") if action: if action == 'join': self.join_room(data['room']) elif action == 'leave': self.leave_room() elif action == 'message': self.send_message_to_room(data['room'], message) except Exception as e: print(e) def on_close(self): print(f"WebSocket closed for user_id: {self.user_id}") self.leave_room() if self.user_id in clients: del clients[self.user_id] self.clear_user_connection() def join_room(self, room_id): if self.room_id: self.leave_room() self.room_id = room_id if room_id not in rooms: rooms[room_id] = set() rooms[room_id].add(self.user_id) print(f"User {self.user_id} joined room {room_id} rooms:", rooms) def leave_room(self): print("self.room_id:", self.room_id) if self.room_id and self.room_id in rooms: rooms[self.room_id].remove(self.user_id) if not rooms[self.room_id]: del rooms[self.room_id] print(f"User {self.user_id} left room {self.room_id} rooms:", rooms) self.room_id = None def send_message_to_room(self, room_id, message): if room_id in rooms: for user_id in rooms[room_id]: if user_id in clients: clients[user_id].send_msg(message) def send_msg(self, data, status=0, binary=False): try: print("success write_message") self.write_message({"status": status, "errmsg": "成功", "error": "", "data": data}, binary=binary) except WebSocketClosedError: print('write_message fail for socket closed') except Exception as e: print('write_message fail', e) def check_and_set_user_connection(self): key = f"{USER_CONNECTION_PREFIX}{self.user_id}" current_time = time.time() connection_info = {"timestamp": current_time, "ip": self.request.remote_ip} # 使用 Redis SETNX 命令原子性地设置用户连接信息 if redis_client.setnx(key, json.dumps(connection_info)): redis_client.expire(key, REDIS_KEY_EXPIRATION) return True # 如果用户已经连接,检查连接是否仍然有效 existing_connection_info = json.loads(redis_client.get(key)) if current_time - existing_connection_info["timestamp"] < 60: # 如果连接在一分钟内活跃,则拒绝新连接 print(f"User {self.user_id} is already connected.") return False else: # 否则,允许新连接并更新连接信息,并重新设置过期时间 redis_client.setex(key, REDIS_KEY_EXPIRATION, json.dumps(connection_info)) return True def clear_user_connection(self): key = f"{USER_CONNECTION_PREFIX}{self.user_id}" redis_client.delete(key) def broadcast_message(message): message_data = json.loads(message) target = message_data['target'] content = message_data['message'] if target == 'all': for client in clients.values(): client.write_message(content) elif target.startswith('room:'): room_id = target.split(':', 1)[1] if room_id in rooms: for user_id in rooms[room_id]: if user_id in clients: clients[user_id].write_message(content) elif target in clients: clients[target].write_message(content) def redis_listener(ioloop): # 监听 Redis 频道 pubsub.subscribe('broadcast') for item in pubsub.listen(): if item['type'] == 'message': message = item['data'].decode('utf-8') print("message:", message) ioloop.add_callback(broadcast_message, message) def make_app(): return tornado.web.Application([ (r"/", MainHandler), (r"/websocket", EchoWebSocket), ]) if __name__ == "__main__": app = make_app() port = 34567 app.listen(port) print(f"Server is running on http://localhost:{port}") # 获取当前IOLoop实例 ioloop = tornado.ioloop.IOLoop.current() # 启动一个线程来运行后台任务 thread = threading.Thread(target=redis_listener, args=(ioloop,)) thread.daemon = True # 设置为守护线程,确保主线程退出时该线程自动退出 thread.start() # 启动Tornado IOLoop try: ioloop.start() except KeyboardInterrupt: print("Tornado IOLoop stopped by user")
通过redis 发送消息
import redis import json def send_message(target, message): redis_client = redis.StrictRedis(host='localhost', port=6379, db=0) message = { "target": target, # 目标可以是 'all', 'room:<room_id>', 或用户ID "message": message } redis_client.publish('broadcast', json.dumps(message)) if __name__ == "__main__": # # 向所有用户发送消息 # send_message('all', 'Hello, everyone!') # # # 向特定房间发送消息 send_message('room:room1', 'Hello, room1!') # 向特定用户发送消息 send_message('12345', {"status":0})
postman 作为客户端连接
新建websocket 类型连接
输入连接地址点击content 连接 通过send 发送消息
反向代理: 使用 Nginx 或 Apache 作为反向代理来处理 HTTPS 和负载均衡。 Nginx 示例配置:
server {
listen 80;
server_name your_domain.com;
location /websocket {
proxy_pass http://localhost:8888;
proxy_http_version 1.1;
proxy_set_header Upgrade $http_upgrade;
proxy_set_header Connection "upgrade";
proxy_set_header Host $host;
}
}
使用守护进程管理器: 使用 supervisor
、systemd
或其他进程管理工具来确保 Tornado 应用在崩溃后自动重启。 Supervisor 示例配置:
[program:tornado]
command=python /path/to/your/app.py
autostart=true
autorestart=true
stderr_logfile=/var/log/tornado.err.log
stdout_logfile=/var/log/tornado.out.log