SpringBoot实现WebSocket服务端

1. 什么是WebSocket

WebSocket协议是基于TCP的一种新的网络协议。它实现了浏览器与服务器全双工(full-duplex)通信——允许服务器主动发送信息给客户端。

2. SpringBoot如何集成WebSocket

SpringBoot集成WebSocket有2种方式

• 第一种是用“@ServerEndPoint”注解来实现，实现简单；
• 第二种是实现WebSocketHandler，可以添加拦截器在WebSocket连接建立和断开前进行一些额外操作，稍显麻烦。
  2种方式都需要添加依赖包：

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

2.1 新建一个maven项目命名为websocket_server

pom文件如下：

<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0"
         xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
         xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
    <modelVersion>4.0.0</modelVersion>

    <groupId>com.owner</groupId>
    <artifactId>websocket_server</artifactId>
    <version>1.0-SNAPSHOT</version>

    <parent>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-parent</artifactId>
        <version>2.0.6.RELEASE</version>
    </parent>

    <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>
</project>

2.2 新增application.yml文件

server:
  port: 8001
#  设置服务名
spring:
  application:
    name: websocket-server

2.3 新增启动类

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

2.4 @ServerEndPoint集成WebSocket

WebSocket 业务逻辑实现如下

@ServerEndpoint("/websocket/{id}/{name}")
@RestController
public class WebSocketController {
    // 用来记录当前连接数的变量
    private static volatile int onlineCount = 0;

    // concurrent包的线程安全Set，用来存放每个客户端对应的MyWebSocket对象
    private static CopyOnWriteArraySet<WebSocketController> webSocketSet = new CopyOnWriteArraySet<>();

    // 与某个客户端的连接会话，需要通过它来与客户端进行数据收发
    private Session session;

    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketController.class);

    @OnOpen
    public void onOpen(Session session, @PathParam("id") long id, @PathParam("name") String name) throws Exception {
        this.session = session;
        System.out.println(this.session.getId());
        webSocketSet.add(this);
        LOGGER.info("Open a websocket. id={}, name={}", id, name);

        sendMessage("hello websocket");
    }

    @OnClose
    public void onClose() {
        webSocketSet.remove(this);
        LOGGER.info("Close a websocket. ");
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        LOGGER.info("Receive a message from client: " + message);
    }

    @OnError
    public void onError(Session session, Throwable error) {
        LOGGER.error("Error while websocket. ", error);
    }

    public void sendMessage(String message) throws Exception {
        if (this.session.isOpen()) {
            this.session.getBasicRemote().sendText(message);
        }
    }

    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        WebSocketController.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        WebSocketController.onlineCount--;
    }
}

@Configuration
public class WebSocketConfig {

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

完成以上步骤以后启动服务，打开websocket在线测试工具，尝试连接服务器

2.5 实现WebSocketHandler方式集成WebSocket

WebSocket 业务逻辑实现如下

@RestController
public class WebSocketController implements WebSocketHandler {
    
    private static AtomicInteger onlineCount = new AtomicInteger(0);
    
    private static final ArrayList<WebSocketSession> sessions = new ArrayList<>();
    
    private final Logger LOGGER = LoggerFactory.getLogger(WebSocketController.class);
    
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        sessions.add(session);
        int onlineNum = addOnlineCount();
        LOGGER.info("Open a WebSocket. Current connection number: " + onlineNum);
    }
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        sessions.remove(session);
        int onlineNum = subOnlineCount();
        LOGGER.info("Close a webSocket. Current connection number: " + onlineNum);
    }

    @Override
    public void handleMessage(WebSocketSession wsSession, WebSocketMessage<?> message) throws Exception {
        LOGGER.info("Receive a message from client: " + message.toString());
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        LOGGER.error("Exception occurs on webSocket connection. disconnecting....");
        if (session.isOpen()) {
            session.close();
        }
        sessions.remove(session);
        subOnlineCount();
    }

    /*
     * 是否支持消息拆分发送：如果接收的数据量比较大，最好打开(true), 否则可能会导致接收失败。
     * 如果出现WebSocket连接接收一次数据后就自动断开，应检查是否是这里的问题。
     */
    @Override
    public boolean supportsPartialMessages() {
        return true;
    }

    public static int getOnlineCount() {
        return onlineCount.get();
    }
    
    public static int addOnlineCount() {
        return onlineCount.incrementAndGet();
    }
    
    public static int subOnlineCount() {
        return onlineCount.decrementAndGet();
    }
}

public class WebSocketInterceptor extends HttpSessionHandshakeInterceptor {

    private final Logger LOGGER = LoggerFactory.getLogger(WebSocketInterceptor.class);

    /*
     * 在WebSocket连接建立之前的操作，以鉴权为例
     */
    @Override
    public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response,
                                   WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
        LOGGER.info("Handle before webSocket connected. ");

        // 获取url传递的参数，通过attributes在Interceptor处理结束后传递给WebSocketHandler
        // WebSocketHandler可以通过WebSocketSession的getAttributes()方法获取参数
        ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
        String id = serverRequest.getServletRequest().getParameter("id");
        String name = serverRequest.getServletRequest().getParameter("name");

        if (StringUtils.isEmpty(id) || StringUtils.isEmpty(name)) {
            LOGGER.error("Validation failed. WebSocket will not connect. ");
            return false;
        } else {
            LOGGER.info("Validation passed. WebSocket connecting.... ");
            attributes.put("id", id);
            attributes.put("name", name);
            return super.beforeHandshake(request, response, wsHandler, attributes);
        }
    }

    @Override
    public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response,
                               WebSocketHandler wsHandler, Exception ex) {
        // 省略
    }
}

@Configuration
@EnableWebMvc
@EnableWebSocket
public class WebSocketConfig extends WebMvcConfigurerAdapter implements WebSocketConfigurer {

    @Autowired
    private WebSocketController controller;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(controller, "/testWebSocket")
                .addInterceptors(new WebSocketInterceptor()).setAllowedOrigins("*");
    }
}

完成以上步骤以后启动服务，打开websocket在线测试工具，尝试连接服务器