Springboot使用WebSocket发送消息

1. 创建springboot项目，引入spring-boot-starter-websocket依赖

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

完整项目依赖

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>
    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>2.5.9</version>
        <relativePath/> <!-- lookup parent from repository -->
    </parent>
    <groupId>com.mutest</groupId>
    <artifactId>websocket-demo</artifactId>
    <version>0.0.1-SNAPSHOT</version>
    <name>websocket-demo</name>
    <description>websocket demo</description>
    <properties>
        <java.version>8</java.version>
    </properties>
    <dependencies>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter</artifactId>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-test</artifactId>
            <scope>test</scope>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>
        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
        </dependency>
    </dependencies>

    <build>
        <plugins>
            <plugin>
                <groupId>org.springframework.boot</groupId>
                <artifactId>spring-boot-maven-plugin</artifactId>
            </plugin>
        </plugins>
    </build>

</project>

2. 对websocket进行配置config
registerWebSocketHandlers()方法参数是WebSocketHandlerRegistry 。调用WebSocketHandlerRegistry 的addHandler方法传递处理器和路径。处理器参数是上一步创建的TextWebSocketHandler 的子类。路径是客户端调用时跟在端口后的路径。再调用WebSocketHandlerRegistry 的setAllowedOrigins方法传递星号，允许跨域访问。

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    private MyWsHandler myWsHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry
                .addHandler(myWsHandler, "myWs")
                //允许跨域
                .setAllowedOrigins("*");
    }
}

WebSocketHandlerRegistry
addHandler()方法第一个传输传递处理器，第一个参数传递路径
setAllowedOrigins()方法设置允许源，传递星号，否则跨域。

3.定义处理器MyWsHandler

@Component
@Slf4j
public class MyWsHandler extends AbstractWebSocketHandler {

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

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 接收客户端/浏览器端的消息 message.getPayload()消息体
        String payload = message.getPayload();
        log.info("server 接收到客户端/浏览器端的消息 " + payload);
        log.info("发送文本消息");
        //服务端准备给客户端/浏览器端发送消息
        session.sendMessage(new TextMessage("server 发送给的消息 " + payload + "，发送时间:" + LocalDateTime.now().toString()));
    }

    @Override
    protected void handleBinaryMessage(WebSocketSession session, BinaryMessage message) throws Exception {
        log.info("发送二进制消息");
    }

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

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

4. 单播/多播发送消息
在MyWsHandler中，handleTextMessage方法使用的是单播发送消息，如果要给客户端/浏览器端发送广播消息该如何发送呢，只需遍历所有客户端/浏览器端的Session即可，通过Session广播消息。

@Service
@Slf4j
public class WsService {
    /**
     * 发送消息
     *
     * @param session 连接信息
     * @param text    消息
     */
    public void sendMsg(WebSocketSession session, String text) throws IOException {
        session.sendMessage(new TextMessage(text));
    }

    /**
     * 广播消息
     *
     * @param text 消息
     */
    public void broadcastMsg(String text) throws IOException {
        for (WebSocketSession session : WsSessionManager.SESSION_POOL.values()) {
            session.sendMessage(new TextMessage(text));
        }
    }
}

5. 设置定时任务发送广播消息

@Slf4j
@Component
public class MessageJob {
    @Autowired
    WsService wsService;

    /**
     * 每5s发送
     */
    @Scheduled(cron = "0/5 * * * * *")
    public void run() {
        try {
            log.info("推送消息===>" + LocalDateTime.now().toString());
            wsService.broadcastMsg("自动生成消息 " + LocalDateTime.now().toString());
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}

6. 管理客户端/浏览器端的Session工具类

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

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

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

    /**
     * 删除并同步关闭连接
     *
     * @param key key
     */
    public static void removeAndClose(String key) {
        WebSocketSession session = remove(key);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                // todo: 关闭出现异常处理
                e.printStackTrace();
            }
        }
    }

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