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;
    }
}

 

使用守护进程管理器: 使用 supervisorsystemd 或其他进程管理工具来确保 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

 

posted on 2024-07-10 16:25  星河赵  阅读(27)  评论(0编辑  收藏  举报

导航