Spring Boot 整合 WebSocket

一. WebSocket 介绍

项目学习案例，仅供参考！如有更好的方案和想法，欢迎互相交流（关于我）！
WebSocket 是 HTML5 开始提供的一种在单个 TCP 连接上进行全双工通讯的协议。

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

在 WebSocket API 中，浏览器和服务器只需要做一个握手的动作，然后，浏览器和服务器之间就形成了一条快速通道。两者之间就直接可以数据互相传送。

websocket

浏览器通过 JavaScript 向服务器发出建立 WebSocket 连接的请求，连接建立以后，客户端和服务器端就可以通过 TCP 连接直接交换数据。
当你获取 Web Socket 连接后，你可以通过 send() 方法来向服务器发送数据，并通过 onmessage 事件来接收服务器返回的数据。
以下 API 用于创建 WebSocket 对象。

var webSocket = new WebSocket(url, [protocol] );

以上代码中的第一个参数 url, 指定连接的 URL。第二个参数 protocol 是可选的，指定了可接受的子协议。

二. 项目结构

sw-websocket                                # WebSocket Demo
    - src                                   # 源文件目录
        -- main                             # 主目录
            -- java                         # Java 源文件目录
                -- com.lmay.websocket       # Java 包路径
                    -- config               # 项目配置类
                    -- controller           # 控制器
                    -- exception            # 自定义异常
                    -- handler              # 处理器
                    -- service              # 服务层
                        -- impl             # 服务实现类
                    -- task                 # Spring 定时任务
                    WebSocketApplication    # 应用启动
            -- resources                    # 项目资源目录
                -- static                   # 静态资源
                -- templates                # HTML模版
                application.yml             # 项目配置文件
                log4j2.xml                  # 日志配置文件
        -- test                             # 测试目录
    pom.xml                                 # Maven 资源库配置文件

三. 项目架构

JDK 8
Spring Boot 2
Maven
Guava
Gson
lombok
log4j2
Spring Scheduled
…

四. 源码实现（核心源码）

WebSocket.java [WebSocket 处理类]

package com.lmay.websocket.handler;

import com.google.common.base.Strings;
import com.lmay.common.common.Response;
import com.lmay.common.exception.CommonException;
import com.lmay.common.utils.GsonUtils;
import com.lmay.websocket.service.WebSocketService;
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.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;

/**
 * <pre>
 *   ┏┓　　　┏┓
 * ┏┛┻━━━┛┻┓
 * ┃　　　　　　　┃
 * ┃　　　━　　　┃
 * ┃　┳┛　┗┳　┃
 * ┃　　　　　　　┃
 * ┃　　　┻　　　┃
 * ┃　　　　　　　┃
 * ┗━┓　　　┏━┛
 * 　　┃　　　┃神兽保佑
 * 　　┃　　　┃代码无BUG！
 * 　　┃　　　┗━━━┓
 * 　　┃　　　　　　　┣┓
 * 　　┃　　　　　　　┏┛
 * 　　┗┓┓┏━┳┓┏┛
 * 　　　┃┫┫　┃┫┫
 * 　　　┗┻┛　┗┻┛
 * </pre>
 * -- Websocket Handler
 *
 * @author lmay.Zhou
 * @date 2018/12/5 18:11 星期三
 * @qq 379839355
 * @email lmay@lmaye.com
 */
@Slf4j
@Component
@ServerEndpoint(value = "/ws/{userId}")
public class WebSocket {
    public static WebSocketService webSocketService;

    private static ConcurrentHashMap<String, WebSocket> webSocket = new ConcurrentHashMap<>();

    /**
     * 用户编号
     */
    private String userId;

    /**
     * session
     */
    private Session session;

    /**
     * 建立连接调用的方法
     *
     * @param userId  用户ID
     * @param session Session
     */
    @OnOpen
    public void onOpen(@PathParam(value = "userId") String userId, Session session) {
        try {
            this.session = session;
            this.userId = userId;
            webSocket.put(this.userId, this);
            long count = webSocket.size();
            log.info("用户[" + this.userId + "]加入连接在线总数[" + count + "]");
            sendMessage(GsonUtils.toJson(Response.success("连接成功")));
        } catch (CommonException e) {
            sendMessage(GsonUtils.toJson(Response.failed(e.getError())));
        }
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        webSocket.remove(this.userId);
        long count = WebSocket.webSocket.size();
        log.info("用户[" + this.userId + "]关闭连接在线总数[" + count + "]");
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param json 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String json, Session session) {
        if (!Strings.isNullOrEmpty(json)) {
            Response<Map<String, Object>> result = webSocketService.selectUserByUserId(userId);
            result.getData().put("msg", json);
            sendMessage(GsonUtils.toJson(result));
        } else {
            sendMessage(GsonUtils.toJson(Response.failed("连接成功")));
        }
    }

    /**
     * 连接错误
     *
     * @param session Session
     * @param e       Throwable
     */
    @OnError
    public void onError(Session session, Throwable e) {
        log.error("websocket IO异常", e);
    }


    /**
     * 发送消息
     *
     * @param message 消息
     */
    private void sendMessage(String message) {
        try {
            this.session.getBasicRemote().sendText(message);
        } catch (IOException e) {
            log.error("发送用户[" + this.userId + "]消息[" + message + "]失败", e);
        }
    }

    /**
     * 给用户发送消息
     *
     * @param userId  用户编号
     * @param message 消息
     * @throws IOException IOException
     */
    private void sendMessage(String userId, String message) {
        if (webSocket.get(userId) != null) {
            webSocket.get(userId).sendMessage(message);
        }
    }

    /**
     * 发送信息
     *
     * @param userIds 用户
     * @param message 消息
     */
    public static void sendMessage(Set<String> userIds, String message) {
        for (String userId : userIds) {
            webSocket.get(userId).sendMessage(message);
        }
    }

    /**
     * 发送信息给所有人
     *
     * @param message 消息
     */
    public static void sendMessageAll(String message) {
        for (String userId : webSocket.keySet()) {
            webSocket.get(userId).sendMessage(message);
        }
    }

    /**
     * 获取在线人数
     *
     * @return long
     */
    public static synchronized long getCount() {
        return webSocket.size();
    }
}

WebSocketConfig.java [WebSocket 配置类]

package com.lmay.websocket.config;

import com.lmay.websocket.handler.WebSocket;
import com.lmay.websocket.service.WebSocketService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

/**
 * <pre>
 *   ┏┓　　　┏┓
 * ┏┛┻━━━┛┻┓
 * ┃　　　　　　　┃
 * ┃　　　━　　　┃
 * ┃　┳┛　┗┳　┃
 * ┃　　　　　　　┃
 * ┃　　　┻　　　┃
 * ┃　　　　　　　┃
 * ┗━┓　　　┏━┛
 * 　　┃　　　┃神兽保佑
 * 　　┃　　　┃代码无BUG！
 * 　　┃　　　┗━━━┓
 * 　　┃　　　　　　　┣┓
 * 　　┃　　　　　　　┏┛
 * 　　┗┓┓┏━┳┓┏┛
 * 　　　┃┫┫　┃┫┫
 * 　　　┗┻┛　┗┻┛
 * </pre>
 * -- Websocket Config
 *
 * @author lmay.Zhou
 * @date 2018/12/5 18:05 星期三
 * @qq 379839355
 * @email lmay@lmaye.com
 */
@Configuration
public class WebSocketConfig {
    /**
     * 首先要注入ServerEndpointExporter，这个bean会自动注册使用了@ServerEndpoint注解声明的Websocket endpoint。
     * 要注意，如果使用独立的servlet容器，而不是直接使用springboot的内置容器，就不要注入ServerEndpointExporter，因为它将由容器自己提供和管理。
     *
     * @return ServerEndpointExporter
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

    /**
     * 因 SpringBoot WebSocket 对每个客户端连接都会创建一个 WebSocketServer（@ServerEndpoint 注解对应的） 对象，Bean 注入操作会被直接略过，因而手动注入一个全局变量
     *
     * @param webSocketService webSocketService
     */
    @Autowired
    public void setMessageService(WebSocketService webSocketService) {
        WebSocket.webSocketService = webSocketService;
    }
}

WebSocketTask.java [WebSocket 定时任务类 - 模拟服务器向前端推送消息]

package com.lmay.websocket.task;

import com.lmay.websocket.service.WebSocketService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;

/**
 * <pre>
 *   ┏┓　　　┏┓
 * ┏┛┻━━━┛┻┓
 * ┃　　　　　　　┃
 * ┃　　　━　　　┃
 * ┃　┳┛　┗┳　┃
 * ┃　　　　　　　┃
 * ┃　　　┻　　　┃
 * ┃　　　　　　　┃
 * ┗━┓　　　┏━┛
 * 　　┃　　　┃神兽保佑
 * 　　┃　　　┃代码无BUG！
 * 　　┃　　　┗━━━┓
 * 　　┃　　　　　　　┣┓
 * 　　┃　　　　　　　┏┛
 * 　　┗┓┓┏━┳┓┏┛
 * 　　　┃┫┫　┃┫┫
 * 　　　┗┻┛　┗┻┛
 * </pre>
 * -- WebSocket Task
 *
 * @author lmay.Zhou
 * @date 2018/12/6 14:06 星期四
 * @qq 379839355
 * @email lmay@lmaye.com
 */
@Slf4j
@Component
public class WebSocketTask {
    private final WebSocketService webSocketService;

    public WebSocketTask(WebSocketService webSocketService) {
        this.webSocketService = webSocketService;
    }

    /**
     * 服务器主动向客户端推送消息
     */
    @Scheduled(cron = "0/30 * * * * ?")
    public void initiativeSendMsg() {
        log.info("服务器主动向客户端推送消息: {}", webSocketService.initiativeSendMsg());
    }
}

index.ftl [WebSocket 页面]

<!DOCTYPE html>
<html lang="zh">
    <head>
        <meta charset="UTF-8">
        <title>WebSocket Demo</title>
    </head>

    <body>
        Welcome WebSocket Demo:<br/>
        <input id="msg" name="msg" type="text" />
        <button onclick="send();">Send</button>
        <button onclick="closeWebSocket();">Close</button>
        <div id="message">

        </div>
    </body>

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

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

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

        //接收到消息的回调方法
        websocket.onmessage = function (event) {
            setMessageInnerHTML("Receive Message: " + 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() {
            var message = document.getElementById('msg').value;
            websocket.send(message);
        }
    </script>
</html>

五. WebSocket 示例

WebSocket 服务连接[ Receive Message: {“code”:200,”msg”:”操作成功”,”data”:”连接成功”} ];
客户端发送消息[ Receive Message: {“code”:200,”msg”:”操作成功”,”data”:{“userName”:”WebSocket Test”,”msg”:”o(ﾟДﾟ)っ啥！”,”userId”:”10000”}} ];
服务端向客户端发送消息[ Receive Message: {“code”:200,”msg”:”操作成功”,”data”:”欢迎使用 WebSocket 服务！”} ];