【经验分享】SpringBoot集成WebSocket开发-03 使用WebSocketSession为每个对话存储独立信息

WebSocketSession为每个对话存储独立信息

要在每个WebSocket会话中存储独立的信息，比如用户信息、对话唯一ID等，可以通过以下几种方式来实现：

1. 使用WebSocketSession存储会话级别的属性

WebSocketSession对象提供了一个getAttributes()方法，可以用来为每个WebSocket会话存储自定义信息。我们可以将用户信息、会话ID等存储在这个属性中，确保每个会话都有独立的数据。

2. 管理WebSocket会话

我们可以在WebSocket处理器中跟踪每个WebSocket会话，并在会话关闭时清理数据。

示例实现

下面是一个完整的示例，展示如何为每个WebSocket会话存储独立的信息：

1. 修改 MyWebSocketHandler，使用 WebSocketSession 存储会话信息

import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.TextMessage;

import java.util.Map;
import java.util.UUID;

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    private final MyService myService;

    // 构造器注入MyService
    public MyWebSocketHandler(MyService myService) {
        this.myService = myService;
    }

    // 处理消息
    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {
        // 从session获取对话唯一ID和用户信息
        String conversationId = (String) session.getAttributes().get("conversationId");
        String userName = (String) session.getAttributes().get("userName");

        // 可以将这些信息传给service进行处理
        String response = myService.processMessage(message.getPayload(), conversationId, userName);

        try {
            // 发送处理后的消息
            session.sendMessage(new TextMessage(response));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    // 在会话建立时设置会话信息
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 设置对话唯一ID和用户信息
        String conversationId = UUID.randomUUID().toString(); // 生成唯一的对话ID
        String userName = "Guest-" + UUID.randomUUID().toString(); // 假设用户名是随机生成的，实际中可以通过认证来获取

        // 将对话ID和用户名存入会话属性
        session.getAttributes().put("conversationId", conversationId);
        session.getAttributes().put("userName", userName);

        System.out.println("Connection established for user: " + userName + " with conversation ID: " + conversationId);
    }

    // 在会话关闭时清理资源
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        // 清理资源
        session.getAttributes().clear();
        System.out.println("Connection closed for session: " + session.getId());
    }
}

2. 修改 MyService，接受对话ID和用户名参数

MyService类需要处理来自不同会话的消息，并根据不同的对话ID和用户信息来做出响应。

import org.springframework.stereotype.Service;

@Service
public class MyService {

    public String processMessage(String message, String conversationId, String userName) {
        // 这里根据对话ID和用户名处理消息
        // 你可以通过conversationId来查找历史对话记录，或者根据用户名做不同的响应
        return "用户: " + userName + " 在对话: " + conversationId + " 中发送的消息: " + message;
    }
}

3. WebSocket 配置类

WebSocketConfig类不需要做修改，只需要保持WebSocket的基本配置即可。

import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.WebSocketHandler;
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 {

    private final MyWebSocketHandler myWebSocketHandler;

    // 注入自定义的 WebSocketHandler
    public WebSocketConfig(MyWebSocketHandler myWebSocketHandler) {
        this.myWebSocketHandler = myWebSocketHandler;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(myWebSocketHandler, "/ws").setAllowedOrigins("*");
    }
}

4. 启动类

与之前一样，我们创建一个Spring Boot启动类来启动应用。

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

@SpringBootApplication
public class WebSocketApplication {
    public static void main(String[] args) {
        SpringApplication.run(WebSocketApplication.class, args);
    }
}

5. 说明

• 在 MyWebSocketHandler 中，使用了 WebSocketSession 对象的 getAttributes() 来存储会话级别的属性，如对话ID和用户名等。这些信息在WebSocket会话建立时被初始化，并且在后续的消息处理中被引用。
• 使用 UUID.randomUUID().toString() 来生成一个唯一的对话ID，每个连接会有一个独立的对话ID。实际应用中，你可以通过身份认证来为用户生成一个真实的用户名。
• 每个WebSocket连接都是独立的，WebSocketSession 会自动维护各自独立的状态，所以每个连接可以保存自己的对话ID和用户信息。
• 在会话关闭时，通过 afterConnectionClosed() 方法清理会话信息，确保没有残留的数据。

6. 进一步扩展

• 认证: 如果你需要根据用户身份来设置用户信息，可以在 afterConnectionEstablished() 方法中进行认证（比如使用JWT令牌来识别用户）。
• 历史消息: 你可以在 MyService 中扩展，利用 conversationId 来查找并返回该对话的历史消息。
• 广播消息: 如果你需要广播消息给所有用户，可以使用一个全局的WebSocket会话管理器来存储每个会话，并根据需求发送消息。

这样，每个WebSocket会话就能够独立地存储和访问用户信息、对话ID等数据，同时能够根据不同的会话做出不同的响应。

客户端主动发送消息

向多个客户端广播消息

如果你希望向所有连接的客户端广播消息，而不仅仅是某个单独的客户端，你需要维护一个所有WebSocket会话的列表。这通常可以通过一个简单的集合来实现，比如使用Set<WebSocketSession>来存储所有活跃的WebSocket会话。

示例：广播消息给所有连接的客户端

import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.stereotype.Component;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;

@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    private final Set<WebSocketSession> sessions = new CopyOnWriteArraySet<>(); // 线程安全的集合，保存所有会话

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        // 将新的会话添加到集合中
        sessions.add(session);
        String userId = (String) session.getAttributes().get("userId");
        sendMessageToClient(session, "Hello, " + userId + "! You are connected.");
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        // 处理接收到的消息后，可以向所有客户端广播
        broadcastMessage("Broadcast message: " + message.getPayload());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus status) throws Exception {
        // 移除关闭的会话
        sessions.remove(session);
    }

    // 向所有连接的客户端广播消息
    private void broadcastMessage(String message) throws Exception {
        TextMessage textMessage = new TextMessage(message); // 创建TextMessage对象
        for (WebSocketSession session : sessions) {
            session.sendMessage(textMessage); // 向所有客户端发送消息
        }
    }

    // 向单个客户端发送消息
    private void sendMessageToClient(WebSocketSession session, String message) throws Exception {
        TextMessage textMessage = new TextMessage(message);
        session.sendMessage(textMessage); // 发送消息
    }
}

关键点：

• sessions：这个集合存储了所有活跃的WebSocket会话，每个客户端的连接都会被保存到这个集合中。

• broadcastMessage：这个方法遍历所有活跃的会话，并向每个会话发送消息，实现了广播消息功能。

• afterConnectionEstablished：在每个客户端连接时，将其WebSocket会话对象添加到sessions集合中。

• afterConnectionClosed：在WebSocket连接关闭时，移除对应的会话对象，确保集合中只保留当前活跃的会话。

• 可以进一步优化为线程安全的map进行存储session信息例如

     /**
       * 使用ConcurrentHashMap来存储WebSocket会话
       */
      private final ConcurrentHashMap<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();