【stomp 实战】spring websocket 接收消息源码分析

后台消息的发送过程，我们通过spring websocket用户消息发送源码分析已经了解了。我们再来分析一下后端接收消息的过程。这个过程和后端发送消息过程有点类似。

前端发送消息

前端发送消息给服务端的示例如下：

发送给目的/app/echo一个消息。

javascript 复制代码

//主动发送消息给服务器，对应的后端topic为/app/echo
function send() {
    var value = document.getElementById("content").value;
    var msg = {
        msgType: 1,
        content: value
    };
    stompClient.send("/app/echo", {}, JSON.stringify(msg));
    //stompClient.send("/app/echo2", {}, JSON.stringify(msg));
}

后端接收消息的配置

java 复制代码

 @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {
        /**
         * 这里表示前端往/app路径推送， 如果后端定义一个controller ->@MessageMapping("/echo"),
         * stompClient.send("/app/echo",{},...)
         * 这时，消息会被推送到注解对应的@MessageMapping("/echo")方法上
         */
        registry.setApplicationDestinationPrefixes("/app");
    }

后端配置/app前缀。

这个前缀和哪里结合起来用呢，来看下面的代码

java 复制代码

@Slf4j
@Controller
@RequiredArgsConstructor(onConstructor_ = {@Autowired})
public class StompController {
    private final SimpMessageSendingOperations msgOperations;
    private final SimpUserRegistry simpUserRegistry;

    /**
     * 回音消息，将用户发来的消息内容加上 Echo 前缀后推送回客户端
     */
    @MessageMapping("/echo")
    public void echo(Principal principal, Msg msg) {
        String username = principal.getName();
        msg.setContent("Echo: " + msg.getContent());
        msgOperations.convertAndSendToUser(username, "/topic/answer", msg);
        int userCount = simpUserRegistry.getUserCount();
        int sessionCount = simpUserRegistry.getUser(username).getSessions().size();
        log.info("当前本系统总在线人数: {}, 当前用户: {}, 该用户的客户端连接数: {}", userCount, username, sessionCount);
    }
}

实际上就和这个echo方法结合一起用的。 @MessageMapping("/echo")中的/echo和前缀结合一起，就是/app/echo。

因此，这个echo方法，就是接收前端发送消息的方法入口。

源码分析

消息处理器的注册

在 spring websocket源码分析之握手请求的处理这一节中，在完成websocket握手请求后，我们看到了如下的代码。

java 复制代码

	public void onOpen(final javax.websocket.Session session, EndpointConfig config) {
		this.wsSession.initializeNativeSession(session);

		// The following inner classes need to remain since lambdas would not retain their
		// declared generic types (which need to be seen by the underlying WebSocket engine)

		if (this.handler.supportsPartialMessages()) {
			session.addMessageHandler(new MessageHandler.Partial<String>() {
				@Override
				public void onMessage(String message, boolean isLast) {
					handleTextMessage(session, message, isLast);
				}
			});
			session.addMessageHandler(new MessageHandler.Partial<ByteBuffer>() {
				@Override
				public void onMessage(ByteBuffer message, boolean isLast) {
					handleBinaryMessage(session, message, isLast);
				}
			});
		}
		else {
			session.addMessageHandler(new MessageHandler.Whole<String>() {
				@Override
				public void onMessage(String message) {
					handleTextMessage(session, message, true);
				}
			});
			session.addMessageHandler(new MessageHandler.Whole<ByteBuffer>() {
				@Override
				public void onMessage(ByteBuffer message) {
					handleBinaryMessage(session, message, true);
				}
			});
		}

		session.addMessageHandler(new MessageHandler.Whole<javax.websocket.PongMessage>() {
			@Override
			public void onMessage(javax.websocket.PongMessage message) {
				handlePongMessage(session, message.getApplicationData());
			}
		});

		try {
			this.handler.afterConnectionEstablished(this.wsSession);
		}
		catch (Exception ex) {
			ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, ex, logger);
		}
	}

代码总结：

这里入参传了一个javax.websocket.Session。这个可以理解为当前Websocket连接。
原来这个Session可以给自己添加messageHandler，那当有消息来的时候，就会经过这些handler来进行处理。
那这个hander就是处理业务消息的重点了
看一下这个hander是怎么处理消息的

java 复制代码

private void handleTextMessage(javax.websocket.Session session, String payload, boolean isLast) {
		TextMessage textMessage = new TextMessage(payload, isLast);
		try {
			this.handler.handleMessage(this.wsSession, textMessage);
		}
		catch (Exception ex) {
			ExceptionWebSocketHandlerDecorator.tryCloseWithError(this.wsSession, ex, logger);
		}
	}

这个handler，对应的实现是：SockJsWebSocketHandler

进入handleMessage看一下处理逻辑，原来是将消息分为三类

文本消息
二进制消息
心跳消息
这三种消息，分别进行处理

java 复制代码

@Override
	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
		if (message instanceof TextMessage) {
			handleTextMessage(session, (TextMessage) message);
		}
		else if (message instanceof BinaryMessage) {
			handleBinaryMessage(session, (BinaryMessage) message);
		}
		else if (message instanceof PongMessage) {
			handlePongMessage(session, (PongMessage) message);
		}
		else {
			throw new IllegalStateException("Unexpected WebSocket message type: " + message);
		}
	}

我们一般处理的是文本消息

java 复制代码

	@Override
	public void handleTextMessage(WebSocketSession wsSession, TextMessage message) throws Exception {
		this.sockJsSession.handleMessage(message, wsSession);
	}

又交给sockJsSession来处理消息。

再看下WebSocketServerSockJsSession的handlerMessage方法。

往下，找到了delegateMessages。

java 复制代码

//WebSocketServerSockJsSession
	public void handleMessage(TextMessage message, WebSocketSession wsSession) throws Exception {
		String payload = message.getPayload();
		if (!StringUtils.hasLength(payload)) {
			return;
		}
		String[] messages;
		try {
			messages = getSockJsServiceConfig().getMessageCodec().decode(payload);
		}
		catch (Exception ex) {
			logger.error("Broken data received. Terminating WebSocket connection abruptly", ex);
			tryCloseWithSockJsTransportError(ex, CloseStatus.BAD_DATA);
			return;
		}
		if (messages != null) {
			delegateMessages(messages);
		}
	}

public void delegateMessages(String... messages) throws SockJsMessageDeliveryException {
		for (int i = 0; i < messages.length; i++) {
			try {
				if (isClosed()) {
					logUndeliveredMessages(i, messages);
					return;
				}
				this.handler.handleMessage(this, new TextMessage(messages[i]));
			}
			catch (Exception ex) {
				if (isClosed()) {
					if (logger.isTraceEnabled()) {
						logger.trace("Failed to handle message '" + messages[i] + "'", ex);
					}
					logUndeliveredMessages(i, messages);
					return;
				}
				throw new SockJsMessageDeliveryException(this.id, getUndelivered(messages, i), ex);
			}
		}
	}

可以看到delegateMessages实际上是把消息一条条处理。交给了handler来处理。

这里的hander是什么？SubProtocolWebSocketHandler。

java 复制代码

//SubProtocolWebSocketHandler
	@Override
	public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
		WebSocketSessionHolder holder = this.sessions.get(session.getId());
		if (holder != null) {
			session = holder.getSession();
		}
		SubProtocolHandler protocolHandler = findProtocolHandler(session);
		protocolHandler.handleMessageFromClient(session, message, this.clientInboundChannel);
		if (holder != null) {
			holder.setHasHandledMessages();
		}
		checkSessions();
	}

这里就是通过session取出子协议处理器，这里实际上就一个实现，是StompSubProtocolHandler。

java 复制代码

//StompSubProtocolHandler
	@Override
	public void handleMessageFromClient(WebSocketSession session,
			WebSocketMessage<?> webSocketMessage, MessageChannel outputChannel) {

		List<Message<byte[]>> messages;
		try {
			ByteBuffer byteBuffer;
			if (webSocketMessage instanceof TextMessage) {
				byteBuffer = ByteBuffer.wrap(((TextMessage) webSocketMessage).asBytes());
			}
			else if (webSocketMessage instanceof BinaryMessage) {
				byteBuffer = ((BinaryMessage) webSocketMessage).getPayload();
			}
			else {
				return;
			}

			BufferingStompDecoder decoder = this.decoders.get(session.getId());
			if (decoder == null) {
				if (!session.isOpen()) {
					logger.trace("Dropped inbound WebSocket message due to closed session");
					return;
				}
				throw new IllegalStateException("No decoder for session id '" + session.getId() + "'");
			}

			messages = decoder.decode(byteBuffer);
			if (messages.isEmpty()) {
				if (logger.isTraceEnabled()) {
					logger.trace("Incomplete STOMP frame content received in session " +
							session + ", bufferSize=" + decoder.getBufferSize() +
							", bufferSizeLimit=" + decoder.getBufferSizeLimit() + ".");
				}
				return;
			}
		}
		catch (Throwable ex) {
			if (logger.isErrorEnabled()) {
				logger.error("Failed to parse " + webSocketMessage +
						" in session " + session.getId() + ". Sending STOMP ERROR to client.", ex);
			}
			handleError(session, ex, null);
			return;
		}

		for (Message<byte[]> message : messages) {
			StompHeaderAccessor headerAccessor =
					MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
			Assert.state(headerAccessor != null, "No StompHeaderAccessor");

			StompCommand command = headerAccessor.getCommand();
			boolean isConnect = StompCommand.CONNECT.equals(command) || StompCommand.STOMP.equals(command);

			boolean sent = false;
			try {

				headerAccessor.setSessionId(session.getId());
				headerAccessor.setSessionAttributes(session.getAttributes());
				headerAccessor.setUser(getUser(session));
				if (isConnect) {
					headerAccessor.setUserChangeCallback(user -> {
						if (user != null && user != session.getPrincipal()) {
							this.stompAuthentications.put(session.getId(), user);
						}
					});
				}
				headerAccessor.setHeader(SimpMessageHeaderAccessor.HEART_BEAT_HEADER, headerAccessor.getHeartbeat());
				if (!detectImmutableMessageInterceptor(outputChannel)) {
					headerAccessor.setImmutable();
				}

				if (logger.isTraceEnabled()) {
					logger.trace("From client: " + headerAccessor.getShortLogMessage(message.getPayload()));
				}

				if (isConnect) {
					this.stats.incrementConnectCount();
				}
				else if (StompCommand.DISCONNECT.equals(command)) {
					this.stats.incrementDisconnectCount();
				}

				try {
					SimpAttributesContextHolder.setAttributesFromMessage(message);
					sent = outputChannel.send(message);

					if (sent) {
						if (this.eventPublisher != null) {
							Principal user = getUser(session);
							if (isConnect) {
								publishEvent(this.eventPublisher, new SessionConnectEvent(this, message, user));
							}
							else if (StompCommand.SUBSCRIBE.equals(command)) {
								publishEvent(this.eventPublisher, new SessionSubscribeEvent(this, message, user));
							}
							else if (StompCommand.UNSUBSCRIBE.equals(command)) {
								publishEvent(this.eventPublisher, new SessionUnsubscribeEvent(this, message, user));
							}
						}
					}
				}
				finally {
					SimpAttributesContextHolder.resetAttributes();
				}
			}
			catch (Throwable ex) {
				if (logger.isDebugEnabled()) {
					logger.debug("Failed to send message to MessageChannel in session " + session.getId(), ex);
				}
				else if (logger.isErrorEnabled()) {
					// Skip unsent CONNECT messages (likely auth issues)
					if (!isConnect || sent) {
						logger.error("Failed to send message to MessageChannel in session " + session.getId() +
								":" + ex.getMessage());
					}
				}
				handleError(session, ex, message);
			}
		}
	}

代码很长，总结一下：

1、消息报文的编码处理，转换成Message对象
2、StompHeaderAccessor的处理，包括设置user、session等
3、调用outputChannel发送消息：outputChannel.send(message);
4、如果发送消息成功，则发送相应的事件消息，有以下几类事件：SessionConnectEvent、SessionSubscribeEvent、SessionUnsubscribeEvent。

MessageChannel发送消息过程

outputChannel.send(message)，发送消息，这个似乎似曾相识。在【stomp 实战】spring websocket用户消息发送源码分析这一节中，我们也看到过这个类。在服务端往客户端发送消息时，也有这个MessageChannel的出现。

java 复制代码

//AbstractMessageChannel
	@Override
	public final boolean send(Message<?> message, long timeout) {
		Assert.notNull(message, "Message must not be null");
		Message<?> messageToUse = message;
		ChannelInterceptorChain chain = new ChannelInterceptorChain();
		boolean sent = false;
		try {
			messageToUse = chain.applyPreSend(messageToUse, this);
			if (messageToUse == null) {
				return false;
			}
			sent = sendInternal(messageToUse, timeout);
			chain.applyPostSend(messageToUse, this, sent);
			chain.triggerAfterSendCompletion(messageToUse, this, sent, null);
			return sent;
		}
		catch (Exception ex) {
			chain.triggerAfterSendCompletion(messageToUse, this, sent, ex);
			if (ex instanceof MessagingException) {
				throw (MessagingException) ex;
			}
			throw new MessageDeliveryException(messageToUse,"Failed to send message to " + this, ex);
		}
		catch (Throwable err) {
			MessageDeliveryException ex2 =
					new MessageDeliveryException(messageToUse, "Failed to send message to " + this, err);
			chain.triggerAfterSendCompletion(messageToUse, this, sent, ex2);
			throw ex2;
		}
	}

构造了一个拦截链，在发送前，可以进行前置处理和后置处理。这个拦截链就是扩展的关键了。我们可以定义自己的拦截器，在发送消息前后进行拦截处理。这里spring给我们的扩展点。
通过sendInternal将消息发送出去
然后我们Debug看看这个sendInternal

看到有三个MessageHandler
WebSocketAnnotationMethodMessageHandler
SimpleBrokerMessageHandler
UserDestinationMessageHandler
这里依次会调用这三个handler来发送消息。一般情况下，只会有一个handler来处理
我们示例中发送的消息destination是/app/echo，对应着一个方法。这里当然是WebSocketAnnotationMethodMessageHandler来处理了。
这里封装成一个Task，执行其run方法。在executor不为空的时候，是异步发送的。

进入SendTask，看一下run方法

java 复制代码

//
public void run() {
	Message<?> message = this.inputMessage;
	try {
		message = applyBeforeHandle(message);
		if (message == null) {
			return;
		}
		this.messageHandler.handleMessage(message);
		triggerAfterMessageHandled(message, null);
	}
	catch (Exception ex) {
		triggerAfterMessageHandled(message, ex);
		if (ex instanceof MessagingException) {
			throw (MessagingException) ex;
		}
		String description = "Failed to handle " + message + " to " + this + " in " + this.messageHandler;
		throw new MessageDeliveryException(message, description, ex);
	}
	catch (Throwable err) {
		String description = "Failed to handle " + message + " to " + this + " in " + this.messageHandler;
		MessageDeliveryException ex2 = new MessageDeliveryException(message, description, err);
		triggerAfterMessageHandled(message, ex2);
		throw ex2;
	}
}

这里的关键点是：this.messageHandler.handleMessage(message);

/app/echo会进入AbstractMethodMessageHandler

java 复制代码

// AbstractMethodMessageHandler
	@Override
	public void handleMessage(Message<?> message) throws MessagingException {
		String destination = getDestination(message);
		if (destination == null) {
			return;
		}
		String lookupDestination = getLookupDestination(destination);
		if (lookupDestination == null) {
			return;
		}

		MessageHeaderAccessor headerAccessor = MessageHeaderAccessor.getMutableAccessor(message);
		headerAccessor.setHeader(DestinationPatternsMessageCondition.LOOKUP_DESTINATION_HEADER, lookupDestination);
		headerAccessor.setLeaveMutable(true);
		message = MessageBuilder.createMessage(message.getPayload(), headerAccessor.getMessageHeaders());

		if (logger.isDebugEnabled()) {
			logger.debug("Searching methods to handle " +
					headerAccessor.getShortLogMessage(message.getPayload()) +
					", lookupDestination='" + lookupDestination + "'");
		}

		handleMessageInternal(message, lookupDestination);
		headerAccessor.setImmutable();
	}

	protected void handleMessageInternal(Message<?> message, String lookupDestination) {
		List<Match> matches = new ArrayList<>();

		List<T> mappingsByUrl = this.destinationLookup.get(lookupDestination);
		if (mappingsByUrl != null) {
			addMatchesToCollection(mappingsByUrl, message, matches);
		}
		if (matches.isEmpty()) {
			// No direct hits, go through all mappings
			Set<T> allMappings = this.handlerMethods.keySet();
			addMatchesToCollection(allMappings, message, matches);
		}
		if (matches.isEmpty()) {
			handleNoMatch(this.handlerMethods.keySet(), lookupDestination, message);
			return;
		}

		Comparator<Match> comparator = new MatchComparator(getMappingComparator(message));
		matches.sort(comparator);
		if (logger.isTraceEnabled()) {
			logger.trace("Found " + matches.size() + " handler methods: " + matches);
		}

		Match bestMatch = matches.get(0);
		if (matches.size() > 1) {
			Match secondBestMatch = matches.get(1);
			if (comparator.compare(bestMatch, secondBestMatch) == 0) {
				Method m1 = bestMatch.handlerMethod.getMethod();
				Method m2 = secondBestMatch.handlerMethod.getMethod();
				throw new IllegalStateException("Ambiguous handler methods mapped for destination '" +
						lookupDestination + "': {" + m1 + ", " + m2 + "}");
			}
		}

		handleMatch(bestMatch.mapping, bestMatch.handlerMethod, lookupDestination, message);
	}

handleMessage 主要做一些消息的处理
handleMessageInternal就是关键点了。
- 根据destination找到mappings，即我们注解中配置的url
- 正常情况下，会找到一个匹配的url，这个url会对应一个method，调用下面的方法执行后续逻辑。handleMatch(bestMatch.mapping, bestMatch.handlerMethod, lookupDestination, message)

java 复制代码

protected void handleMatch(T mapping, HandlerMethod handlerMethod, String lookupDestination, Message<?> message) {
		if (logger.isDebugEnabled()) {
			logger.debug("Invoking " + handlerMethod.getShortLogMessage());
		}
		handlerMethod = handlerMethod.createWithResolvedBean();
		InvocableHandlerMethod invocable = new InvocableHandlerMethod(handlerMethod);
		if (this.handlerMethodLogger != null) {
			invocable.setLogger(this.handlerMethodLogger);
		}
		invocable.setMessageMethodArgumentResolvers(this.argumentResolvers);
		try {
			Object returnValue = invocable.invoke(message);
			MethodParameter returnType = handlerMethod.getReturnType();
			if (void.class == returnType.getParameterType()) {
				return;
			}
			if (returnValue != null && this.returnValueHandlers.isAsyncReturnValue(returnValue, returnType)) {
				ListenableFuture<?> future = this.returnValueHandlers.toListenableFuture(returnValue, returnType);
				if (future != null) {
					future.addCallback(new ReturnValueListenableFutureCallback(invocable, message));
				}
			}
			else {
				this.returnValueHandlers.handleReturnValue(returnValue, returnType, message);
			}
		}
		catch (Exception ex) {
			processHandlerMethodException(handlerMethod, ex, message);
		}
		catch (Throwable ex) {
			Exception handlingException =
					new MessageHandlingException(message, "Unexpected handler method invocation error", ex);
			processHandlerMethodException(handlerMethod, handlingException, message);
		}
	}

这里最重要的就是 invocable.invoke(message);。即调用反射来执行目标方法。这里代码之所以比较复杂，是处理入参和返回值。这里不是我们研究的重点。就不再分析了。

整个流程总结如下