如何通过Spring Boot Service向指定WebSocket连接用户发送消息

核心思路 #

要实现从@Service类向指定用户发送WebSocket消息，关键是维护WebSocket会话与用户标识的映射关系，通过映射找到目标会话后即可发送消息。

步骤1：创建线程安全的会话管理器 #

实现一个组件存储用户ID/连接ID与WebSocketSession的对应关系，同时处理会话关闭后的清理：

@Component
public class WebSocketSessionManager {
    // 用ConcurrentHashMap保证多线程环境下的线程安全
    private final Map<String, WebSocketSession> sessionMap = new ConcurrentHashMap<>();

    // 注册会话：支持按用户ID或连接ID存储
    public void registerSession(String identifier, WebSocketSession session) {
        sessionMap.put(identifier, session);
        // 会话关闭时自动移除映射，避免内存泄漏
        session.onClose().doFinally(signalType -> sessionMap.remove(identifier)).subscribe();
    }

    // 根据标识查找会话
    public Optional<WebSocketSession> getSession(String identifier) {
        return Optional.ofNullable(sessionMap.get(identifier));
    }
}

步骤2：修改WebSocketHandler注册会话 #

在MyHandler中注入会话管理器，客户端建立连接时，将会话与用户标识绑定（可从握手请求中获取用户ID或直接用会话自身的ID）：

public class MyHandler implements WebSocketHandler {
    private final WebSocketSessionManager sessionManager;

    // 构造注入会话管理器
    public MyHandler(WebSocketSessionManager sessionManager) {
        this.sessionManager = sessionManager;
    }

    @Override
    public Mono<Void> handle(WebSocketSession webSocketSession) {
        // 示例：从握手URL的请求参数中获取用户ID，也可通过Header/认证信息获取
        String userId = webSocketSession.getHandshakeInfo().getUri().getQuery().split("userId=")[1];
        // 获取WebSocket会话自身的唯一ID
        String connectionId = webSocketSession.getId();

        // 同时注册用户ID和连接ID的映射，方便后续灵活查找
        sessionManager.registerSession(userId, webSocketSession);
        sessionManager.registerSession(connectionId, webSocketSession);

        // 保持连接并处理客户端消息（若不需要处理可省略doOnNext）
        return webSocketSession.receive()
                .doOnNext(message -> {
                    // 可选逻辑：处理客户端发送的消息
                })
                .then();
    }
}

步骤3：更新WebSocketConfig注入依赖 #

将MyHandler改为Bean注入，确保会话管理器能被正确传入：

@Configuration
public class WebSocketConfig {

    @Bean
    public MyHandler myHandler(WebSocketSessionManager sessionManager) {
        return new MyHandler(sessionManager);
    }

    @Bean
    public HandlerMapping handlerMapping(MyHandler myHandler) {
        Map<String, WebSocketHandler> map = new HashMap<>();
        map.put("/data", myHandler);
        SimpleUrlHandlerMapping mapping = new SimpleUrlHandlerMapping();
        mapping.setUrlMap(map);
        mapping.setOrder(Ordered.HIGHEST_PRECEDENCE);
        return mapping;
    }
}

步骤4：在Service中实现定向消息发送 #

注入会话管理器，根据用户ID或连接ID找到目标会话并发送消息：

@Service
public class MyService {
    private final WebSocketSessionManager sessionManager;

    // 构造注入会话管理器
    public MyService(WebSocketSessionManager sessionManager) {
        this.sessionManager = sessionManager;
    }

    public void sendToUser(String userId, String connectionId, String message) {
        // 优先按用户ID查找会话，找不到则尝试用连接ID
        Optional<WebSocketSession> targetSession = sessionManager.getSession(userId);
        if (targetSession.isEmpty()) {
            targetSession = sessionManager.getSession(connectionId);
        }

        targetSession.ifPresent(session -> {
            // 创建文本消息并发送，subscribe触发响应式操作执行
            WebSocketMessage webSocketMessage = session.textMessage(message);
            session.send(Mono.just(webSocketMessage)).subscribe();
        });
    }
}

关键注意事项 #

• 线程安全：必须使用ConcurrentHashMap存储会话，避免并发场景下的映射异常。
• 会话清理：务必在会话关闭时移除映射，防止无效会话占用内存。
• 响应式触发：WebFlux的Mono是懒加载的，必须调用subscribe()才会实际发送消息。
• 用户标识安全：生产环境优先从认证信息（如webSocketSession.getHandshakeInfo().getPrincipal()）获取用户ID，避免用请求参数传递敏感信息。

内容的提问来源于stack exchange，提问作者Vito Lipari