分布式下websocket+rabbitmq实现

一、Websocket简介

WebSocket是一种在单个TCP连接上进行全双工通信的协议。WebSocket通信协议于2011年被IETF定为标准RFC 6455，并由RFC7936补充规范。WebSocket API也被W3C定为标准。WebSocket使得客户端和服务器之间的数据交换变得更加简单，允许服务端主动向客户端推送数据。在WebSocket API中，浏览器和服务器只需要完成一次握手，两者之间就直接可以创建持久性的连接，并进行双向数据传输。(以上来自百度百科)

二、这里给出关键思路以及分布式下的解决方案：

问题：

        (1)不同的会话请求打到不同的Servlet容器中，可能存储在不同的Servlet容器中。单点下，可以从WebSocketSession缓存中获取，但是分布式下，就一定可以从本地获取。

        (2) WebSocketSession如何获取HttpSession的信息

解决方案：

(1)三次握手之前从ServerHttpRequest中获取HttpSession，将HttpSession中获取用户等业务信息放入WebSocketSession中。

import java.util.Map;
import javax.servlet.http.HttpSession;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

@Component
public class MessageHandshakeInterceptor implements HandshakeInterceptor {

    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
    	HttpSession session = this.getSession(request);
    	System.out.println("sessionUser:"+session.getAttribute("user"));
    	Long userId = (Long)session.getAttribute("user");
    	if(userId == null){
    		return false;
    	}
        attributes.put("uid", userId);
        return true;
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler, Exception exception) {

    }
    
    // 获取HttpSession
    private HttpSession getSession(ServerHttpRequest request) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest serverRequest = (ServletServerHttpRequest)request;
            return serverRequest.getServletRequest().getSession();
        } else {
            return null;
        }
    }
}

(2)保存会话信息到map中，将用户id作为key，将WebSocketSession作为value存储到本地hashMap中。WebSocketSession没有序列化，所以不能放到redis、mongodb等nosql中。

(3)当一个请求发送到nginx，转发到不同的tomcat servlet容器，每个servlet容器存储不同的WebSocketSession。从当前WebSocketSession中获取用户id，根据用户id从会话缓存中获取该用户的WebSocketSession。若本地会话缓存有，那么拿到WebSocketSession发送消息。

@Component
public class MessageHandler extends TextWebSocketHandler {

    @Autowired
    private MessageDAO messageDAO;
    @Autowired
    private RabbitSender rabbitSender;

    private static final ObjectMapper MAPPER = new ObjectMapper();

    public static final Map<Long, WebSocketSession> SESSIONS = new HashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");
        System.out.println("set session uid:"+uid);
        // 将当前用户的session放置到map中，后面会使用相应的session通信
        SESSIONS.put(uid, session);
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage
            textMessage) throws Exception {
        Long uid = (Long) session.getAttributes().get("uid");

        JsonNode jsonNode = MAPPER.readTree(textMessage.getPayload());
        Long toId = jsonNode.get("toId").asLong();
        String msg = jsonNode.get("msg").asText();
        System.out.println("toId:"+toId+",msg:"+msg);
        Message message = Message.builder()
                .from(UserData.USER_MAP.get(uid))
                .to(UserData.USER_MAP.get(toId))
                .msg(msg)
                .build();

        // 将消息保存到MongoDB
        message = this.messageDAO.saveMessage(message);

        // 判断to用户是否在线
        WebSocketSession toSession = SESSIONS.get(toId);
        if (toSession != null && toSession.isOpen()) {
        	System.out.println("get from local, toId:"+toId+",msg:"+msg);
            //TODO 具体格式需要和前端对接
            toSession.sendMessage(new
                    TextMessage(MAPPER.writeValueAsString(message)));
            // 更新消息状态为已读
            this.messageDAO.updateMessageState(message.getId(), 2);
        } else {
        	rabbitSender.sendMessage(message);
        }
    }
}

(4)若本地会话缓存没有找到，那么发送信息到mq中，去其他的servlet容器中查找会话信息，只要建立了会话，那么肯定在某个servlet容器中，找到某个用户的WebSocketSession，发送消息。

@Component
public class RabbitReceiver {
	
    @Autowired
    private MessageDAO messageDAO;
	
	private static final ObjectMapper MAPPER = new ObjectMapper();
	/**
	 * 
	 * 	spring.rabbitmq.listener.log.queue.name=queue-2
		spring.rabbitmq.listener.log.queue.durable=true
		spring.rabbitmq.listener.log.exchange.name=exchange-1
		spring.rabbitmq.listener.log.exchange.durable=true
		spring.rabbitmq.listener.log.exchange.type=topic
		spring.rabbitmq.listener.log.exchange.ignoreDeclarationExceptions=true
		spring.rabbitmq.listener.log.key=springboot.*
	 * @param log
	 * @param channel
	 * @param headers
	 * @throws Exception
	 */
	@RabbitListener(bindings = @QueueBinding(
			value = @Queue(value = "${spring.rabbitmq.listener.websocket.queue.name}", 
			durable="${spring.rabbitmq.listener.websocket.queue.durable}"),
			exchange = @Exchange(value = "${spring.rabbitmq.listener.websocket.exchange.name}", 
			durable="${spring.rabbitmq.listener.websocket.exchange.durable}", 
			type= "${spring.rabbitmq.listener.websocket.exchange.type}", 
			ignoreDeclarationExceptions = "${spring.rabbitmq.listener.websocket.exchange.ignoreDeclarationExceptions}"),
			key = "${spring.rabbitmq.listener.websocket.key}"
			)
	)
	@RabbitHandler
	public void onLogMessage(@Payload pojo.Message msg, Channel channel, 
			@Headers Map<String, Object> headers) throws Exception {
		Long deliveryTag = (Long)headers.get(AmqpHeaders.DELIVERY_TAG);
		//手工ACK
		channel.basicAck(deliveryTag, false);
		System.out.println("rabbitmq收到msg"+msg);
		Long toId = msg.getTo().getId();
		// 判断to用户是否在线
		WebSocketSession toSession = MessageHandler.SESSIONS.get(toId);
		if (toSession != null && toSession.isOpen()) {
			System.out.println("toSession: sessionId:"+toId);
			toSession.sendMessage(new
                    TextMessage(MAPPER.writeValueAsString(msg)));
			// 更新消息状态为已读
			this.messageDAO.updateMessageState(msg.getId(), 2);
		}
	}
}