webscoket實戰之利用httpsession定向推送

開發框架

springboot

場景

在利用websocket主動推送信息給客戶端的過程中，經常會遇到一個普遍需求，就是推送的消息要定向推送給不同的用戶，或者解釋的再普通一點，不同的消息推送給不同的session。例如一個用戶admin，可以在多台設備登錄，此時就有多個session，當一個設備向后台發起一個請求，處理時間較長（未采用客戶端進行ajax輪訓或者long poll時），采用websocket協議時，要針對不同session的操作定向返回操作結果。
下面來具體看看如何實現上述需求。

實例

springboot對websocket支持很友好，只需要繼承webSocketHandler類，重寫幾個方法就可以了

public class ResultHandler extends TextWebSocketHandler {
    private static final Logger logger = LoggerFactory.getLogger(ResultHandler.class);

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception { 
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketsession, CloseStatus status) throws Exception {
       
    }
}

然后寫一個簡單的配置類

@EnableWebSocket
@Configuration
public class WebSocketConfig implements WebSocketConfigurer{

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {
        webSocketHandlerRegistry.addHandler(resultHandler (), "/getResult").withSockJS();
    }

    @Bean
    public ResultHandler resultHandler(){
        return new ResultHandler();
    }

}

這樣一個websocket 服務端就寫好了

我們希望能夠把websocketSession和httpsession對應起來，這樣就能根據當前不同的session，定向對websocketSession進行數據返回
如何解決這個問題呢，考慮到websocket是在http的基礎上建立起來的（具體websocket建立，自行百度），能不能在websocket握手的時候，把當前的sessionId放入websocketSession的屬性中呢？

在查詢資料之后，發現spring中有一個攔截器接口，HandshakeInterceptor，可以實現這個接口，來攔截握手過程，向其中添加屬性

public class WebSocketHandshakeInterceptor implements HandshakeInterceptor{

    private static final Logger logger = LoggerFactory.getLogger(WebSocketHandshakeInterceptor.class);
    @Override
    public boolean beforeHandshake(ServerHttpRequest serverHttpRequest,ServerHttpResponse serverHttpResponse,
		 WebSocketHandler webSocketHandler, Map<String, Object> map) throws Exception {
        if(serverHttpRequest instanceof ServletServerHttpRequest){
            ServletServerHttpRequest servletRequest = (ServletServerHttpRequest) serverHttpRequest;
             HttpSession httpSession = servletRequest.getServletRequest().getSession(true);
            if(null != httpSession){
                map.put(Constants.SESSION_ID,httpSession.getId());
            }
        }
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest serverHttpRequest, 
		ServerHttpResponse serverHttpResponse, WebSocketHandler webSocketHandler, Exception e) {

    }
}

beforeHandShake方法中的Map參數 就是對應websocketSession里的屬性，向里面添加內容后，可以在上面的resultHandler里面利用websocketSession參數將其取出來 String sessionId = (String)webSocketsession.getAttributes().get(Constants.SESSION_ID);

參考下面代碼

public class ResultHandler extends TextWebSocketHandler {
    private static final Logger logger = LoggerFactory.getLogger(ResultHandler.class);

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
		
    }

    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception { 
	if(null != webSocketSession ){
            String sessionId  = (String)webSocketSession.getAttributes().get(Constants.SESSION_ID);
            addConnectionCount();
            logger.info("WebSocket connection established, sessionId={} ConnectCount={}", sessionId, getConnectionCount());
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketsession, CloseStatus status) throws Exception {
        if(null != webSocketsession ){
            String sessionId = (String)webSocketsession.getAttributes().get(Constants.SESSION_ID);
            SocketSessionManager.socketSessionMap.remove(sessionId);
            subConnectionCount();
            logger.info("WebSocket connection close, Status={}, connectionCount={}", status, getConnectionCount());
        }
    }
    public static synchronized  int getConnectionCount (){
        return ScriptTaskResultHandler.connectionCount;
    }

    private static synchronized void addConnectionCount(){
        ScriptTaskResultHandler.connectionCount++;
    }

    private static synchronized void subConnectionCount(){
        ScriptTaskResultHandler.connectionCount--;
    }
	
}

還可以自己寫一個webSocketSession管理類，當連接建立好之后，保存在這個管理類中，用sessionId當做Key，webSocketSession當做value，當結果處理完成后，根據sessionId找到相應的webSocketSession，進行數據推送。

攔截器加入的方法

@EnableWebSocket
@Configuration
public class WebSocketConfig implements WebSocketConfigurer{

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry webSocketHandlerRegistry) {
        webSocketHandlerRegistry.addHandler(resultHandler (), "/getResult").
        addInterceptors(new WebSocketHandshakeInterceptor()).withSockJS();
    }

    @Bean
    public ResultHandler resultHandler(){
        return new ResultHandler();
    }

}