SpringBoot集成WebSocket支持流式消息推送

背景

在日常工作中，经常遇到需要前后端消息交互的场景，例如

• 客服软电话，需要通知前端弹出工作台。
• 业务异常报警，提醒用户及时关注。
• 用户账户余额不足，提醒其及时充值。
• 持续向用户推送应用日志等等。

---

以上场景，自然而然就需要webSocket支持。 而且，需要支持单条消息推送，也要支撑流式消息推送。废话不多说，开搞！

一、后端使用SpringBoot集成Websocket。

1.引入依赖

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

2.创建WebsocketServerEndpoint。

说明：由于spring默认单例，这里需要用集合记录Endpoint。

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.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @Author Doit
 * @Date 2022/8/21 10:25
 * @Desc websocket server
 * @Version 1.0
 * @Slogan Just do it.
 */

@Slf4j
@Component
@ServerEndpoint("/doit/websocket/{target}") //创建ws的请求路径。
public class WebsocketServerEndpoint {
    private Session session;
    private String target;
    //支持持续流推送
    private InputStream inputStream;
    private final static CopyOnWriteArraySet<WebsocketServerEndpoint> websockets = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session , @PathParam("target") String target){
        this.session = session;
        this.target = target;
        websockets.add(this);
        log.info("websocket connect server success , target is {},total is {}",target,websockets.size());
    }

    @OnMessage
    public void onMessage(String message) {
        log.info("message is {}",message);
    }

    @OnClose
    public void onClose(){
        log.info("connection has been closed ,target is {},total is {}" ,this.target, websockets.size());
        this.destroy();
    }

    @OnError
    public void onError(Throwable throwable){
        this.destroy();
        log.info("websocket connect error , target is {} ,total is {}, error is {}",this.target ,websockets.size(),throwable.getMessage());
    }

    /**
     * 根据目标身份推送消息
     * @param target
     * @param message
     * @throws IOException
     */
    public void sendMessageOnce(String target, String message) throws IOException {
        this.sendMessage(target,message,false,null);
    }

    /**
     * stream 同步日志输出，通过websocket推送至前台。
     * @param target
     * @param is
     * @throws IOException
     */
    public void sendMessageSync(String target, InputStream is) throws IOException {
        this.sendMessage(target,null,true , is);
    }

    /**
     * Send message.
     * @param target 通过target获取{@link WebsocketServerEndpoint}.
     * @param message message
     * @param continuous 是否通过inputStream持续推送消息。
     * @param is 输入流
     * @throws IOException
     */
    private void sendMessage(String target , String message ,Boolean continuous , InputStream is) throws IOException {
        WebsocketServerEndpoint websocket = getWebsocket(target);
        if(Objects.isNull(websocket)){
            throw new RuntimeException("The websocket does not exists or has been closed.");
        }
        if(continuous){
            if(Objects.isNull(is)){
                throw new RuntimeException("InputStream can not be null when continuous is true.");
            }else{
                websocket.inputStream = is;
                CompletableFuture.runAsync(websocket::sendMessageWithInputSteam);
            }
        }else{
            websocket.session.getBasicRemote().sendText(message);
        }
    }

    /**
     * 通过inputStream 持续推送消息。
     * 支持文件、消息、日志等。
     */
    private void sendMessageWithInputSteam(){
        String message;
        BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(this.inputStream));
        try {
            while ((message = bufferedReader.readLine()) !=null){
                if(websockets.contains(this)){
                    this.session.getBasicRemote().sendText(message);
                }
            }
        }catch(IOException e){
            log.warn("SendMessage failed {}",e.getMessage());
        }finally {
            this.closeInputStream();
        }
    }

    /**
     * 根据目标获取对应的{@link WebsocketServerEndpoint}。
     * @param target 约定标的
     * @return WebsocketServerEndpoint
     */
    private WebsocketServerEndpoint getWebsocket(String target){
        WebsocketServerEndpoint websocket = null;
        for (WebsocketServerEndpoint ws : websockets) {
            if (target.equals(ws.target)) {
                websocket = ws;
            }
        }
        return websocket;
    }

    /**
     * close inputStream.
     * @Author Doit
     * @Date 20221/08/23 15:30:00
     */
    private void closeInputStream(){
        if(Objects.nonNull(inputStream)){
            try {
                inputStream.close();
            } catch (Exception e) {
                log.warn("websocket close failed {}",e.getMessage());
            }
        }
    }

    /**
     * destroy {@link WebsocketServerEndpoint}
     * @Author Doit
     * @Date 20221/08/23 15:30:00
     */
    private void destroy(){
        websockets.remove(this);
        this.closeInputStream();
    }
}

3.创建Configuration，把WebsocketServerEndpoint交给SpringBoot作支持。

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

/**
 * @Author Doit
 * @Date 2022/8/21 10:25
 * @Desc websocket configuration
 * @Version 1.0
 * @Slogan Just do it.
 */
@Configuration
public class WebsocketConfig {
    @Bean
    public ServerEndpointExporter serverEndpointExporter(){
        return new ServerEndpointExporter();
    }
}

此时，后段内容已完成。

4.前端采用SocketJS，帮你解决浏览器兼容问题，推荐使用。

<script src="./socket.js"></script>
<script>
    //@ServerEndpoint 里的地址。 ws用于http，wss用于https，后者更安全。
    var ws = new WebSocket("ws//127.0.0.1:10881/doit/websocket/{target}");
    //建立连接
    ws.onopen = function(msg) {
        console.log("Connection open ...");
    };
    //发送消息
    ws.onmessage = function(msg) {
        console.log("Received msg: " + msg.data);
    };
    //关闭连接
    ws.onclose = function(msg) {
        console.log("Connection closed.");
        ws.close();
    };
    ws.onerror = function(msg){
        console.log("Connection error.");
        ws.close();
    };
</script>

5.开发完成，可以快乐的单机使用了。

1. 首先通过target约定身份建立ws，一般采用请求方IP组合用户ID的方式。
2. 通过接口推送消息，支持sendOnce 和 sendWithInputStream，分别用来支持单次推送和流式推送（如日志等）。

后记

以上内容仅供单机玩玩，生产使用需要集群方案。后续有进阶方案，欢迎关注。