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實現