基于SpringBoot WebSocket的一套消息推送系统

1、需求

需要实现业务系统实时发送任务消息至移动客户端，由于原单体架构基于socket的消息推送复杂难以在链接建立时获取客户信息、并发差难以实现分布式负载均衡。考虑采用Springboot-WebSocket+Kafka+Redis构建分布式高可用消息推送系统。

2、思路

整体思路如图所示，左边为WebSocket构建基本发送协议，根据业务系统产生的消息统一写入Kafka消息通知Topic，采用SpringBoot-WebSocket 消费Kafka中消息并发送到客户端：

消息发送的目标为客户端连接时传入的clientId，同时将ClientId写入外部Redis缓存作为当前在线用户，并检查Redis中是否存在此用户离线信息，如有则读取离线信息并发送；采用HashMap在单个WS推动程序中维护clientId-WsSession——当用户连入时保存其ClientId及其对应的WebSocketSession，当业务系统产生响应消息时，通过查询Redis中clientId Set判断是否用户当前在线，如在线取出其WebSocketSession发送对应信息，如不在线通过Redis存储离线信息（Key：clientId+Msg.timestamp，Value：离线信息，Expire：7Days）；当用户退出或发生错误时，移除Redis中用户Set对应的clientId，并删除HashMap中的Session键值对。

关于分布式，本文目前采用Redis外部缓存用户信息——当多个WebSocket实例运行时，可采用不同的Kafka Consumer GroupId以达到同时消费信息（同一条消息各实例都会消费到），由于用户对应的WebSocket Session只能存在于一个实例当中并且系统不知道是哪一个实例（这部分也可通过编写代理路由实现clientId对应唯一确认的实例），因此当有消息流入后，首先判断Redis中客户是否在线，如不在线写入key-value的离线信息（多个实例都会写入，由于是同一条信息因此最终redis也只存在一条），如用户在线，每个实例都会尝试从HashMap中取当前用户的WSSession，此时当然只会有一个实例获取到，然后发送信息。

更好的做法是当用户连接websocket时通过外部存储保存其当前连接指向那一台实例，保存clientId——WebSocketInstanceId，当消息流入后，根据clientId判断由哪一个实例消费信息（相应的KafkaCosnumer要做对应的路由处理），这样可以避免多个实例重复消费同一条消息。

3、Demo

1）Nginx实现WebSocket代理及负载均衡

Nginx较高版本开始支持WebSocket代理，下载Nginx后只需修改conf在http配置块中增加如下内容：

	#增加WebSocket代理设置
	map $http_upgrade $connection_upgrade{
		default upgrade;
		'' close;
	}
	
	
	upstream ws {
		server localhost:8888 weight=1; //websocket运行实例
		server localhost:9999 weight=1;
	}
	
	server{
		listen 8080;
		location /ws {
			proxy_pass http://ws;
			proxy_http_version 1.1;
			proxy_connect_timeout 5s; // 连接超时
			proxy_read_timeout 120s; // 超过此时间断开连接——采用1min的心跳保持连接
			proxy_send_timeout 15s;
			proxy_set_header Upgrade $http_upgrade;
			proxy_set_header Connection $connection_upgrade;
			
		}
	}

2）SpringBoot-WebSocket

kafka部分建立好对应的Topic并写入测试信息，此部分不赘述。

a) Pom文件主要部分如下：

    <groupId>com.test</groupId>
    <artifactId>WebSocketService</artifactId>
    <version>1.0-SNAPSHOT</version>

    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>2.0.0.RELEASE</version>
    </parent>

    <properties>
        <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
        <project.reporting.outputEncoding>UTF-8</project.reporting.outputEncoding>
        <java.version>1.8</java.version>
        <swagger2.version>2.7.0</swagger2.version>
        <beanutils.version>1.9.3</beanutils.version>
        <lang3.version>3.3.2</lang3.version>
        <start-class>com.test.Application</start-class>
        <skipTests>true</skipTests>
        <directory>./target</directory>
    </properties>

    <dependencies>
        <!--springBoot 启动jar-->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-web</artifactId>
            <exclusions>
                <exclusion>
                    <artifactId>jackson-databind</artifactId>
                    <groupId>com.fasterxml.jackson.core</groupId>
                </exclusion>
                <exclusion>
                    <groupId>org.springframework.boot</groupId>
                    <artifactId>spring-boot-starter-tomcat</artifactId>
                </exclusion>
            </exclusions>
        </dependency>

        <!--替换Tomcat Web容器为Undertow-->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-undertow</artifactId>
            <version>2.0.0.RELEASE</version>
        </dependency>

        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-actuator</artifactId>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-configuration-processor</artifactId>
            <optional>true</optional>
        </dependency>
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

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

        <dependency>
            <groupId>io.springfox</groupId>
            <artifactId>springfox-swagger2</artifactId>
            <version>${swagger2.version}</version>
        </dependency>
        <dependency>
            <groupId>io.springfox</groupId>
            <artifactId>springfox-swagger-ui</artifactId>
            <version>${swagger2.version}</version>
        </dependency>

        <dependency>
            <groupId>org.projectlombok</groupId>
            <artifactId>lombok</artifactId>
            <version>1.16.20</version>
        </dependency>

        <dependency>
            <groupId>org.apache.commons</groupId>
            <artifactId>commons-lang3</artifactId>
        </dependency>

        <!--kafka 依赖-->
        <dependency>
            <groupId>org.springframework.kafka</groupId>
            <artifactId>spring-kafka</artifactId>
        </dependency>
        <!-- https://mvnrepository.com/artifact/redis.clients/jedis -->
        <dependency>
            <groupId>redis.clients</groupId>
            <artifactId>jedis</artifactId>
        </dependency>

        <dependency>
            <groupId>com.alibaba</groupId>
            <artifactId>fastjson</artifactId>
            <version>1.2.54</version>
        </dependency>
    </dependencies>

 </project>

b) 配置部分，WebSocketConfig如下，Kafka Consumer及Redis连接池的配置限于篇幅就不贴了，网上相关资料比较多。

@Slf4j
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {
    @Autowired
    private MyHandShake handShake;
    @Autowired
    private MyHandler handler;

    /**
     * 实现 WebSocketConfigurer 接口，重写 registerWebSocketHandlers 方法，这是一个核心实现方法，配置 websocket 入口，允许访问的域、注册 Handler、SockJs 支持和拦截器。
     * <p>
     * registry.addHandler()注册和路由的功能，当客户端发起 websocket 连接，把 /path 交给对应的 handler 处理，而不实现具体的业务逻辑，可以理解为收集和任务分发中心。
     * <p>
     * addInterceptors，顾名思义就是为 handler 添加拦截器，可以在调用 handler 前后加入我们自己的逻辑代码。
     * <p>
     * setAllowedOrigins(String[] domains),允许指定的域名或 IP (含端口号)建立长连接，如果只允许自家域名访问，这里轻松设置。如果不限时使用”*”号，如果指定了域名，则必须要以 http 或 https 开头。
     * @param registry
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry){
        //部分 支持websocket 的访问链接,允许跨域
        registry.addHandler(handler,"/ws")
                .addInterceptors(handShake).setAllowedOrigins("*");
        //部分 不支持websocket的访问链接,允许跨域
        registry.addHandler(handler,"/sockjs/ws")
                .addInterceptors(handShake).setAllowedOrigins("*").withSockJS();
    }

}

c）编写请求处理类，Handler：

@Slf4j
@Service
public class MyHandShake implements HandshakeInterceptor {

    @Autowired
    private IJedisService jedisService;
    /**
     * 通过请求的WS路径参数：ws://ip:port/ws?parameter=xxx
     * 获取用户信息
     * （通过http session获取用户信息——需要在请求ws之前通过登录接口，保存服务器生成session，并共享session至ws服务）
     * @param request
     * @param response
     * @param webSocketHandler
     * @param attributes
     * @return
     * @throws Exception
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler webSocketHandler,
                                   Map<String,Object> attributes) throws Exception{
        HttpServletRequest servletRequest = null;
        if(request instanceof ServletServerHttpRequest){
            //Transfer WebSocketServer Request to HttpServletRequest to Get Session
            servletRequest = ((ServletServerHttpRequest) request).getServletRequest();
        }
        //Get User Identity from Session And Put it To ParamMap
        String userName = (String) servletRequest.getParameter(WebSocketConstants.ATTRI_USERID);  //ATTRI_USERID为自定义的ws连接参数字符串，比如 ws://localhost:8888/ws?userid=XXX 中的userid是也
        //TODO 判断userId是否允许连接websocket  if userName in Database 整体架构内，鉴权由网关层处理
//        if(jedisService.isUserSetContainsUser(USERS_KEY,userName)){
//            log.error("|----- 同一用户重复请求连接！ -----|");
//            return false;
//        }
        attributes.put(WebSocketConstants.ATTRI_USERID,userName);
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request,ServerHttpResponse response,WebSocketHandler webSocketHandler,Exception e){

    }
}

d）实现WebSocketSession发送的逻辑：

其中WebSocketConstant为常量Integer类型用于判断消息类型——0：心跳，1：业务消息。

LogFormatStringConstants.PREFIX_FORMAT 及LogFormatStringConstants.SUFFIX_FORMAT为自定义的特殊字符串用来标识比较重要的log行用。WebSocketMsg为封装的简单消息类，结构如下：

{
    "from": "消息来源",
    "msg": "消息内容",
    "timestamp":消息产生的时间戳,
    "to": "发送目标，与客户端的ATTR_USERID对应",
    "type": 消息类型，0—心跳，1—业务消息
}

@Slf4j
@Service
public class MyHandler implements WebSocketHandler {

    @Autowired
    private IJedisService iJedisService;

    /**
     * Static list to store exist websocket sessions
     */
    private final static Map<String,WebSocketSession> USER_SESSION_MAP = Collections.synchronizedMap(new HashMap<>());

    /**
     * Things to do after clients connected
     * @param webSocketSession
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession webSocketSession) throws Exception{
        String clientId = (String) webSocketSession.getAttributes().get(WebSocketConstants.ATTRI_USERID);
        log.info(LogFormatStringConstants.PREFIX_FORMAT+"Connected Successfully. User: "+clientId+LogFormatStringConstants.SUFFIX_FORMAT);

        if(USER_SESSION_MAP.containsKey(clientId)){
            WebSocketSession session = USER_SESSION_MAP.get(clientId);
            if(null != session && session.isOpen()){
                WebSocketMsg msg = new WebSocketMsg();
                msg.setType(WebSocketConstants.KICKED);
                msg.setFrom("server");
                msg.setTo(clientId);
                msg.setMsg("用户已在其他设备登录，本次连接终止");
                msg.setTimestamp(System.currentTimeMillis());
                session.sendMessage(new TextMessage(msg.toString()));
                log.info("发送消息至用户:"+clientId + ", "+msg.toString());
                session.close();
                USER_SESSION_MAP.remove(clientId);
            }
        }

        iJedisService.saveUserSet(clientId);
        USER_SESSION_MAP.put(clientId,webSocketSession);

        if(iJedisService.isExistOfflineMsg(clientId)){
            try{
                WebSocketSession session = USER_SESSION_MAP.get(clientId);
                if(null != session && session.isOpen()){
                    for(String msg : iJedisService.getOfflineMsg(clientId)){
                        session.sendMessage(new TextMessage(msg));
                        log.info("发送消息至用户:"+clientId + ", "+msg);
                    }
                }
            }catch (Exception e){
                e.printStackTrace();
            }finally {
                iJedisService.delOfflineMsg(clientId);
            }
        }
    }

    /**
     * Logic when new messages arrived
     * @param webSocketSession
     * @param message
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession webSocketSession, WebSocketMessage<?> message) throws Exception{
        log.info(LogFormatStringConstants.PREFIX_FORMAT+"Handling Messages to be sent"+LogFormatStringConstants.SUFFIX_FORMAT);
        //Parse socket message to JSONObject
        JSONObject msg = JSON.parseObject(message.getPayload().toString());
        log.info(LogFormatStringConstants.PREFIX_FORMAT+"Got Message :{}"+LogFormatStringConstants.SUFFIX_FORMAT,msg);
        sendMessage(msg);
    }

    /**
     * Message send logic
     * @param object
     */
    public void sendMessage(JSONObject object){
        JSONObject result = new JSONObject();
        try{
            WebSocketMsg webSocketMsg = WebSocketMsg.parseFromJson(object.toJSONString());
            if(webSocketMsg.getType() == WebSocketConstants.HEART_BEAT && webSocketMsg.getMsg().equals("ping")){
                String clientId = webSocketMsg.getFrom();
                WebSocketMsg response = new WebSocketMsg();
                response.setType(WebSocketConstants.HEART_BEAT);
                response.setFrom("server");
                response.setTo(clientId);
                response.setMsg("pong");
                response.setTimestamp(System.currentTimeMillis());
                if(iJedisService.getUserSet().contains(clientId)){
                    WebSocketSession session = USER_SESSION_MAP.get(clientId);
                    if(null != session && session.isOpen()){
                        session.sendMessage(new TextMessage(response.toString()));
                    }
                }
            }else if(webSocketMsg.getType() == WebSocketConstants.MESSAGE){
                String clientIds = webSocketMsg.getTo();
                for(String clientId : GlobalUtils.parseStringToList(clientIds)){
                    if(null !=iJedisService.getUserSet() && iJedisService.getUserSet().contains(clientId)){
                        WebSocketSession session = USER_SESSION_MAP.get(clientId);
                        if(null != session && session.isOpen()){
                            session.sendMessage(new TextMessage(webSocketMsg.toString()));
                            log.info("Sent Msg to Natural User: {}, msg: {}",
                                    clientId,webSocketMsg.toString());
                        }
                    }else {
                        iJedisService.saveOfflineMsg(clientId,webSocketMsg.getTimestamp(),webSocketMsg.toString());
                        log.info("Save Offline Msg to Natural User: {}, msg: {}",
                                clientId,webSocketMsg.toString());
                    }
                }
            }
        }catch (IOException e){
            e.printStackTrace();
        }
    }

    /**
     * Close connection when error occurs
     * @param webSocketSession
     * @param throwable
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession webSocketSession, Throwable throwable) throws Exception{
        String clientId = (String) webSocketSession.getAttributes().get(WebSocketConstants.ATTRI_USERID);
        if(webSocketSession.isOpen()){
            webSocketSession.close();
        }
        log.info(LogFormatStringConstants.PREFIX_FORMAT+"Connecting error, closing connection"+LogFormatStringConstants.PREFIX_FORMAT);
        USER_SESSION_MAP.remove(clientId);
        iJedisService.removeUser(clientId);
    }

    /**
     * Close connection normally
     * @param webSocketSession
     * @param closeStatus
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession webSocketSession, CloseStatus closeStatus) throws Exception{
        String client = (String) webSocketSession.getAttributes().get(WebSocketConstants.ATTRI_USERID);
        log.info(LogFormatStringConstants.PREFIX_FORMAT+"Connection closed: "+closeStatus.toString()+LogFormatStringConstants.PREFIX_FORMAT);
        USER_SESSION_MAP.remove(client);
        iJedisService.removeUser(client);
    }

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

}

e) kafka 消费者代码中引入sendMessage方法即可实现消费topic消息发送至相应用户：

@Slf4j
@Component
public class MsgConsumer {
    @Autowired
    private MyHandler handler;

    @KafkaListener(topicPattern = "${kafka.topic.testTopic}",containerFactory = "kafkaListenerContainerFactory")
    public void consumeAndPushToWS(ConsumerRecord<?,?> record){
        String data = "";
        try{
            Optional<?> kafkaMessage = Optional.ofNullable(record.value());
            if(kafkaMessage.isPresent()){
                Object message = kafkaMessage.get();
                data = message.toString();
                JSONObject jsonData = JSON.parseObject(data);
                handler.sendMessage(jsonData);
            }
        }catch (Exception e){
            e.printStackTrace();
        }
     }

}

服务端整体实现就到此结束了，相应的客户端为保持连接，需定时（2min内）发送给服务端心跳，上述代码中简单定义了很简单的心跳协议——ping、pong . 

原创不易转载请说明谢谢。

================================================================================================