FdroidRepository/src/main/java/cn/seecoder/fdroidrepository/websocket/WebSocketService.java

package cn.seecoder.fdroidrepository.websocket;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.LinkedList;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/message/{userId}")
public class WebSocketService {
    // 用户连接池存储
    private static ConcurrentHashMap<Integer, Session> SESSION_POLL = new ConcurrentHashMap<>();

    // 用户历史消息保存，保存未登录时存储的消息
    private static ConcurrentHashMap<Integer, Queue<String>> HISTORY_MESSAGE = new ConcurrentHashMap<>();

    private static final Logger logger = LoggerFactory.getLogger(WebSocketService.class);

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") Integer userId) {
        if (!SESSION_POLL.contains(userId)) {
            SESSION_POLL.put(userId, session);
        }
        // 历史数据
        if (HISTORY_MESSAGE.contains(userId)) {
            Queue<String> history = HISTORY_MESSAGE.get(userId);
            if (history.size() > 0) {
                for (String text : history) {
                    try {
                        session.getBasicRemote().sendText(text);
                    } catch (IOException exception) {
                        logger.error("发送历史消息失败，userId: {}, message: {}, exception: {}", userId, text, exception.getMessage());
                    }
                }
            }
            HISTORY_MESSAGE.remove(userId);
        }
    }

    @OnClose
    public void onClose(@PathParam("userId") Integer userId) {
        SESSION_POLL.remove(userId);
    }

    @OnMessage
    public void onMessage(String message) {
        logger.info("WebSocket 收到信息：{}", message);
        if (message.indexOf("ToUser:") == 0) {
            String targetId = message.substring(message.indexOf("ToUser:") + 7, message.indexOf(";"));
            message = message.substring(message.indexOf(";"));
            Integer target = Integer.valueOf(targetId);
            if (SESSION_POLL.contains(target)) {
                try {
                    SESSION_POLL.get(target).getBasicRemote().sendText(message);
                } catch (IOException exception) {
                    logger.error("发送历史消息失败, message: {}, exception: {}", message, exception.getMessage());
                }
            } else {
                // 用户未登录，加入历史记录列表
                if (HISTORY_MESSAGE.contains(target)) {
                    HISTORY_MESSAGE.get(target).add(message);
                } else {
                    Queue<String> queue = new LinkedList<>();
                    queue.add(message);
                    HISTORY_MESSAGE.put(target, queue);
                }
            }
        }
    }

    @OnError
    public void onError(Session session, Throwable error) {
        logger.error("WebSocket Server出现错误，{}", error.getMessage());
    }
}