【经验分享】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;// 构造器注入MyServicepublic MyWebSocketHandler(MyService myService) {this.myService = myService;}// 处理消息@Overridepublic 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();}}// 在会话建立时设置会话信息@Overridepublic void afterConnectionEstablished(WebSocketSession session) throws Exception {// 设置对话唯一ID和用户信息String conversationId = UUID.randomUUID().toString(); // 生成唯一的对话IDString 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);}// 在会话关闭时清理资源@Overridepublic 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;// 注入自定义的 WebSocketHandlerpublic WebSocketConfig(MyWebSocketHandler myWebSocketHandler) {this.myWebSocketHandler = myWebSocketHandler;}@Overridepublic 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<>(); // 线程安全的集合，保存所有会话@Overridepublic void afterConnectionEstablished(WebSocketSession session) throws Exception {// 将新的会话添加到集合中sessions.add(session);String userId = (String) session.getAttributes().get("userId");sendMessageToClient(session, "Hello, " + userId + "! You are connected.");}@Overridepublic void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {// 处理接收到的消息后，可以向所有客户端广播broadcastMessage("Broadcast message: " + message.getPayload());}@Overridepublic 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<>();