服务器消息推送-WebSocket

在做 Web 系统时，经常会遇到一些需要服务端主动向客户端发送通知的场景，一般我们可以使用 ajax 轮询来查询服务端信息，但是这样会对服务器造成较大的压力，对于以上问题，我们可以使用 WebSocket 或者 SSE（Server Sent Event）来解决

WebSocket：基于 TCP 的一种网络协议，它实现了浏览器与服务器全双工通信，而且数据格式比较轻量，性能开销小，不用频繁创建及销毁 TCP 请求

1. 依赖

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

2. Config

@Configuration
public class WebSocketConfig {

    @Bean
    public ServerEndpointExporter getServerEndpointExporter() {
        return new ServerEndpointExporter();
    }
}

3. Server

@Slf4j
@ServerEndpoint(value = "/websocket/{userId}")
@Component
public class WebSocketServer {

    /**
     * 建立连接
     *
     * @param userId
     * @param session
     */
    @OnOpen
    public void onOpen(@PathParam("userId") Integer userId, Session session) {
        WebSocketUtil.userCache.put(userId, session);
        WebSocketUtil.sendMsg(userId, "连接成功");
        log.info("有新连接加入，当前连接数为：" + WebSocketUtil.userCache.size());
    }

    /**
     * 关闭连接
     *
     * @param userId
     * @param session
     */
    @OnClose
    public void onClose(@PathParam("userId") Integer userId, Session session) {
        try {
            session.close();
        } catch (IOException e) {
            log.error("关闭连接异常", e);
        }
        WebSocketUtil.sendMsg(userId, "关闭连接");
        WebSocketUtil.userCache.remove(userId);
        log.info("有一个连接关闭，当前连接数为：" + WebSocketUtil.userCache.size());
    }

    /**
     * 客户端发送消息
     *
     * @param userId
     * @param msg
     */
    @OnMessage
    public void onMessage(@PathParam("userId") Integer userId, String msg) {
        log.info("收到来自用户：[{}]的消息：{}", userId, msg);
    }

    /**
     * 连接异常
     *
     * @param session
     * @param error
     */
    @OnError
    public void onError(Session session, Throwable error) {
        try {
            session.close();
        } catch (IOException ignored) {
        }
        log.error("连接异常：{}", error.getMessage());
    }
}

4. Util

@Slf4j
public class WebSocketUtil {

    public static final Map<Integer, Session> userCache = new ConcurrentHashMap<>();

    public static void sendAllMsg(String msg) {
        userCache.forEach((k, v) -> sendMsg(k, msg));
    }

    public static void sendMsg(Integer userId, String msg) {
        log.info("发送消息：{}，到用户：[{}]", msg, userId);
        try {
            userCache.get(userId).getBasicRemote().sendText(msg);
        } catch (IOException e) {
            log.error("发送消息失败：", e);
        }
    }
}

5. Controller

@RestController
public class WebSocketController {

    @GetMapping("/send/{userId}")
    public void send(@PathVariable Integer userId, @RequestParam String msg) {
        WebSocketUtil.sendMsg(userId, msg);
    }

    @GetMapping("/send")
    public void send(@RequestParam String msg) {
        WebSocketUtil.sendAllMsg(msg);
    }

    @GetMapping("/onlineUser")
    public String onlineUser() {
        return WebSocketUtil.userCache.keySet().toString();
    }
}

6. html

<!DOCTYPE html>
<html>
<head>
    <meta charset="utf-8">
    <title>websocket测试</title>
    <script type="text/javascript">
      let ws;

      function init() {
            ws = new WebSocket("ws://localhost:8080/websocket/1");

            ws.onopen = function () {
                printMsg("onOpen");
            };
            ws.onmessage = function (e) {
                printMsg("onMessage: " + e.data);
            };
            ws.onclose = function () {
                printMsg("onClose");
            };
            ws.onerror = function () {
                printMsg("onError");
            };
        }

        function onSubmit() {
          const input = document.getElementById("input");
          ws.send(input.value);
            printMsg("sendMessage: " + input.value);
            input.value = "";
            input.focus();
        }

        function onCloseClick() {
            ws.close();
        }

        function printMsg(str) {
          const log = document.getElementById("log");
          log.innerHTML = str + "<br>" + log.innerHTML;
        }
    </script>
</head>

<body onload="init();">
<form onsubmit="onSubmit(); return false;">
    <input type="text" id="input">
    <input type="submit" value="send">
    <button onclick="onCloseClick(); return false;">close</button>
</form>
<div id="log"></div>
</body>
</html>