Spring Boot+React环境下Stomp协议多订阅会话管理问题咨询

我正在Spring Boot、React环境下开发WebSocket，采用Stomp协议，原本使用ConcurrentHashMap存储会话。这么做是因为当订阅点没有订阅会话时需停止数据传输，存在一个及以上会话时则需定期推送数据。

现有代码如下：

public final static Map<String, String> subscribeSessions = new ConcurrentHashMap<>();

@Override
public void postSend(Message<?> message, MessageChannel channel, boolean sent) {

    StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
    String sessionId = accessor.getSessionId();

    switch (accessor.getCommand()) {
        case CONNECT:
            break;
        case DISCONNECT:
            break;
        case SUBSCRIBE:
            if(!accessor.getDestination().contains("/user/queue")) {
                subscribeSessions.put(accessor.getSessionId(), 
                                      accessor.getDestination());
            }
            break;
        case UNSUBSCRIBE:
            subscribeSessions.remove(accessor.getSessionId());
        default:
            break;
    }
}

我希望支持单个会话订阅多个订阅点（如会话A同时订阅/Topic/ex和/topic/ex/2），但由于HashMap不允许重复键，当前实现无法满足需求。请问是否有解决方法？或者是否有办法统计各订阅点的订阅人数？

1. 调整存储结构支持单会话多订阅 #

将原有的Map<String, String>改为Map<String, Set<String>>，键为会话ID，值为该会话订阅的所有目标地址集合，同时使用线程安全的集合实现保证并发场景下的正确性：

// 替换原有存储结构
public final static Map<String, Set<String>> subscribeSessions = new ConcurrentHashMap<>();

@Override
public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
    StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
    String sessionId = accessor.getSessionId();
    String destination = accessor.getDestination();

    switch (accessor.getCommand()) {
        case CONNECT:
            break;
        case DISCONNECT:
            // 会话断开时移除该会话的所有订阅记录
            subscribeSessions.remove(sessionId);
            break;
        case SUBSCRIBE:
            if (destination != null && !destination.contains("/user/queue")) {
                // 不存在会话对应的集合则创建，再添加订阅点
                subscribeSessions.computeIfAbsent(sessionId, k -> new CopyOnWriteArraySet<>())
                                  .add(destination);
            }
            break;
        case UNSUBSCRIBE:
            if (destination != null) {
                Set<String> userDestinations = subscribeSessions.get(sessionId);
                if (userDestinations != null) {
                    userDestinations.remove(destination);
                    // 集合为空时移除会话键，避免无效内存占用
                    if (userDestinations.isEmpty()) {
                        subscribeSessions.remove(sessionId);
                    }
                }
            }
            break;
        default:
            break;
    }
}

2. 统计各订阅点的订阅人数 #

可以通过维护反向映射关系实现，提供两种可选方式：

方式一：用计数Map直接统计人数 #

// 新增订阅点-订阅人数映射
public final static Map<String, Integer> destinationSubscriberCount = new ConcurrentHashMap<>();

@Override
public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
    StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
    String sessionId = accessor.getSessionId();
    String destination = accessor.getDestination();

    switch (accessor.getCommand()) {
        case CONNECT:
            break;
        case DISCONNECT:
            Set<String> userDests = subscribeSessions.remove(sessionId);
            if (userDests != null) {
                // 遍历会话所有订阅点，逐个减少计数
                for (String dest : userDests) {
                    destinationSubscriberCount.computeIfPresent(dest, (k, v) -> v > 1 ? v - 1 : null);
                }
            }
            break;
        case SUBSCRIBE:
            if (destination != null && !destination.contains("/user/queue")) {
                subscribeSessions.computeIfAbsent(sessionId, k -> new CopyOnWriteArraySet<>())
                                  .add(destination);
                // 订阅点计数+1，不存在则初始化为1
                destinationSubscriberCount.merge(destination, 1, Integer::sum);
            }
            break;
        case UNSUBSCRIBE:
            if (destination != null) {
                Set<String> userDestinations = subscribeSessions.get(sessionId);
                if (userDestinations != null && userDestinations.remove(destination)) {
                    // 仅成功移除订阅点时减少计数
                    destinationSubscriberCount.computeIfPresent(destination, (k, v) -> v > 1 ? v - 1 : null);
                    if (userDestinations.isEmpty()) {
                        subscribeSessions.remove(sessionId);
                    }
                }
            }
            break;
        default:
            break;
    }
}

方式二：用会话集合统计（可查看具体订阅会话） #

这种方式能直观看到哪些会话订阅了指定点，通过集合大小获取人数：

// 新增订阅点-会话ID集合映射
public final static Map<String, Set<String>> destinationSessions = new ConcurrentHashMap<>();

@Override
public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
    StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);
    String sessionId = accessor.getSessionId();
    String destination = accessor.getDestination();

    switch (accessor.getCommand()) {
        case CONNECT:
            break;
        case DISCONNECT:
            Set<String> userDests = subscribeSessions.remove(sessionId);
            if (userDests != null) {
                for (String dest : userDests) {
                    Set<String> destSessionSet = destinationSessions.get(dest);
                    if (destSessionSet != null) {
                        destSessionSet.remove(sessionId);
                        if (destSessionSet.isEmpty()) {
                            destinationSessions.remove(dest);
                        }
                    }
                }
            }
            break;
        case SUBSCRIBE:
            if (destination != null && !destination.contains("/user/queue")) {
                // 更新会话-订阅点映射
                subscribeSessions.computeIfAbsent(sessionId, k -> new CopyOnWriteArraySet<>())
                                  .add(destination);
                // 更新订阅点-会话映射
                destinationSessions.computeIfAbsent(destination, k -> new CopyOnWriteArraySet<>())
                                  .add(sessionId);
            }
            break;
        case UNSUBSCRIBE:
            if (destination != null) {
                Set<String> userDestinations = subscribeSessions.get(sessionId);
                if (userDestinations != null && userDestinations.remove(destination)) {
                    Set<String> destSessionSet = destinationSessions.get(destination);
                    if (destSessionSet != null) {
                        destSessionSet.remove(sessionId);
                        if (destSessionSet.isEmpty()) {
                            destinationSessions.remove(destination);
                        }
                    }
                    if (userDestinations.isEmpty()) {
                        subscribeSessions.remove(sessionId);
                    }
                }
            }
            break;
        default:
            break;
    }
}

内容的提问来源于stack exchange，提问作者김진영