Springboot——整合websocket并根据type区别处理

前言

之前写过一篇类似的博客，但之前写的逻辑过于简单，而且所有的websocket处理都在一个处理器中完成。如果需要按照消息类型等做区分操作时，会导致所有的逻辑处理都在一个处理类中，显得过于冗余。

最近一直在想一个问题，采取websocket通信处理时，能否根据某个变量，比如type，区别进入不同的处理器中。

往期博客参考：Springboot——websocket使用

架构思想

目前考虑的是用一个公共的方法接收所有的ws请求，根据传递参数type的不同，进行分发路由到不同的处理器上完成。

项目结构

代码实现

依赖引入

使用到的技术点：websocket、aop、thymeleaf 。

<!-- aop -->
<dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-aop</artifactId>
</dependency><dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-thymeleaf</artifactId>
</dependency>
<dependency><groupId>org.springframework.boot</groupId><artifactId>spring-boot-starter-websocket</artifactId>
</dependency><!-- 转换相关 -->
<dependency><groupId>com.fasterxml.jackson.core</groupId><artifactId>jackson-databind</artifactId>
</dependency>

自定义注解

新建自定义注解，标注具体的实现类，并指定唯一的类型type。

package cn.xj.aspect;import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
import java.lang.annotation.RetentionPolicy;
import java.lang.annotation.Target;/*** @author xj*/
@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE) // 仅作用于类上
public @interface WebSocketTypeHandler {/*** 指定类型* @return*/String type() default "";
}

定义具体的处理类

定义 TypeAWebSocketHandler

定义TypeAWebSocketHandler，使用自定义注解标注，并指定类型为typeA。

package cn.xj.wshandler;import cn.xj.aspect.WebSocketTypeHandler;
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.TextWebSocketHandler;@Component
@WebSocketTypeHandler(type= "typeA")
public class TypeAWebSocketHandler extends TextWebSocketHandler {@Overrideprotected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {// 处理Type A类型的消息逻辑System.out.println("处理Type A消息: " + message.getPayload());session.sendMessage(new TextMessage("Type A处理结果: " + message.getPayload()));}@Overridepublic void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {// 连接关闭处理逻辑System.out.println("Type A连接关闭: " + status.getReason());}
}

定义 TypeBWebSocketHandler

与TypeAWebSocketHandler一样，只是输出日志不同。

package cn.xj.wshandler;import cn.xj.aspect.WebSocketTypeHandler;
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.TextWebSocketHandler;@Component
@WebSocketTypeHandler(type = "typeB")
public class TypeBWebSocketHandler extends TextWebSocketHandler {@Overrideprotected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {// 处理Type A类型的消息逻辑System.out.println("处理Type B消息: " + message.getPayload());session.sendMessage(new TextMessage("Type B处理结果: " + message.getPayload()));}@Overridepublic void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {// 连接关闭处理逻辑System.out.println("Type B连接关闭: " + status.getReason());}
}

定义路由处理类

定义路由处理类。需要实现的功能点如下：

• 所有的ws请求都会进入当前route中，并按照指定字段进行请求的分发。
• 在容器启动加载完成后，就能将对应的bean识别和加载到本地缓存中。

package cn.xj.wshandler;import java.io.IOException;
import java.lang.reflect.AnnotatedElement;
import java.util.HashMap;
import java.util.Map;import cn.xj.aspect.WebSocketTypeHandler;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.stereotype.Component;
import org.springframework.util.CollectionUtils;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;@Component
public class WebSocketRouter extends TextWebSocketHandler implements InitializingBean, ApplicationContextAware,DisposableBean {/*** 类型对应实例化对象缓存*/private final Map<String, TextWebSocketHandler> handlerMap = new HashMap<>();/*** 容器*/private ApplicationContext applicationContext;/*** 用于设置 applicationcontext 属性 ApplicationContextAware实现类* @param applicationContext the ApplicationContext object to be used by this object* @throws BeansException*/@Overridepublic void setApplicationContext(ApplicationContext applicationContext) throws BeansException {this.applicationContext = applicationContext;}/*** 应用启动完成  InitializingBean 实现类* @throws Exception*/@Overridepublic void afterPropertiesSet() throws Exception {// 获取所有指定注解标识的beanMap<String, Object> handlerBeans = applicationContext.getBeansWithAnnotation(WebSocketTypeHandler.class);if (!CollectionUtils.isEmpty(handlerBeans)) {for (Object handlerBean : handlerBeans.values()) {AnnotatedElement annotatedElement = handlerBean.getClass();// 获取注解标识的值WebSocketTypeHandler webSocketTypeHandler = annotatedElement.getAnnotation(WebSocketTypeHandler.class);String type = webSocketTypeHandler.type();// 按照类型 存对应的beanhandlerMap.put(type, (TextWebSocketHandler) handlerBean);}}}/*** 应用销毁 DisposableBean 实现类* @throws Exception*/@Overridepublic void destroy() throws Exception {// help GChandlerMap.clear();}private final ObjectMapper objectMapper = new ObjectMapper();/*** ws 消息处理转发* @param session* @param message* @throws Exception*/@Overrideprotected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {try {// 转换Map<String, Object> request = objectMapper.readValue(message.getPayload(), Map.class);String type = (String) request.get("type");TextWebSocketHandler handler = handlerMap.get(type);if (handler == null) {session.sendMessage(new TextMessage("不支持的type: " + type));return;} handler.handleMessage(session,message);} catch (IOException e) {session.sendMessage(new TextMessage("消息解析错误"));}}/*** 断开连接处理* @param session* @param status* @throws Exception*/@Overridepublic void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {for (TextWebSocketHandler textWebSocketHandler : handlerMap.values()) {textWebSocketHandler.afterConnectionClosed(session,status);}}
}

配置类，绑定point

package cn.xj.config;import cn.xj.wshandler.WebSocketRouter;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {@Autowiredprivate WebSocketRouter webSocketRouter;/*** 绑定前端请求地址 /websocket-endpoint* @param registry*/@Overridepublic void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {registry.addHandler(webSocketRouter, "/websocket-endpoint").setAllowedOrigins("*");// 其他绑定//registry.addHandler(webSocketRouter, "/websocket-endpoint2").setAllowedOrigins("*");}
}

制定前端页面

此处使用的是thymeleaf框架，必须引入对应的依赖。其次在src/main/resources 中需要创建一个新的目录 templates 存放前端文件。

<!DOCTYPE html>
<html lang="en"><head><meta charset="UTF - 8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><title>WebSocket Example </title>
</head><body>
<script>const socket = new WebSocket('ws://localhost:8080/websocket-endpoint');socket.onopen = function (event) {console.log('WebSocket连接已建立');};socket.onmessage = function (event) {console.log('收到消息:', event.data);};socket.onclose = function (event) {console.log('WebSocket连接已关闭');};function sendMessage(type, content) {const message = {type: type,content: content};socket.send(JSON.stringify(message));}
</script>
传递来的数据值cid:
<input type="text" th:value="${cid}" id="cid" readonly />
<button onclick="sendMessage('typeA', '这是一条typeA的测试消息')">发送typeA消息</button>
<button onclick="sendMessage('typeB', '这是一条typeB的测试消息')">发送typeB消息</button>
</body></html>

编写测试接口方便跳转进入前端页面

package cn.xj.controller;import org.springframework.stereotype.Controller;
import org.springframework.ui.Model;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;/*** 前端页面进入*/
@Controller
@RequestMapping("/view")
public class ViewController {/*** 测试页面跳转，携带自定义的cid信息* @param cid* @param model* @return*/@GetMapping("/socket/{cid}")public String socket(@PathVariable String cid, Model model){model.addAttribute("cid", cid);return "index";}
}

测试验证

请求前端地址，按f12打开浏览器控制台。

http://localhost:8080/view/socket/1

结语

本代码仅供测试使用，afterConnectionClosed逻辑测试中存在问题，正常使用需要进行调整。