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 {

    @Override
    protected 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()));
    }

    @Override
    public 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 {

    @Override
    protected 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()));
    }

    @Override
    public 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
     */
    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }


    /**
     * 应用启动完成  InitializingBean 实现类
     * @throws Exception
     */
    @Override
    public void afterPropertiesSet() throws Exception {
        // 获取所有指定注解标识的bean
        Map<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();
                // 按照类型 存对应的bean
                handlerMap.put(type, (TextWebSocketHandler) handlerBean);
            }
        }
    }

    /**
     * 应用销毁 DisposableBean 实现类
     * @throws Exception
     */
    @Override
    public void destroy() throws Exception {
        // help GC
        handlerMap.clear();
    }

    private final ObjectMapper objectMapper = new ObjectMapper();

    /**
     * ws 消息处理转发
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    protected 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
     */
    @Override
    public 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 {

    @Autowired
    private WebSocketRouter webSocketRouter;

    /**
     * 绑定前端请求地址 /websocket-endpoint
     * @param registry
     */
    @Override
    public 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逻辑测试中存在问题，正常使用需要进行调整。