sanguo_vnpy_v2/sanguo_web/static/js/websocket.js

/**
 * WebSocket 连接管理
 * 处理实时数据推送、重连机制
 */

const WS_URL = `ws://${window.location.host}/ws`;

/**
 * WebSocket 客户端类
 */
class WebSocketClient {
    constructor() {
        this.ws = null;
        this.reconnectInterval = null;
        this.reconnectDelay = 3000; // 3秒后重连
        this.maxReconnectAttempts = 10;
        this.reconnectAttempts = 0;
        this.isConnected = false;
        this.subscriptions = new Set();
        this.messageHandlers = new Map();
    }

    /**
     * 连接 WebSocket
     */
    connect(token) {
        if (this.ws && this.ws.readyState === WebSocket.OPEN) {
            console.log('WebSocket already connected');
            return;
        }

        const url = token ? `${WS_URL}?token=${token}` : WS_URL;

        this.ws = new WebSocket(url);

        this.ws.onopen = () => {
            console.log('WebSocket connected');
            this.isConnected = true;
            this.reconnectAttempts = 0;
            this.clearReconnectInterval();

            // 重新订阅之前的内容
            if (this.subscriptions.size > 0) {
                this.subscribe(Array.from(this.subscriptions));
            }

            // 触发连接事件
            this.emit('connected');
        };

        this.ws.onmessage = (event) => {
            try {
                const message = JSON.parse(event.data);
                this.handleMessage(message);
            } catch (e) {
                console.error('Failed to parse WebSocket message:', e);
            }
        };

        this.ws.onclose = () => {
            console.log('WebSocket disconnected');
            this.isConnected = false;
            this.emit('disconnected');

            // 尝试重连
            this.scheduleReconnect(token);
        };

        this.ws.onerror = (error) => {
            console.error('WebSocket error:', error);
        };
    }

    /**
     * 断开连接
     */
    disconnect() {
        this.clearReconnectInterval();
        if (this.ws) {
            this.ws.close();
            this.ws = null;
        }
        this.isConnected = false;
    }

    /**
     * 安排重连
     */
    scheduleReconnect(token) {
        if (this.reconnectInterval) {
            return;
        }

        if (this.reconnectAttempts >= this.maxReconnectAttempts) {
            console.error('Max reconnect attempts reached');
            this.emit('reconnect-failed');
            return;
        }

        this.reconnectAttempts++;
        console.log(`Scheduling reconnect attempt ${this.reconnectAttempts}`);

        this.reconnectInterval = setTimeout(() => {
            this.reconnectInterval = null;
            this.connect(token);
        }, this.reconnectDelay);
    }

    /**
     * 清除重连定时器
     */
    clearReconnectInterval() {
        if (this.reconnectInterval) {
            clearTimeout(this.reconnectInterval);
            this.reconnectInterval = null;
        }
    }

    /**
     * 处理消息
     */
    handleMessage(message) {
        const { type, data } = message;

        // 根据消息类型调用对应的处理器
        if (this.messageHandlers.has(type)) {
            this.messageHandlers.get(type).forEach(handler => handler(data));
        }

        // 通用消息处理器
        if (this.messageHandlers.has('*')) {
            this.messageHandlers.get('*').forEach(handler => handler(message));
        }
    }

    /**
     * 发送消息
     */
    send(message) {
        if (!this.isConnected) {
            console.warn('WebSocket not connected, message not sent');
            return false;
        }

        try {
            this.ws.send(JSON.stringify(message));
            return true;
        } catch (e) {
            console.error('Failed to send WebSocket message:', e);
            return false;
        }
    }

    /**
     * 订阅消息类型
     */
    subscribe(types) {
        types = Array.isArray(types) ? types : [types];

        types.forEach(type => {
            this.subscriptions.add(type);
        });

        if (this.isConnected) {
            this.send({
                action: 'subscribe',
                types: types
            });
        }
    }

    /**
     * 取消订阅
     */
    unsubscribe(types) {
        types = Array.isArray(types) ? types : [types];

        types.forEach(type => {
            this.subscriptions.delete(type);
        });

        if (this.isConnected) {
            this.send({
                action: 'unsubscribe',
                types: types
            });
        }
    }

    /**
     * 订阅品种
     */
    subscribeSymbol(symbols) {
        symbols = Array.isArray(symbols) ? symbols : [symbols];

        if (this.isConnected) {
            this.send({
                action: 'subscribe_symbol',
                symbols: symbols
            });
        }
    }

    /**
     * 取消订阅品种
     */
    unsubscribeSymbol(symbols) {
        symbols = Array.isArray(symbols) ? symbols : [symbols];

        if (this.isConnected) {
            this.send({
                action: 'unsubscribe_symbol',
                symbols: symbols
            });
        }
    }

    /**
     * 注册消息处理器
     */
    on(type, handler) {
        if (!this.messageHandlers.has(type)) {
            this.messageHandlers.set(type, new Set());
        }
        this.messageHandlers.get(type).add(handler);
    }

    /**
     * 移除消息处理器
     */
    off(type, handler) {
        if (this.messageHandlers.has(type)) {
            this.messageHandlers.get(type).delete(handler);
        }
    }

    /**
     * 触发事件
     */
    emit(type, data) {
        if (this.messageHandlers.has(type)) {
            this.messageHandlers.get(type).forEach(handler => handler(data));
        }
    }

    /**
     * 获取连接状态
     */
    getConnectionState() {
        if (!this.ws) {
            return 'disconnected';
        }

        switch (this.ws.readyState) {
            case WebSocket.CONNECTING:
                return 'connecting';
            case WebSocket.OPEN:
                return 'connected';
            case WebSocket.CLOSING:
                return 'closing';
            case WebSocket.CLOSED:
                return 'closed';
            default:
                return 'unknown';
        }
    }
}

// 创建全局 WebSocket 客户端实例
const wsClient = new WebSocketClient();