Spring WebSocket实现消息订阅

1、目标

客户端通过WebSocket长连接与服务端实现消息订阅，并且根据连接参数订阅对应的消息！

2、服务端

2.1 核心依赖坐标

我这里是使用的springboot启动器依赖

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

2.2 WebSocket工具类

主要实现客户端到服务端WebSocketSession的增删改查

import lombok.extern.slf4j.Slf4j;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;


@Slf4j
public class WsSessionUtil {
    /**
     * 保存连接 session 的地方
     */
    private static final ConcurrentMap<String, WebSocketSession> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 添加 session
     *
     * @param key
     */
    public static void add(String key, WebSocketSession session) {
        // 添加 session
        SESSION_POOL.put(key, session);
    }

    /**
     * 删除 session,会返回删除的 session
     *
     * @param key
     * @return
     */
    public static WebSocketSession remove(String key) {
        // 删除 session
        return SESSION_POOL.remove(key);
    }

    /**
     * 删除并同步关闭连接
     *
     * @param key
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                log.error("websocket关闭失败。。。");
            }
        }
    }

    /**
     * 获得 session
     *
     * @param key
     * @return
     */
    public static WebSocketSession get(String key) {
        // 获得 session
        return SESSION_POOL.get(key);
    }

    public static ConcurrentMap<String, WebSocketSession> get(){
        return SESSION_POOL;
    }

}

2.3 WebSocket消息处理类

继承AbstractWebSocketHandler方法，实现连接创建，关闭，以及消息的发送和异常处理方法

import com.paratera.console.biz.utils.WsSessionUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.AbstractWebSocketHandler;

/**
 * ws消息处理类
 */
@Component
@Slf4j
public class WarningWsHandler extends AbstractWebSocketHandler {

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        log.info("建立ws连接");
        WsSessionUtil.add(session.getId(),session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        log.info("发送文本消息");
        // 获得客户端传来的消息
        String payload = message.getPayload();
        log.info("server 接收到消息 " + payload);
        session.sendMessage(new TextMessage(payload));
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        log.error("异常处理");
        WsSessionUtil.removeAndClose(session.getId());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        log.info("关闭ws连接");
        WsSessionUtil.removeAndClose(session.getId());
    }
}

2.4 WebSocket配置类

将上一步创建的handler注册到WebSocket服务，并且自定义过滤器，实现beforeHandshake方法，对参数进行校验

import com.paratera.console.biz.handler.WarningWsHandler;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.server.HandshakeInterceptor;

import java.util.Map;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Autowired
    private WarningWsHandler warningWsHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                .addHandler(warningWsHandler, "warningWs")
                //允许跨域，可以具体设置客户端IP
                .setAllowedOrigins("*")
                .addInterceptors(new MyHandshakeInterceptor());
    }

    private static class MyHandshakeInterceptor implements HandshakeInterceptor {
        @Override
        public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                       WebSocketHandler wsHandler, Map<String, Object> attributes)
                throws Exception {

            String collector = ((ServletServerHttpRequest) request).getServletRequest().getParameter("userid");
            if (StringUtils.isEmpty(collector)) {
                return false;
            } else {
                attributes.put("userid", collector);
                return true;
            }
        }
        @Override
        public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Exception exception) {
        }
    }
}

2.5 消息发送

/**
 * 广播消息
 * @param text   要推送的消息
 * @param userId 推送的用户
 * @return
 * @throws IOException
 */
private void broadcastMsg(String text, String userId) {
    for (WebSocketSession session : WsSessionUtil.get().values()) {
        //如果是对应推送的用户，则推送
        if (session.getAttributes().containsValue(userId)) {
            try {
                session.sendMessage(new TextMessage(text));
            } catch (IOException e) {
                log.error("发送报警信息到客户端异常：" + e.getMessage());
            }
        }
    }
}

3、客户端

通过浏览器内置WebSocket插件实现消息的订阅

DOCTYPE html>
<html lang="en">
<head>
    <title>Java WebSocket实现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:20030/warningWs?userid=SELF-cDk-xh6Ph34nLqxY_OvhxWTEYbueEeji-_jW5i1ecEc");
    }
    else {
        alert('当前浏览器 Not support websocket')
    }

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

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

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

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

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

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

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

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

4、测试