springboot-websocket实现消息的主动推送

前端时间做项目，有一个场景是在后端接收到信号的时候将消息推送给前端，前端在展示给用户。就是后端主动向前端推送消息的功能。这样的场景可以让前端轮询，但是因为消息比较多，比较频繁，所以如果让前端轮询的话，就有可能实时性不太好，而且由于http协议的请求头比较长，进行轮询会浪费大量的带宽，服务器的压力也会巨大。所以就采用的是websocket实现。

1.什么是websocket

WebSocket是一种在单个TCP连接上进行全双工通信的协议。WebSocket使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在WebSocket API中，浏览器和服务器只需要完成一次握手，两者之间就直接可以创建持久性的连接，并进行双向数据传输。

2.原理

轮询是采用固定的时间周期（比如1秒）请求后端接口，采用http协议进行通信，浪费资源。
而websocket是在首次建立连接的时候采用http协议，之后这个连接通道就一直存在着，直到自动关闭或者是手动关闭。当获取 websocket 连接后，就可以通过 send() 方法来发送数据，并通过 onmessage 事件来接收数据。

3.代码

首先引入pom依赖

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

整个springboot项目的结构：

WebSocketConfig.java

package cn.fzkj.bootquick.config;

import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Component
public class WebSocketConfig {

    /**
     * 使用spring boot时，使用的是spring-boot的内置容器，
     * 如果要使用WebSocket，需要注入ServerEndpointExporter
     *
     * @return
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

SocketUser.java

package cn.fzkj.bootquick.websocket;

import lombok.Data;

import javax.websocket.Session;

@Data
public class SocketUser {
    private Session session;
    private String uid;
}

这个类主要是用来封装连接的信息。
接下来才是硬类：

package cn.fzkj.bootquick.websocket;


import cn.fzkj.bootquick.dto.APIResultCode;
import cn.fzkj.bootquick.dto.JsonDtoWrapper;
import cn.fzkj.bootquick.dto.MessageInfo;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.auth0.jwt.internal.org.apache.commons.lang3.StringUtils;
import com.google.gson.Gson;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * @ServerEndpoint 注解是一个类层次的注解，它的功能主要是将目前的类定义成一个websocket服务器端,
 * 注解的值将被用于监听用户连接的终端访问URL地址,客户端可以通过这个URL来连接到WebSocket服务器端
 */
@Slf4j
@ServerEndpoint(value = "/mysocket/{userId}")
@Component
public class WebSocketService {
    //静态变量，用来记录当前在线连接数。应该把它设计成线程安全的。
    private static int onlineCount = 0;

    //concurrent包的线程安全Set，用来存放每个客户端对应的MyWebSocket对象。
    // 若要实现服务端与单一客户端通信的话，可以使用Map来存放，其中Key可以为用户标识
    public static Map<String, SocketUser> userMap = new HashMap<>();
    private static Map<String, String> sessionUID = new ConcurrentHashMap<>();


    /**
     * 连接建立成功调用的方法
     *
     * @param session 可选的参数。session为与某个客户端的连接会话，需要通过它来给客户端发送数据
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("userId") String userId) {
        System.out.println("连接接收到的userId："+userId);
        SocketUser socketUser = new SocketUser();
        socketUser.setSession(session);
        socketUser.setUid(userId);
        userMap.put(socketUser.getUid(),socketUser);
        sessionUID.put(session.getId(), socketUser.getUid());
        onlineCount ++;
        log.info("有新连接加入！当前在线人数为" + getOnlineCount());
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public synchronized void onClose(Session session) {
        String uid = sessionUID.get(session.getId());

        if (userMap.containsKey(uid)) {
            userMap.remove(uid);
            sessionUID.remove(session.getId());
            onlineCount --;
            log.info("success remove on connect=========================================>");
        }

        log.info("有一连接正常关闭！当前在线人数为" + uid + ":" + getOnlineCount());
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message 客户端发送过来的消息
     * @param session 可选的参数
     */
    @OnMessage
    public void onMessage(String message, Session session) {
        if(StringUtils.isNotBlank(message) && message.charAt(0) == '{'){
            //解析发送的报文
            JSONObject jsonObject = JSONObject.parseObject(message);
            // 判断接收到的消息
            String msg = jsonObject.getString("msg");
            System.out.println("从客户端接收到的消息是："+msg);
            if ("heartbeat".equals(msg)){
                JsonDtoWrapper jo = new JsonDtoWrapper();
                jo.setCodeMsg(APIResultCode.HERAT_RESPONSE);
                try {
                    session.getBasicRemote().sendText(JSON.toJSONString(jo));
                } catch (Exception e) {
                    log.error("heart beat:", e);
                }
            }
            // 获取接收者的id
            String recevierID = jsonObject.getString("recevierID");
            if (recevierID != null && !recevierID.equals("")){
                // 有接收者
                try{
                    SocketUser socketUser = userMap.get(recevierID);
                    if (null != socketUser){
                        socketUser.getSession().getBasicRemote().sendText(message);
                        log.info("success send");
                    }
                }catch (Exception e){
                    log.error("socket sendMessage", e);
                }
            }else {
                // 没有接收者
                log.error("no recevier");
            }
        }else{
            try {
                session.getBasicRemote().sendText(message);
            } catch (Exception e) {
                log.error("heart beat:", e);
            }
        }
        log.info("来自客户端的消息010:" + message);
    }

    /**
     * 发生错误时调用
     *
     * @param session
     * @param error
     */
    @OnError
    public synchronized void onError(Session session, Throwable error) {
        log.info("发生错误");
        String uid = sessionUID.get(session.getId());

        if (userMap.containsKey(uid)) {
            userMap.remove(uid);
            sessionUID.remove(session.getId());
            onlineCount --;
            log.info("success remove on connect=========================================>");
        }

        log.info("有一连接错误关闭！当前在线人数为" + uid + ":" + getOnlineCount());

        error.printStackTrace();
    }

    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static JsonDtoWrapper sendMessage(MessageInfo messageInfo) {
        JsonDtoWrapper jo = new JsonDtoWrapper();
        try {
            SocketUser socketUser = userMap.get(messageInfo.getReceiveID());

            if (null != socketUser) {
                socketUser.getSession().getBasicRemote().sendText(new Gson().toJson(messageInfo));
                log.info("success send");
            } else {
                jo.setCodeMsg(APIResultCode.NOT_ONLINE);
                log.info("socket not online");
                return jo;
            }
            log.info("uid:" + messageInfo.getReceiveID());
            log.info("message:" + new Gson().toJson(messageInfo));
            jo.setCodeMsg(APIResultCode.SUCCESS);
            return jo;
        } catch (Exception e) {
            log.error("socket sendMessage", e);
            jo.setCodeMsg(APIResultCode.FAILURE);
            jo.setExtraMsg(e.toString());
            return jo;
        }
    }

}

onMessage方法为收到消息的时候调用的方法，具体业务自行扩展。

前端代码：

<!DOCTYPE html>
<html>
<head>
    <title>Java后端WebSocket的Tomcat实现</title>
</head>
<body>
    Welcome<br/><input id="text" type="text"/>
    <button onclick="send()">发送消息</button>
    <hr/>
    <button onclick="closeWebSocket()">关闭WebSocket连接</button>
    <hr/>
    <div id="message"></div>
</body>

<script type="text/javascript">
    var websocket = null;
    //判断当前浏览器是否支持WebSocket
    if ('WebSocket' in window) {
        websocket = new WebSocket("ws://localhost:8081/mysocket/10");
    }
    else {
        alert('当前浏览器 Not support websocket')
    }

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

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

    //接收到消息的回调方法
    websocket.onmessage = function (event) {
		console.log(event.data);
		var msg = JSON.parse(event.data);
        setMessageInnerHTML(msg['msg']+"&#13;");
    }

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

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

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

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

    //发送消息
    function send() {
        var message = document.getElementById('text').value;
        websocket.send('{"recevierID":"11","msg":"'+message+'"}');
		setMessageInnerHTML(message+"&#13;");
    }
</script>
</html>

效果展示：

关于我使用到的工具类：
JsonDtoWrapper.java

package cn.fzkj.bootquick.dto;

import lombok.Data;

@Data
public class JsonDtoWrapper {
    private String code;
    private String msg;
    private String extraMsg;

    public void setCodeMsg(APIResultCode co) {
        this.code = co.getCode();
        this.msg = co.getMsg();
        this.extraMsg = co.getExtraMsg();
    }
}

---

持续更新~~