SpringBoot整合WebSocket实现在线聊天室

前言

WebSocket也是一种应用层协议，也是建立在TCP协议之上，类似HTTP，并且兼容HTTP。相比HTTP，它可以实现双向通信，如聊天室场景，使用HTTP就必须客户端轮训查询服务器有没有新的消息，而使用WebSocket就可以服务器直接通知客户端。

Tomcat支持

Tomcat自7.0.5版本开始支持WebSocket，并实现了WebSocket规范（JSR356）。JSR356规定WebSokcet应用由一系列Endpoint组成，类似于Servlet，Tomcat支持两种定义Endpoint的方式。

注解方式

import java.io.IOException;
import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;

@ServerEndpoint("/chat")
public class ChatEndpoint {

  @OnOpen
  public void onOpen(Session session) {
    System.out.println(session.getId());
  }

  @OnClose
  public void onClose(Session session) {
    System.out.println(session.getId());
  }

  @OnError
  public void onError(Session session) {
    System.out.println(session.getId());
  }

  @OnMessage
  public void onMessage(Session session, String msg) throws IOException {
    System.out.println(msg);
    //向客户端发送消息
    session.getBasicRemote().sendText("this is " + session.getId());
  }
}

使用注解@ServerEndpoint声明Endpoint类，并配置请求路径。

编程方式

import javax.websocket.CloseReason;
import javax.websocket.Endpoint;
import javax.websocket.EndpointConfig;
import javax.websocket.MessageHandler.Whole;
import javax.websocket.Session;

public class ChatEndpoint2 extends Endpoint {

  @Override
  public void onOpen(Session session, EndpointConfig config) {
    System.out.println(session.getId());
    //添加消息处理器
    session.addMessageHandler(new Whole<>() {
      @Override
      public void onMessage(Object message) {
        System.out.println(message);
      }
    });
  }

  @Override
  public void onClose(Session session, CloseReason closeReason) {
    System.out.println(closeReason.getReasonPhrase());
  }

  @Override
  public void onError(Session session, Throwable throwable) {
    System.out.println(throwable.getMessage());
  }
}

定义一个类继承Endpoint。

import java.util.HashSet;
import java.util.Set;
import javax.websocket.Endpoint;
import javax.websocket.server.ServerApplicationConfig;
import javax.websocket.server.ServerEndpointConfig;

public class MyServerApplicationConfig implements ServerApplicationConfig {

  @Override
  public Set<ServerEndpointConfig> getEndpointConfigs(Set<Class<? extends Endpoint>> scanned) {
    //根据查询到的Endpoint实现类创建ServerEndpointConfig
    Set<ServerEndpointConfig> result = new HashSet<>();
    if (scanned.contains(ChatEndpoint2.class)) {
      result.add(ServerEndpointConfig.Builder.create(ChatEndpoint2.class, "/chat").build());
    }
    return result;
  }

  @Override
  public Set<Class<?>> getAnnotatedEndpointClasses(Set<Class<?>> scanned) {
    //用来过滤使用注解ServerEndpoint定义的Endpoint
    return scanned;
  }
}

还需要定义一个ServerApplicationConfig的实现类，用来配置Endpoint的请求路径。

Endpoint生命周期

Endpoint的生命周期方法如下：

• onOpen：当开启一个新的会话时调用，这是客户端与服务器握手成功后调用的方法，等同于注解@OnOpen。
• onClose：会话关闭时调用，等同于注解@OnClose。
• onError：传输过程异常时调用，等同于注解@OnError。

原理

Tomcat会在启动时查找所有使用注解@ServerEndpoint声明的类和Endpoint的子类，创建WebSocketContainer(类似ServletContext)并将所有Endpoint注入其中。
但如果使用嵌入式的Tomcat(如SpringBoot内嵌的Tomcat)就不会进行此查找，具体原因可以看Embedded Tomcat does not honor ServletContainerInitializers。根本原因是因为SpringBoot创建Tomcat的Context时没有添加ContextConfig这个LifecycleListener(不清楚是出于什么考虑)，ContextConfig会使用java的SPI技术查找所有ServletContainerInitializer的实现类。

SpringBoot支持

添加依赖

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

SpringBoot自动注入了一个TomcatWebSocketServletWebServerCustomizer，向Context中添加了WsContextListener监听器

它也会创建WebSocketContainer并初始化。

配置处理器

import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

public class ChatHandler2 extends TextWebSocketHandler {

  @Override
  public void afterConnectionEstablished(WebSocketSession session) throws Exception {
    System.out.println(session.getId());
  }

  @Override
  public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
    System.out.println(session.getId());
  }

  @Override
  public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {

  }

  @Override
  protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
    System.out.println(message.getPayload());
  }

}

类似于继承Endpoint的类实现。

添加配置

import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;

@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

  @Override
  public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
    registry.addHandler(new ChatHandler2(), "/chat");
  }

}

一定要添加@EnableWebSocket注解，它会处理所有WebSocketConfigurer的实现类。并且会创建WebSocketHandlerMapping的Bean类，之前我们处理HTTP请求时使用的是RequestMappingHandlerMapping，WebSocket请求是在HTTP请求的基础上升级的，在握手阶段还是HTTP请求，所以还是会交给SpringMVC的DispatcherServlet处理。

实现原理

使用的HandlerAdapter实现为HttpRequestHandlerAdapter，它会持有一个WebSocketHttpRequestHandler对象(其中封装了HandshakeHandler握手处理器和我们自己定义的ChatHandler2)。实际的握手处理器实现为DefaultHandshakeHandler，核心的处理逻辑就是握手的过程，其中会创建StandardWebSocketHandlerAdapter对象。

可以看到StandardWebSocketHandlerAdapter就是一个Endpoint，最终会将这个Endpoint添加到WebSocketContainer中，后续的WebSocket请求就交给Endpoint来处理了。

聊天室实现效果

项目地址springboot_chatroom，页面部分参考ChatRoom项目。

参考

WebSocket 教程
学习WebSocket协议—从顶层到底层的实现原理（修订版）
Embedded Tomcat does not honor ServletContainerInitializers
Tomcat实现Web Socket
websocket之三：Tomcat的WebSocket实现