Springboot版本的WebSocket向app端推送消息实现

业务背景：app端需要实时接收到服务端的交易价格信息，肯定无法实时主动去轮询拉取pull操作，需要一种push实时数据到app端，因此websocket技术正好应用于此。

技术栈选型：

1 传统型 使用@ServerEndpoint注解实现

2 前沿性 使用springboot 三个核心 一配二拦三监听

一配指的 是注册并配置websocket的handler监听器和拦截器

二拦指的 是拦截客户端请求的websocket的参数和权鉴服务

三监听指的 是监听各种事件(请求连接/接收消息/关闭连接)

为什么要用springboot集成的websocket，主要是因为springboot对websocket做了很好的支持，而且提供了在建立连接前可以权鉴和做一些其他的事情。

代码实现：

maven 主要依赖：

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

一 配置服务

/**
 * 注册websocket服务
 * @author panmeng
 */
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(new EpaiWebSocketHandler(),"/websocket")
                .setAllowedOrigins("*")
                .addInterceptors(new HandShakeInterceptor())
        ;
    }
}

二 拦截器 （可以接收url的数据）

@Slf4j
public class HandShakeInterceptor extends HttpSessionHandshakeInterceptor {
    /**
     * 在WebSocket连接建立之前的操作，以鉴权为例
     * @param request
     * @param response
     * @param wsHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        log.info("Handle before webSocket connected. ");
        // 获取url传递的参数，通过attributes在Interceptor处理结束后传递给WebSocketHandler
        // WebSocketHandler可以通过WebSocketSession的getAttributes()方法获取参数
        ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
        String catId = serverRequest.getServletRequest().getParameter("object_id");
        if (TokenValidation.validateSign()) {
            log.info("Validation passed. WebSocket connecting.... ");
            attributes.put("catId", catId);
            return super.beforeHandshake(request, response, wsHandler, attributes);
        } else {
            log.error("Validation failed. WebSocket will not connect. ");
            return false;
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception ex) {
        // 省略
    }
}

三 监听 并处理

主要代码实现是通过一个商品的id作为key，然后每一个查看当前商品的用户都会给他建立一个websession会话连接，记录下来，然后当服务端有最新的交易数据过来时，调用当前的sendWebsocketMessage dubbo接口即可实现将最新的交易消息发送到各个连接客户端。

/**
 * @author panmeng
 */
@Slf4j
@Service("websocketService")
public class EpaiWebSocketHandler implements WebSocketHandler, WebsocketService {

    private static ConcurrentMap<String,List<WebSocketSession>> existSocketClientMap = new ConcurrentHashMap<>();
    /**
     * 当前在线数
     */
    private static int onlineCount = 0;
    
    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception {
        log.info("success connect the web socket server...the remote sessionId is {}",webSocketSession.getId());
        Map<String, Object> attributes = webSocketSession.getAttributes();
        Object catAucId = attributes.get("catId");
        if(catAucId!=null&& !StringUtils.isEmpty(catAucId.toString())){
            List<WebSocketSession> historySessionList= existSocketClientMap.get(catAucId.toString());
            List<WebSocketSession> sessionList=new ArrayList<>();
            sessionList.add(webSocketSession);
            if(!CollectionUtils.isEmpty(historySessionList)){
                sessionList.addAll(historySessionList);
            }
            existSocketClientMap.put(catAucId.toString(),sessionList);
        }else{
            throw new BusinessException(GlobalConstant.CAT_AUC_ID_PARAM_NOT_NULL);
        }
        addOnlineCount();
        log.info("current connect total count is :{}",onlineCount);
    }

    private static synchronized void addOnlineCount() {
        EpaiWebSocketHandler.onlineCount++;
    }


    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> webSocketMessage) throws Exception {
        String msg = webSocketMessage.getPayload().toString();
        log.info("the web socket receive client message,info is {} prepare start send...",msg);
        BaseMessageInfo baseMessageInfo = JSONObject.parseObject(msg, BaseMessageInfo.class);
        Integer catId=null;
        if(baseMessageInfo.getData() !=null){
            JSONObject data= (JSONObject) baseMessageInfo.getData();
            if(data !=null &&  data.get("catId") !=null){
                catId= Ints.tryParse(data.get("catId").toString());
            }
        }
        if(catId ==null){
            throw new BusinessException(GlobalConstant.CAT_INFO_NOT_FIND_ERROR);
        }
        log.info("the web socket receive client message,catId is {}",catId);
        List<WebSocketSession> historySessionList= existSocketClientMap.get(catId.toString());
        if(!CollectionUtils.isEmpty(historySessionList)){
            log.info("the web socket should send client count is {}",historySessionList.size());
            historySessionList.stream().forEach(session->{
                try {
                    if(session.isOpen()){
                        session.sendMessage(new TextMessage(msg));
                    }else{
                        log.info("the session {} is already closed,now should remove this session from history memory",session);
                        afterConnectionClosed(session,CloseStatus.NORMAL);
                    }
                } catch (IOException e) {
                    log.error("the sessionId {} an error occurred while sending the message...",session.getId());
                    e.printStackTrace();
                } catch (Exception e) {
                    log.error("the sessionId {} an error occurred while the session closed...",session.getId());
                    e.printStackTrace();
                }
            });
        }
        log.info("the web socket receive clint message,info is {} send to all complete...",msg);
    }

    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) throws Exception {
        log.info("the connection happen unknow error,the remote address is {}",webSocketSession.getRemoteAddress());
        if (webSocketSession.isOpen()) {
            webSocketSession.close();
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) throws Exception {
        log.info("the connection will close,sessionId is:{}",webSocketSession.getId());
        if (webSocketSession.isOpen()) {
            webSocketSession.close();
        }
        Iterator<Map.Entry<String, List<WebSocketSession>>> iterator = existSocketClientMap.entrySet().iterator();
        while (iterator.hasNext()){
            Map.Entry<String, List<WebSocketSession>> entry = iterator.next();
            List<WebSocketSession> sessionList = entry.getValue();
            if(sessionList.contains(webSocketSession)){
                sessionList.remove(webSocketSession);
            }
        }
        log.info("the connection already closed,sessionId is:{}",webSocketSession.getId());
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    @Override
    public ReturnJson sendWebsocketMessage(String msg) {
        log.info("the web socket receive clint message,info is {}",msg);
        BaseMessageInfo baseMessageInfo = JSONObject.parseObject(msg, BaseMessageInfo.class);
        Integer catId=null;
        if(baseMessageInfo.getData() !=null){
            JSONObject data= (JSONObject) baseMessageInfo.getData();
            if(data !=null &&  data.get("catId") !=null){
                catId= Ints.tryParse(data.get("catId").toString());
            }
        }
        if(catId ==null){
            throw new BusinessException(GlobalConstant.CAT_INFO_NOT_FIND_ERROR);
        }
        log.info("the web socket receive clint message,catId is {}",catId);
        List<WebSocketSession> historySessionList= existSocketClientMap.get(catId.toString());
        if(!CollectionUtils.isEmpty(historySessionList)){
            historySessionList.stream().forEach(session->{
                try {
                    if(session.isOpen()){
                        session.sendMessage(new TextMessage(msg));
                    }else{
                        log.info("the session {} is already closed,now should remove this session from history memory",session);
                        afterConnectionClosed(session,CloseStatus.NORMAL);
                    }
                } catch (IOException e) {
                    e.printStackTrace();
                } catch (Exception e) {
                    e.printStackTrace();
                }
            });
        }
        log.info("the web socket receive clint message,info is {} send to all complete...",msg);
        return ReturnJson.success();
    }
}