如何在Spring GraphQL中编程式关闭所有WebSocket会话？

由于Spring GraphQL的GraphQlWebSocketHandler内部维护的sessionInfoMap是私有不对外暴露的，我们可以通过反射访问内部会话或自定义拦截器跟踪会话两种方式实现需求，以下是具体实现：

方法一：反射访问内部会话并关闭 #

该方式直接通过反射获取框架私有维护的会话信息，适合快速实现，但需注意版本兼容性（框架内部结构变更可能导致失效）。

实现代码 #

1. 定义会话管理组件：

import org.springframework.graphql.server.webmvc.GraphQlWebSocketHandler;
import org.springframework.stereotype.Component;
import java.lang.reflect.Field;
import java.util.Map;
import org.springframework.web.socket.WebSocketSession;

@Component
public class WebSocketSessionManager {

    private final GraphQlWebSocketHandler graphQlWebSocketHandler;
    private Field sessionInfoMapField;

    public WebSocketSessionManager(GraphQlWebSocketHandler graphQlWebSocketHandler) throws NoSuchFieldException {
        this.graphQlWebSocketHandler = graphQlWebSocketHandler;
        // 获取私有sessionInfoMap字段并设置可访问
        this.sessionInfoMapField = GraphQlWebSocketHandler.class.getDeclaredField("sessionInfoMap");
        sessionInfoMapField.setAccessible(true);
    }

    public void closeAllSessions() throws IllegalAccessException {
        // 获取所有会话信息集合
        Map<String, ?> sessionInfoMap = (Map<String, ?>) sessionInfoMapField.get(graphQlWebSocketHandler);
        for (Object sessionInfo : sessionInfoMap.values()) {
            // 反射获取WebMvcSessionInfo内部的WebSocketSession实例
            Field sessionField = sessionInfo.getClass().getDeclaredField("session");
            sessionField.setAccessible(true);
            WebSocketSession session = (WebSocketSession) sessionField.get(sessionInfo);
            // 关闭活跃会话
            if (session.isOpen()) {
                session.close();
            }
        }
    }
}

2. 在事件触发时调用：

@Autowired
private WebSocketSessionManager sessionManager;

// 示例：在特定事件处理方法中调用
public void handleShutdownEvent() {
    try {
        sessionManager.closeAllSessions();
    } catch (IllegalAccessException e) {
        // 捕获并处理反射异常
        e.printStackTrace();
    }
}

方法二：自定义WebSocket拦截器跟踪会话 #

该方式通过拦截器主动维护活跃会话列表，不依赖框架内部私有结构，稳定性和兼容性更强，推荐使用。

实现代码 #

1. 自定义会话跟踪拦截器：

import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.WebSocketHandlerDecoratorFactory;
import org.springframework.stereotype.Component;
import java.util.concurrent.CopyOnWriteArrayList;

@Component
public class TrackableWebSocketHandlerDecoratorFactory implements WebSocketHandlerDecoratorFactory {

    // 线程安全的活跃会话列表
    private final CopyOnWriteArrayList<WebSocketSession> activeSessions = new CopyOnWriteArrayList<>();

    @Override
    public WebSocketHandler decorate(WebSocketHandler handler) {
        return new WebSocketHandler() {
            @Override
            public void afterConnectionEstablished(WebSocketSession session) throws Exception {
                activeSessions.add(session);
                handler.afterConnectionEstablished(session);
            }

            @Override
            public void handleMessage(WebSocketSession session, org.springframework.web.socket.WebSocketMessage<?> message) throws Exception {
                handler.handleMessage(session, message);
            }

            @Override
            public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
                activeSessions.remove(session);
                handler.handleTransportError(session, exception);
            }

            @Override
            public void afterConnectionClosed(WebSocketSession session, org.springframework.web.socket.CloseStatus closeStatus) throws Exception {
                activeSessions.remove(session);
                handler.afterConnectionClosed(session, closeStatus);
            }

            @Override
            public boolean supportsPartialMessages() {
                return handler.supportsPartialMessages();
            }
        };
    }

    public void closeAllSessions() {
        for (WebSocketSession session : activeSessions) {
            if (session.isOpen()) {
                try {
                    session.close();
                } catch (Exception e) {
                    // 处理单个会话关闭失败的异常
                    e.printStackTrace();
                }
            }
        }
        activeSessions.clear();
    }
}

2. 配置拦截器生效：

import org.springframework.context.annotation.Configuration;
import org.springframework.graphql.server.webmvc.GraphQlWebSocketHandler;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
public class WebSocketConfig implements WebSocketConfigurer {

    private final GraphQlWebSocketHandler graphQlWebSocketHandler;
    private final TrackableWebSocketHandlerDecoratorFactory decoratorFactory;

    public WebSocketConfig(GraphQlWebSocketHandler graphQlWebSocketHandler, TrackableWebSocketHandlerDecoratorFactory decoratorFactory) {
        this.graphQlWebSocketHandler = graphQlWebSocketHandler;
        this.decoratorFactory = decoratorFactory;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(decoratorFactory.decorate(graphQlWebSocketHandler), "/graphql")
                .setAllowedOrigins("*"); // 根据实际业务配置跨域规则
    }
}

3. 在事件触发时调用：

@Autowired
private TrackableWebSocketHandlerDecoratorFactory sessionTracker;

// 示例：在特定事件处理方法中调用
public void handleTerminateEvent() {
    sessionTracker.closeAllSessions();
}

注意事项 #

• 方法一依赖框架内部私有字段，升级Spring GraphQL版本前需验证内部结构是否变更；
• 方法二更符合标准扩展方式，稳定性和兼容性更优；
• 关闭会话时建议捕获异常，避免单个会话关闭失败影响其他会话；
• 跨域配置（setAllowedOrigins）需根据实际生产环境调整，避免使用通配符带来安全风险。

内容的提问来源于stack exchange，提问作者Mustard Tiger