springboot2.0 websocket连接和集群

1. 引入websocket的starter

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

这里引入rabbitmq的ampq starter是为了websocket服务的集群，集群之间主要通过rabbitmq来通信。

2. 配置websocket的Exporter

包含一个异步线程池，用于异步处理rabbitmq消息

@Configuration
public class WebSocketConfig {

    @Bean
    public ServerEndpointExporter serverEndpointExporter(){
        return new ServerEndpointExporter();
    }

    @Bean("taskExecutor")
    public Executor taskExecutor(){
        ThreadPoolTaskExecutor taskExecutor = new ThreadPoolTaskExecutor();
        taskExecutor.setWaitForTasksToCompleteOnShutdown(true);//线程池关闭的时候等待所有任务都完成再继续销毁其他的Bean
        taskExecutor.setAwaitTerminationSeconds(60);//设置线程池中任务的等待时间，如果超过这个时候还没有销毁就强制销毁，以确保应用最后能够被关闭，而不是阻塞住
        taskExecutor.setCorePoolSize(4);
        taskExecutor.setMaxPoolSize(8);
        taskExecutor.setQueueCapacity(Integer.MAX_VALUE);//缓冲执行任务的队列
        taskExecutor.setKeepAliveSeconds(60);//当超过了核心线程出之外的线程在空闲时间到达之后会被销毁
        taskExecutor.setThreadNamePrefix("kanjia-websocket-thread-");
        return taskExecutor;
    }
}

websocket的rabbitmq集群配置文件

websocket.rabbitmq.queue.kanjia.routing-key=kanjia.#
#每个集群服务器监听一个有UUID的队列
websocket.rabbitmq.queue.kanjia.name=websocket.kanjia.${random.uuid}
websocket.rabbitmq.exchange.topic=websocket.topic

配置rabbitmq的yml

spring:
  rabbitmq:
    host: ${rabbitmq.host}
    port: ${rabbitmq.port}
    username: ${rabbitmq.i5x.username}
    password: ${rabbitmq.i5x.password}
    virtual-host: ${rabbitmq.i5x.vhost}

3. 配置websocket的endpoint

@Controller
@ServerEndpoint("/ws/{biz}/{key}") 
public class WebSocket {
    private static Logger log = LoggerFactory.getLogger(WebSocket.class);
    private static ConcurrentMap <String, CopyOnWriteArrayList<WebSocket>> webSocketMap = new ConcurrentHashMap<String, CopyOnWriteArrayList<WebSocket>>();
    private Session session;
    private String address;//消息的地址:业务区分.key，例如duorendiancai.4991_123
    //WebSocket和RabbitMq的消息互通bridge
    private WebSocketRabbitMqBridge webSocketRabbitMqBridge = SpringContext.getBean(WebSocketRabbitMqBridge.class);


    @OnOpen
    public void onOpen(@PathParam("biz") String biz, @PathParam("key") String key, Session session) {
        //参数合法性check
        if (StringUtils.isBlank(biz) || StringUtils.isBlank(key)) {
            String msg = "websocket连接参数不合法，biz=" + biz + ",key=" + key;
            log.error(msg);
            try {
                session.close(new CloseReason(CloseReason.CloseCodes.CANNOT_ACCEPT, msg));
            } catch (IOException e) {
                log.error("关闭websocket连接异常:" + e.toString());
            }
        }

        //保存websocket连接
        this.address = biz + "." + key;
        this.session = session;

        //按address区分websocket的session连接
        CopyOnWriteArrayList<WebSocket> webSocketList = webSocketMap.get(address);
        if (webSocketList == null) {
            CopyOnWriteArrayList<WebSocket> tempList = new CopyOnWriteArrayList<WebSocket>();
            tempList.add(this);
            webSocketMap.put(address, tempList);
        } else {
            webSocketList.add(this);
        }
    } 

    @OnMessage
    public void onMessage(String message, Session session) {
        try {
            //message合法性check
            JSONObject msgJson = JSON.parseObject(message);
            if (msgJson != null) {
                msgJson.put(KanjiaMsgHandler.ADDRESS_KEY_MQ_MSG, this.address);
                //发送mq消息
                webSocketRabbitMqBridge.sendRabbitMqMsg(this.address, msgJson.toJSONString());
                log.info("server receive msg:address=" + this.address);
            } else {
                log.error("server receive msg=" + message + "，address=" + this.address);
                return;
            }
        } catch (Exception e) {
            log.error("处理接收到的信息异常：" + e.toString() + "，msg=" + message + "，address=" + this.address);
        }
    }

    @OnClose
    public void onClose(Session session, CloseReason closeReason) {
        CopyOnWriteArrayList<WebSocket> webSocketList = webSocketMap.get(this.address);
        if (webSocketList != null) {
            webSocketList.remove(this);
            if (webSocketList.isEmpty()) {
                webSocketMap.remove(this.address);
            }
        }
        log.info("onClose: address=" + this.address + ",id=" + session.getId() + ",reason=" + closeReason.getReasonPhrase() );
    }

    //连接错误时执行
    @OnError
    public void onError(Throwable t) {
        log.error("websocket onError:" + t.toString());
    }


    /**
     * publish <br/>
     * 广播消息 <br/>
     *
     * @author Mobile Web Group-lff
     * @date 2018年2月11日 上午10:12:27
     *
     * @param address
     * @param message
     * @return void
     */
    public static void publish(String address, String message){
        CopyOnWriteArrayList<WebSocket> webSocketList = webSocketMap.get(address);
        if (webSocketList == null) {
            return;
        }
        for (WebSocket webSocket : webSocketList) {
            try {
                //发送消息
                webSocket.session.getBasicRemote().sendText(message);
            } catch (IOException e) {
                //输出log，继续下一个webSocket的msg发送
                log.error("发送消息失败：原因=" + e.toString() + ",id=" + webSocket.session.getId() + ",msg=" + message);
            }
        }
    }

    @Override
    public int hashCode() {
        final int prime = 31;
        int result = 1;
        result = prime * result + ((this.session == null) ? 0 : this.session.getId().hashCode());
        return result;
    }

    @Override
    public boolean equals(Object obj) {
        if (this == obj) {
            return true;
        }
        if (obj == null) {
            return false;
        }
        if (!(obj instanceof WebSocket)) {
            return false;
        }
        WebSocket other = (WebSocket) obj;
        if (this.session == null) {
            if (other.session != null) {
                return false;
            }
        } else if (!this.session.getId().equals(other.session.getId())) {
            return false;
        }
        return true;
    }

onOpen：连接成功的回调函数。安照url连接中的参数，拼成一个websocket的地址address，可以代表着一个聊天室或者一个群聊。利用这个address来转发publish消息。

onMessage：接收到消息的回调函数。这里为了多服务器部署时的集群，通过rabbitmq来过渡消息。接收到websocket消息后，通过rabbitmq转发给所有的websocket的服务器。

onClose：连接关闭的回调函数。

onError：错误回调。

hashCode/equals：为了定位session必须要重写。

自定义的publish方法：根据websocket的地址address转发消息。

4. websocket和rabbitmq消息过渡的桥梁

用来发送websocket和rabbitmq消息

@Component
public class WebSocketRabbitMqBridge {

    //websocket和rabbitmq传递消息的exchange
    @Value("${websocket.rabbitmq.exchange.topic}")
    private String websocket_rabbitmq_exchange_topic;

    @Autowired
    private RabbitTemplate rabbitTemplate;


    /**
     * sendRabbitMqMsg <br/>
     * 将websocket消息转发到rabbitmq上，进行服务间传播 <br/>
     *
     * @author Mobile Web Group-lff
     * @date 2018年2月12日 上午9:02:19
     *
     * @param exchange
     * @param routingKey
     * @param object
     * @return void
     */
    public void sendRabbitMqMsg(String routingKey, final Object object) {
        rabbitTemplate.convertAndSend(websocket_rabbitmq_exchange_topic, routingKey, object);
    }


    /**
     * publish <br/>
     * 广播websocket消息 <br/>
     *
     * @author Mobile Web Group-lff
     * @date 2018年2月11日 下午3:09:21
     *
     * @param address
     * @param msg
     * @return void
     */
    public void publishWebSocketMsg(String address, String msg) {
        WebSocket.publish(address, msg);
    }


}

5. 监听处理rabbitmq消息

创建一个含有uuid的队列并监听，通过异步任务处理消息

@Component
@RabbitListener(bindings = @QueueBinding(
                value = @Queue(value = "${websocket.rabbitmq.queue.kanjia.name}", autoDelete = "true"),
                exchange = @Exchange(value = "${websocket.rabbitmq.exchange.topic}", type = ExchangeTypes.TOPIC, durable="true"),
                key = "${websocket.rabbitmq.queue.kanjia.routing-key}")
                )
public class RabbitMqKanJiaMsgReceiver {
    @Autowired
    private KanjiaMsgHandler kanjiaMsgHandler;


    /**
     * processMsg <br/>
     * 处理mq消息，消息格式： <br/>
     *
     * @author Mobile Web Group-lff
     * @date 2018年2月11日 下午3:11:04
     *
     * @param msg
     * @return void
     */
    @RabbitHandler
    public void processMsg(String msg) {
        kanjiaMsgHandler.handler(msg);
    }

}

异步任务类

@Component
public class KanjiaMsgHandler {
    //mq消息中的address key，代表websocket的地址
    public static final String ADDRESS_KEY_MQ_MSG = "_address";
    private static Logger log = LoggerFactory.getLogger(KanjiaMsgHandler.class);
    //WebSocket和RabbitMq的消息互通bridge
    @Autowired
    private WebSocketRabbitMqBridge webSocketRabbitMqBridge;

    @Async("taskExecutor")
    public void handler(String msg) {
        if (StringUtils.isBlank(msg)) {
            //空msg
            return;
        }
        //msg格式：{_address:'kanjia.1234567', type:'barrage', body:{nickname:,photoImg:,placeholder:}}
        try {
            JSONObject msgJson = JSON.parseObject(msg);
            String address = msgJson.getString(ADDRESS_KEY_MQ_MSG);
            webSocketRabbitMqBridge.publishWebSocketMsg(address, msg);
            log.info("server send msg:address=" + address);
        } catch (Exception e) {
            log.error("消息处理异常：" + e.toString() + ",msg=" + msg);
        }
    }

}

6. 其它设置的代码

由于用到了异步任务，需要通过@EnableAsync开启异步功能

@SpringCloudApplication
@EnableAsync
public class Application {

    public static void main(String[] args) {
        SpringApplication.run(Application.class, args);
    }

}

另外，由于@ServerEndpoint中不能用@Autowired等注入方式注入其它Bean，所以需要通过ApplicationContex显示获取其它Bean，SpringContext的util类

@Component
public class SpringContext implements ApplicationContextAware {
    private static Logger log = LoggerFactory.getLogger(SpringContext.class);
    private static ApplicationContext applicationContext;

    /**
     * Description: 获取bean的管理器
     * @author: Mobile Web Group-lff
     * @date: 2018年2月11日 上午10:52:58
     * 
     * @see org.springframework.context.ApplicationContextAware#setApplicationContext(org.springframework.context.ApplicationContext)
     */
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        log.info("setApplicationContext，applicationContext=" + applicationContext);
        SpringContext.applicationContext = applicationContext;
    }

    //获取applicationContext
    private static ApplicationContext getApplicationContext() {
        return applicationContext;
    }

    //通过name获取 Bean.
    public static Object getBean(String name){
        return getApplicationContext().getBean(name);
    }

    //通过class获取Bean.
    public static <T> T getBean(Class<T> clazz){
        return getApplicationContext().getBean(clazz);
    }

    //通过name,以及Clazz返回指定的Bean
    public static <T> T getBean(String name,Class<T> clazz){
        return getApplicationContext().getBean(name, clazz);
    }

}

另外需要注意：
1. @ServerEndpoint类上不加@Controller的话，websocket无法连接
2. websocket的Session默认没有超时时间session.getMaxIdleTimeout()=0；如果发现过一段时间后，websocket连接会自动断掉，应该是nginx的配置proxy_read_timeout 90；proxy_send_timeout 90；