解决WebSocket程序中无法注入Spring管理的Bean

1.原因

起初我想在关闭连接的时候将内存中的数据写入到数据库中，发现无法注入dao

理由是：
WebSocket API 是独立于任何特定框架的标准，它的生命周期和管理是由 WebSocket 容器负责的，而不是由 Spring 容器负责。因此，WebSocket 容器不会识别或处理 Spring 的依赖注入注解。

 

在 @ServerEndpoint 注解中，@OnClose 方法是由 WebSocket 容器调用的，而不是由 Spring 容器调用。WebSocket 容器是独立于 Spring 容器的，它不了解或处理 Spring 的依赖注入机制。当 WebSocket 连接关闭时，WebSocket 容器会调用 @OnClose 方法，但此时 Spring 容器并不参与调用过程。因此，在 @OnClose 方法中无法直接使用 Spring 的依赖注入机制来注入普通的 Bean。

2.背景

@Configuration
public class WebSocketConfig implements WebSocketConfigurer {
    // 主要是这一步，注册bean，扫描并注册这些 WebSocket 端点，
    // 使它们可以被正确地管理和调用
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
    }
}

这里一般只存在三个方法

@onOpen
@onMessage

@onClose

分别在建立连接、接收消息、断开连接调用

@Slf4j
@Component
@ServerEndpoint("/websocket/{name}")
public class WebSocketEndpoint{

    private String name;
    private Session session;


    private static MessageMapper messageMapper;

    @Autowired
    public void setMessageMapper(MessageMapper messageMapper) {
        this.messageMapper = messageMapper;
    }
    // 单线程的线程池
    private ExecutorService executorService = Executors.newSingleThreadExecutor();

    // 一个线程安全的集合
    private static ConcurrentHashMap<String, Session> sessionPools = WebSocketServer.sessionPools;

    /**
     * 处理客户端连接的逻辑
     *
     * @param session
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("name") String name) {
        this.session = session;
        this.name = name;

        sessionPools.put(name, session);
        log.info("{}【websocket消息】{}建立新的连接, 在线人数为:{}",
                LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")),
                name, sessionPools.size());
    }

    /**
     * 处理收到的消息的逻辑
     *
     * @param msg
     */
    @OnMessage
    public void OnMessage(String msg) {
        Message data = JSON.parseObject(msg, Message.class);
        data.setSendTime(LocalDateTime.now());

        // 如果传输过来的是文本信息
        if (data.getDataType().equals(MsgType.TEXT)) {
            if (data.getReceiveUser().equals("all"))
                p2All(data.getData(), data.getFromUser());
            else
                p2Peer(data.getData(), data.getFromUser(), data.getReceiveUser());
        }

        // 缓存
        RedisUtil.set(RedisUtil.generateKey(), data);
        log.info("{}【websocket消息】收到客户端发来的消息:{},是{}向{}发送消息",
                LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")),
                data.getData(), data.getFromUser(), data.getReceiveUser());
    }

    /**
     * 处理客户端关闭连接的逻辑
     * 从内存中写入数据库
     */
    @OnClose
    public void OnClose() {
        sessionPools.remove(this.name);
        log.info("{}【websocket消息】{}连接断开, 总数:{}",
                LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")),
                this.name, sessionPools.size());

        log.info("{} 清除缓存中...", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));
        // 清除缓存
        List<Message> list = RedisUtil.getAllaRemove(Message.class);
        log.info("{} 写入数据库中...", LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")));// 异步写入数据库

        writeInputDB(list, Message.class);

    }

    protected <T> void writeInputDB(List<T> list, Class<T> target) {
        list.forEach(msg -> {
            Message data = (Message) target.cast(msg);
            int insert = messageMapper.insert(data);
            System.out.println(insert);
        });
    }

    /**
     * @param msg      消息
     * @param fromName 发送者
     * @param toName   接收者
     */
    public void p2Peer(String msg, String fromName, String toName) {
        Message data = Message.builder()
                .data(msg)
                .dataType(MsgType.TEXT)
                .fromUser(fromName)
                .receiveUser(toName)
                .sendTime(LocalDateTime.now())
                .build();


        try {
            Session to = sessionPools.get(toName);
            // 如果在线
            if (to != null) {
                log.info("{}【websocket消息】{}向{}发送了一条消息",
                        LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss")),
                        data.getFromUser(), data.getReceiveUser());

                to.getBasicRemote().sendText(JSON.toJSONString(data, true));
            }
            // 如果不在线
            else {
                sessionPools.get(fromName).getBasicRemote().
                        sendText(JSON.toJSONString(Message.builder().data("该用户不在线...").build()));
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void p2All(String msg, String fromName) {
        Message data = Message.builder()
                .data(msg)
                .receiveUser("all")
                .fromUser(fromName)
                .sendTime(LocalDateTime.now())
                .build();

        sessionPools.forEach((name, session) -> {
            if (!name.equals(fromName)) {
                try {
                    session.getBasicRemote().sendText(JSONObject.toJSONString(data));
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        });

    }

}

3.解决方案

1.使用静态方法或单例模式：将需要注入的 Bean 定义为静态的或使用单例模式，在 WebSocket 端点中直接访问该实例。

@Autowired
    public void setMessageMapper(MessageMapper messageMapper) {
        this.messageMapper = messageMapper;
    }

2.直接把websocket容器交给spring管理

在启动类上加上

@Import(WebSocketConfig.class)