spring boot websocket 的使用

背景

业务上有个新需求，需要后台接口通过websocket的方式将数据推送给前端。
用于提示未处理的消息。

引入依赖

<dependency>
   <groupId>org.springframework.bootgroupId>
    <artifactId>spring-boot-starter-websocketartifactId>
    <version>2.7.1version>
dependency>

使用

配置spring-boot-starter-websocket

创建 WebSocketConfig 配置类

/**
 * WebSocket 配置
 */
@EnableWebSocket   // 开启WebSocket
@Configuration
public class WebSocketConfig implements WebSocketConfigurer {
    @Bean
    public ServerEndpointExporter serverEndpointExporter(){
        return new ServerEndpointExporter();
    }
	
	/**
	*  注入WebSocket的处理器
	*/
    @Autowired
    private WebSocketServerHandler webSocketServerHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
    	// webSocketDemo 为前端请求的地址，前端具体地址组成结构为：ws://ip:接口启动的端口/webSocketDemo
        registry.addHandler(webSocketServer,"webSocketDemo")
                .setAllowedOriginPatterns("*");
    }
}

创建WebSocketServerHandler

@Slf4j
@Component
public class WebSocketServer extends AbstractWebSocketHandler {
    @Autowired
    private WebSocketMessageUtil webSocketMessageUtil;

    /**
     * 连接成功后触发
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("客户端连接成功");
        WebSocketSessionManager.add(session);// 客户端建立连接,将客户端的session对象加入到WebSocketSessionManager的sessionGroup中
        //将连接结果返回给客户端
        webSocketMessageUtil.sendMsg(session,session.getId()+" 连接成功"+ LocalDateTime.now().toString());
    }

    /**
     * 关闭socket连接后触发
     * @param session
     * @param status
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("客户端关闭连接成功");
        WebSocketSessionManager.removeAndClose(session);// 关闭连接,从WebSocketSessionManager的sessionGroup中移除连接对象
    }

    /**
     * 接收客户端发送的文本消息
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        log.info("客户端发送："+message.getPayload());
        //将连接结果返回给客户端
        webSocketMessageUtil.sendMsg(session,session.getId()+" 客户端发送【"+LocalDateTime.now().toString()+"】:"+ message.getPayload());
    }

    /**
     * 接收客户端发送的二进制消息
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        log.info("客户端二进制发送："+message.getPayload());
    }

    /**
     * 异常处理
     * @param session
     * @param exception
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("socket 异常："+exception.getMessage(),exception);
        WebSocketSessionManager.removeAndClose(session); // 出现异常则关闭连接,从WebSocketSessionManager的sessionGroup中移除连接对象
    }
}

创建WebSocketSessionManager

由于每个连接对象都会存在一个session。通过session可以确定每个不同的连接对象。
此方法主要用于对session的管理和操作

@Slf4j
public class WebSocketSessionManager {
    /**
     * 保存连接对象的 session 到集合中
     */
    public  static ConcurrentHashMap<String, WebSocketSession> sessionGroup = new ConcurrentHashMap<>();

    /**
     * 添加 session 到集合中
     *
     * @param session 连接对象的session
     */
    public static void add(WebSocketSession session) {
        // 添加 session
        sessionGroup.put(session.getId(), session);
    }

    /**
     * 从集合中删除 session,会返回删除的 session
     *
     * @param key 通过session.getId()得到
     * @return
     */
    public static WebSocketSession remove(String key) {
        // 删除 session
        return sessionGroup.remove(key);
    }
    
    /**
     * 从集合中删除 session,会返回删除的 session
     *
     * @param session 连接对象的session
     * @return
     */
    public static WebSocketSession remove(WebSocketSession session) {
        // 删除 session
        return sessionGroup.remove(session.getId());
    }
    
    /**
     * 删除并关闭 连接
     *
     * @param key 通过session.getId()得到
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                log.error("关闭出现异常处理",e)
                e.printStackTrace();
            }
        }
    }
    
 	/**
     * 删除并关闭 连接
     *
     * @param session 连接对象的session
     */
    public static void removeAndClose(WebSocketSession session) {
        WebSocketSession session = remove(session);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                log.error("关闭出现异常处理",e)
                e.printStackTrace();
            }
        }
    }
    
    /**
     * 从集合中获得 session
     *
     * @param key 通过session.getId()得到
     * @return
     */
    public static WebSocketSession get(String key) {
        // 获得 session
        return sessionGroup.get(key);
    }
	
}

创建 WebSocketMessageUtil

用于发送消息的工具类

@Service
@Slf4j
public class WsService {

    /**
     * 发送消息给指定客户端
     * @param session 客户端session
     * @param text 发送消息的内容
     * @return
     * @throws IOException
     */
    public synchronized  void sendMsg(WebSocketSession session, String text) throws IOException {
        session.sendMessage(new TextMessage(text));
    }

    /**
     * 发送消息给所有客户端，客户端的session必须在WebSocketSessionManager的sessionGroup中
     * @param text 发送消息的内容
     * @return
     * @throws IOException
     */
    public synchronized void broadcastMsg(String text) throws IOException {
        for (WebSocketSession session: WsSessionManager.sessionGroup.values()) {
            session.sendMessage(new TextMessage(text));
        }
    }

}

前端的html

DOCTYPE HTML>
<html>
<head>
    <title>WebSocket 测试页面title>
head>
<body>
<input id="text" type="text" />
<button onclick="send()">发送消息button>
<button onclick="closeWebSocket()">关闭webScoket连接button>
<div id="message">div>
body>

<script type="text/javascript">
    // 连接到WebSocket的url地址。格式为  ws://ip:接口启动的端口/webSocketDemo
    // 如果 sping boot 配置文件中中配置了server.servlet.context-path ，则格式为： ws://ip:接口启动的端口/server.servlet.context-path的名称/webSocketDemo
    // 此处demo中 sping boot 配置了server.servlet.context-path 为 webSocketDemoApi
    let connectionUrl = "ws://localhost:8003/webSocketDemoApi/webSocketDemo"
    let ws = null;
    //判断当前浏览器是否支持WebSocket
    if ('WebSocket' in window) {
        ws = new WebSocket(connectionUrl);
    }else {
        alert('当前浏览器不支持 websocket')
    }

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

    //连接成功建立的回调方法
    ws.onopen = function(event) {
        console.log("ws调用连接成功回调方法")
        //ws.send("")
    }
    //接收到消息的回调方法
    ws.onmessage = function(message) {
        console.log("接收消息：" + message.data);
        if (typeof(message.data) == 'string') {
            setMessageInnerHTML(message.data);
        }
    }
    //ws连接断开的回调方法
    ws.onclose = function(e) {
        console.log("ws连接断开")
        //console.log(e)
        setMessageInnerHTML("ws close");
    }

    //将消息显示在网页上
    function setMessageInnerHTML(innerHTML) {
        console.log(innerHTML)
        document.getElementById('message').innerHTML += '接收的消息:' + innerHTML + '
';
    }

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


    //发送消息
    function send(msg) {
        if(!msg){
            msg = document.getElementById('text').value;
            document.getElementById('message').innerHTML += "发送的消息:" + msg + '
';
            ws.send(msg);
        }
    }
script>

结果