《SpringBoot2.0 实战》系列-整合WebSocket实现高性能实时消息推送

前言

实时推送消息，在web项目中还是很常用的，比如最近接到一个需求，需要给所有登录的人推送一条广播，就可以使用websocket实现。当然实时消息通知这些也都可以用WebSocket实现。
本文使用的框架是一位大佬封装过的 netty+websocket框架，使用起来和websocket一样，但是性能高很多。有兴趣的可以看源码：https://gitee.com/Yeauty/netty-websocket-spring-boot-starter.git。

如何开始

增加相应依赖

<dependency>
    <groupId>org.yeauty</groupId>
    <artifactId>netty-websocket-spring-boot-starter</artifactId>
    <version>0.9.5</version>
</dependency>

配置下主要配置，其他默认即可

# websocket
netty-websocket:
  port: 8090

websocket端点，内部维护了一个用户map，用于统计在线人数和广播功能。路径参数arg 作为用户id参数。

/**
 * websocket端点
 *
 * @author gourd.hu
 */
@ServerEndpoint(path = "/ws/{arg}", port = "${netty-websocket.port}")
@Component
@Slf4j
public class NioWebSocket {


    private static final String successFlag = "OK";

    /**
     * 用户channel映射关系
     */
    public static final ConcurrentHashMap<String, Session> userChannelMap = new ConcurrentHashMap<>();

    /**
     * 当有新的连接进入时，对该方法进行回调
     * 可做简单的鉴权
     * @param session
     * @param headers
     * @param req
     * @param reqMap
     * @param arg
     * @param pathMap
     */
    @BeforeHandshake
    public void handshake(Session session, HttpHeaders headers,
                          @RequestParam String req,
                          @RequestParam MultiValueMap reqMap,
                          @PathVariable String arg,
                          @PathVariable Map pathMap) {
        session.setSubprotocols("stomp");
//        if (!successFlag.equalsIgnoreCase(req)) {
//            session.close();
//        }
    }

    /**
     * 当有新的WebSocket连接完成时，对该方法进行回调
     *
     * @param session
     * @param headers
     * @param req
     * @param reqMap
     * @param arg
     * @param pathMap
     * @throws IOException
     */
    @OnOpen
    public void onOpen(Session session, HttpHeaders headers,
                       @RequestParam String req,
                       @RequestParam MultiValueMap reqMap,
                       @PathVariable String arg,
                       @PathVariable Map pathMap){
        log.info("new connection");
        // 模拟获取到的userId
        String userId = arg;
        if (userChannelMap.get(userId) == null) {
            userChannelMap.put(userId, session);
        }
        log.info("当前在线人数：" + userChannelMap.size());
    }

    /**
     * 当有WebSocket连接关闭时，对该方法进行回调
     *
     * @param session
     * @param arg
     * @throws IOException
     */
    @OnClose
    public void onClose(Session session,@PathVariable String arg) {
        log.info("one connection closed");
        if (userChannelMap != null) {
            if(userChannelMap.get(arg) != null){
                userChannelMap.remove(arg);
            }
        }
        session.close();
        log.info("当前在线人数：" + userChannelMap.size());
    }

    /**
     * 当有WebSocket抛出异常时，对该方法进行回调
     *
     * @param session
     * @param throwable
     */
    @OnError
    public void onError(Session session,@PathVariable String arg,Throwable throwable) {
        log.info("one connection error");
        log.error(throwable.getMessage(), throwable);
        this.onClose(session,arg);
    }

    /**
     * 当接收到字符串消息时，对该方法进行回调
     *
     * @param session
     * @param message
     */
    @OnMessage
    public void onMessage(Session session, String message) {
        log.info("接收到消息: {}",message);
        sendMessageToUser(session, message);
    }

    /**
     * 当接收到二进制消息时，对该方法进行回调
     * @param session
     * @param bytes
     */
    @OnBinary
    public void onBinary(Session session, byte[] bytes) {
        session.sendBinary(bytes);
    }

    /**
     * 当接收到Netty的事件时，对该方法进行回调
     *
     * @param session
     * @param evt
     */
    @OnEvent
    public void onEvent(Session session, Object evt) {
        if (evt instanceof IdleStateEvent) {
            IdleStateEvent idleStateEvent = (IdleStateEvent) evt;
            switch (idleStateEvent.state()) {
                case READER_IDLE:
                    log.info("Client: " + session.channel().id() + " READER_IDLE 读超时");
                    break;
                case WRITER_IDLE:
                    log.info("Client: " + session.channel().id() + " WRITER_IDLE 写超时");
                    break;
                case ALL_IDLE:
                    log.info("Client: " + session.channel().id() + " ALL_IDLE 总超时");
                    break;
                default:
                    break;
            }
        }
    }


    /**
     * 单点推送给某个人
     *
     * @param session
     * @param msg
     */
    public static final void sendMessageToUser(Session session, String msg) {
        if(session.isOpen() && session.isActive()){
            TextWebSocketFrame tws = new TextWebSocketFrame(msg);
            session.sendText(tws);
        }
    }

    /**
     * 广播给所有在线的人
     *
     * @param msg
     */
    public static final void broadcast(String msg) {
        if (userChannelMap != null) {
            userChannelMap.forEach((key, value) -> sendMessageToUser(value, msg));
        }
    }
}

前端页面

通过随机数生成用户id，用于测试多人连接。

<!DOCTYPE HTML>
<html xmlns:layout="http://www.w3.org/1999/xhtml" layout:decorator="layout">
<head>
    <title>My WebSocket</title>
</head>

<body>
Welcome<br/>
<input id="text" type="text" /><button onclick="send()">Send</button>    <button onclick="closeWebSocket()">Close</button>
<div id="message">
</div>
</body>

<script type="text/javascript">
    let websocket = null;
    // 随机生成模拟用户id
    let userId = Math.ceil(Math.random()*10000);
    // 判断当前浏览器是否支持WebSocket
    if('WebSocket' in window){
        websocket = new WebSocket("ws://localhost:8090/ws/"+userId);
    }
    else{
        alert('Not support websocket')
    }

    // 连接发生错误的回调方法
    websocket.onerror = function(){
        setMessageInnerHTML("error");
    };

    // 连接成功建立的回调方法
    websocket.onopen = function(event){
        setMessageInnerHTML("open");
    }

    // 接收到消息的回调方法
    websocket.onmessage = function(event){
        setMessageInnerHTML(event.data);
    }

    // 连接关闭的回调方法
    websocket.onclose = function(){
        setMessageInnerHTML("close");
    }

    // 监听窗口关闭事件，当窗口关闭时，主动去关闭websocket连接，防止连接还没断开就关闭窗口，server端会抛异常。
    window.onbeforeunload = function(){
        websocket.close();
    }

    // 将消息显示在网页上
    function setMessageInnerHTML(innerHTML){
        document.getElementById('message').innerHTML += innerHTML + '<br/>';
    }

    // 关闭连接
    function closeWebSocket(){
        websocket.close();
    }

    // 发送消息
    function send(){
        let message = document.getElementById('text').value;
        websocket.send(message);
    }
</script>
</html>

广播推送接口

    @PostMapping("/broadcast")
    @ApiOperation(value = "广播消息")
    public void broadcastMsg(String msg) {
        NioWebSocket.broadcast(msg);
    }

测试效果

结语

至此，实时推送消息功能已完成，如果本文有错误的地方，欢迎评论指正。

===============================================

代码均已上传至本人的开源项目
cloud-plus：https://blog.csdn.net/HXNLYW/article/details/104635673