SpringBoot开发系列（7）-开发WebSocket的一点经验-阿里云开发者社区

1、前言

在某些项目场景中，WebSocket是个利器，但毕竟常规应用场景不多。趁现在还记得些，把一些开发过程中总结的一些经验记下来，以免过个一年半载再次需要用到时忘却了。之前已经写过一篇《WebSocket，不再轮询》，讲了一些WebSocket的概念和应用场景，而本文这次偏实战，讲解的代码会比较多一些。

代码包括WebSocket的服务端和客户端，以及如何写WebSocket的单元测试。其中还会针对一些 “坑” ，做重点分析。

2、WebSocket服务端

WebSocket服务端，即提供WebSocket服务的程序。SpringBoot开发WebSocket，常规有两种方式 - 申明式和编程式，前者最简单，我用的就是申明式。

2.1、pom.xml

        <!--websocket 服务端-->
        <dependency>
            <groupId>org.springframework.boot</groupId>
            <artifactId>spring-boot-starter-websocket</artifactId>
        </dependency>

2.2、注册Bean

在包含@Configuration类（启动类也包含该注解）中配置ServerEndpointExporter，配置后会自动注册所有“@ServerEndpoint”注解声明的Websocket Endpoint。

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

2.3、WebSocket服务端类

根据上文，写WebSocket的类需要通过“@ServerEndpoint”注解声明。
MyWebSocketService .java

@Component
@ServerEndpoint(value = "/xxx/{userId}")
@Slf4j
public class MyWebSocketService {
    private String userId = "anonymous";
    private static int onlineCount = 0;
    private Session session;
    private static ConcurrentHashMap<String, Session> sessionPool = new ConcurrentHashMap<>();


    @OnOpen
    public void onOpen(Session curSession, @PathParam("userId") String curUserId) {
        this.session = curSession;
        this.userId = curUserId;
        sessionPool.put(curUserId, curSession);
        addOnlineCount();
        log.info(curUserId + "有一连接加入！当前在线人数为" + onlineCount);
    }

    /**
     * 连接关闭调用的方法
     */
    @OnClose
    public void onClose() {
        if (sessionPool.get(this.userId) != null) {
            sessionPool.remove(userId);
            subOnlineCount();
            log.info(userId + "有一连接关闭！当前在线人数为" + getOnlineCount());
        }
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message 客户端发送过来的消息
     */
    @OnMessage
    public void onMessage(String message) {
        handleMessage(message);
    }

    /**
     * @param curSession
     * @param error
     */
    @OnError
    public void onError(Session curSession, Throwable error) {
        log.error(error.getMessage(), error);
    }

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

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

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

    /**
     * 只发送给单一用户
     *
     * @param curUserId
     * @param socketMessage
     */
    public void sendMessageSingle(String curUserId, SocketMessage socketMessage) {
        Session curSession = sessionPool.get(curUserId);
        if (curSession != null) {
            try {
                String response = JSON.toJSONString(socketMessage);
                curSession.getBasicRemote().sendText(response);
            } catch (Exception e) {
                log.error(e.getMessage(), e);
            }
        }
    }

  
    /**
     * 处理客户端发送的消息
     * websocket 初始化早，无法注入Bean
     *
     * @param message
     * @return
     */
    private void handleMessage(String message) {
        try {
            SocketMessage request = JSON.parseObject(message, SocketMessage.class);
            switch (ModuleEnum.valueOf(request.getModule())) {
                case HEART_CHECK:
                    this.session.getBasicRemote().sendText(
                            JSON.toJSONString(new SocketMessage(ModuleEnum.HEART_CHECK.name(), "回复心跳检查")));
                    break;
                case ACTION_MAP_SWITCH:
                    MapMapper mapMapper =ApplicationContextRegister.getApplicationContext().getBean(MapMapper.class);
                    mapMapper.updateAllBranchVisible();
                    sendMessageSingle(chlWeb, request);
                    break;
                 //case 等等，其他处理逻辑
                default:
                    break;
            }

        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

}

2.4、单例与多例的冲突

上述代码中有一步是要调用dao层的方法，handleMessage方法中。

//...
                case ACTION_MAP_SWITCH:
                    MapMapper mapMapper =ApplicationContextRegister.getApplicationContext().getBean(MapMapper.class);
                    mapMapper.updateAllBranchVisible();
//...

正常我们开发SpringBoot，都是借助Spring容器的IOC的特性，将Service、Dao等直接依赖注入，类似于下面。

    @Autowired
    private MapMapper mapMapper;

//...
                case ACTION_MAP_SWITCH:
                    mapMapper.updateAllBranchVisible();
//...

但是这么写会报错，在执行 mapMapper.updateAllBranchVisible(); 方法时报空指针，即MapMapper的Bean没有注入进来。所以本文是通过Spring容器上下文，用工厂类的方式创建MapMapper的Bean。
ApplicationContextRegister.java

@Component
public class ApplicationContextRegister implements ApplicationContextAware {
    private static ApplicationContext applicationContext;

    @Override
    public void  setApplicationContext(ApplicationContext curApplicationContext) {
        applicationContext = curApplicationContext;
    }

    public static ApplicationContext getApplicationContext() {
        return applicationContext;
    }

}

要弄懂原因，首先要了解Spring注入Bean的方式：

有些人可能不知道，Spring默认实例化的Bean是单例模式，这就意味着在Spring容器加载时，就注入了MapMapper的实例，不管再调用多少次接口，加载的都是这个Bean同一个实例。

而WebSocket是多例模式，在项目启动时第一次初始化实例时，MapMapper的实例的确可以加载成功，但可惜这时WebSocket是无用户连接的。当有第一个用户连接时，WebSocket类会创建第二个实例，但由于Spring的Dao层是单例模式，所以这时MapMapper对应的实例为空。后续每连接一个新的用户，都会再创建新的WebSocket实例，当然MapMapper的实例都为空。

3、WebSocket客户端

一般很少有人在SpringBoot里面写WebSocket的客户端，通常都是后端提供服务，前端来作为客户端通讯。但是如果你的应用场景是后端之间的长连接交互，还是会用到的。或者，当你需要给你的服务端写单元测试时，这个后面再说。

3.1、pom.xml

        <!--websocket 客户端-->
        <dependency>
            <groupId>org.java-websocket</groupId>
            <artifactId>Java-WebSocket</artifactId>
            <version>1.3.8</version>
        </dependency>

3.2、WebSocket客户端类

配置WebSocket客户端的方法更简单，继承并实现WebSocketClient 类。

MyWebSocketClient.java

import lombok.extern.slf4j.Slf4j;
import org.java_websocket.client.WebSocketClient;
import org.java_websocket.handshake.ServerHandshake;

@Slf4j
public class MyWebSocketClient extends WebSocketClient {
public MyWebSocketClient(URI uri){
    super(uri);
}

    @Override
    public void onOpen(ServerHandshake serverHandshake) {
        log.info("客户端连接成功");
    }

    @Override
    public void onMessage(String s) {
        log.info("客户端接收到消息："+s);
    }

    @Override
    public void onClose(int i, String s, boolean b) {
        log.info("客户端关闭成功");
    }

    @Override
    public void onError(Exception e) {
        log.error("客户端出错");
    }

    public static void main(String[] args) {
        try {
            MyWebSocketClient myWebSocketClient = new MyWebSocketClient(new URI("ws://localhost:9000/xxx/user1"));
            myWebSocketClient.connect();
            while (!WebSocket.READYSTATE.OPEN.equals(myWebSocketClient.getReadyState())) {
                log.info("WebSocket客户端连接中，请稍等...");
                Thread.sleep(500);
            }
            myWebSocketClient.send("{\"module\":\"HEART_CHECK\",\"message\":\"请求心跳\"}");
            myWebSocketClient.close();
        } catch (Exception e) {
            log.error("error", e);
        }
    }

}

4、WebSocket单元测试

客户要求我们的SpringBoot程序发布前，要通过sonar的质量检查，其中有一项就是 “保证单元测试的覆盖率超过50%” 。普通Http接口的单元测试我们都知道，实在不会也可以百度出来。可是你很难百度出来，WebSocket接口如何做单元测试？

后来我想，单元测试嘛，无非就是监听后端服务的路由，调用一下程序的方法。那我能不能写一个测试类，通过创建WebSocket客户端的方式，模拟前端来测试服务端的逻辑？实际上我研究 “3、WebSocket客户端“ ，就是为了提高这个单元测试的覆盖率。

4.1、WebEnvironment

我们在写Junit的测试类时，通常都会如下文一样，通过@SpringBootTest获取启动类，加载SpringBoot配置。但是如果我们的项目里面有WebSocket，这样会报无法启动WebSocket的错误。

@RunWith(SpringRunner.class)
@SpringBootTest
public class CompositeControllerTest{
    @Test
    public void websocketClient() {
        int num = new Integer(1);
        Assert.assertEquals(num, 1);
    }
}

@SpringBootTest注解实际有一个webEnvironment的属性，SpringBootTest.WebEnvironment有下列四种：

MOCK（默认）：加载一个WebApplicationContext并提供一个模拟servlet环境。嵌入式servlet容器在使用此注释时不会启动。如果servlet API不在你的类路径上，这个模式将透明地回退到创建一个常规的非web应用程序上下文。可以与@AutoConfigureMockMvc结合使用，用于基于MockMvc的应用程序测试。
RANDOM_PORT ：加载一个EmbeddedWebApplicationContext并提供一个真正的servlet环境。嵌入式servlet容器启动并在随机端口上侦听。
DEFINED_PORT ：加载一个EmbeddedWebApplicationContext并提供一个真正的servlet环境。嵌入式servlet容器启动并监听定义的端口（即从application.properties或默认端口8080）。
NONE ：使用SpringApplication加载ApplicationContext，但不提供任何servlet环境（模拟或其他）。

我们在测试使用websocket时是需要完整的容器，所以可以选 RANDOM_PORT或DEFINED_PORT。

4.2、测试类

为了方便测试我们使用 SpringBootTest.WebEnvironment.DEFINED_PORT，监听固定的端口。

@RunWith(SpringRunner.class)
@SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.DEFINED_PORT)
@AutoConfigureMockMvc
@Slf4j
class CompositeControllerTest {
    @Autowired
    private MockMvc mockMvc;
    @Autowired
    private WebApplicationContext webApplicationContext;

    @Before
    public void before(){
        mockMvc= MockMvcBuilders.webAppContextSetup(webApplicationContext).build();
    }

// ...  等等

    /**
     * 创建 WebSocket 的客户端做测试
     * @throws Exception
     */
    @Test
    void websocketClient() throws Exception{
        MyWebSocketClient myWebSocketClient = new MyWebSocketClient(new URI("ws://localhost:port/xxxx/user1"));
        myWebSocketClient.connect();
        while (!WebSocket.READYSTATE.OPEN.equals(myWebSocketClient.getReadyState())){
            log.info("WebSocket客户端连接中，请稍等...");
            Thread.sleep(500);
        }
        Map<String,String> requestMap=new HashMap<>();
        requestMap.put("HEART_CHECK","{\"module\":\"HEART_CHECK\",\"message\":\"请求心跳\"}");
        requestMap.put("KEY1","VALUE1");
        requestMap.put("KEY2","VALUE2");
        requestMap.put("KEY3","VALUE3");

        for(String key: requestMap.keySet()){
            myWebSocketClient.send(requestMap.get(key));
        }
        //测试 onError、onMessage、onClose
        // ...  等等
        myWebSocketClient.close();
    }
}

OK，在最后生成的sonar测试报告里面，我们可以看到WebSocket的代码基本都被覆盖到，单元测试的覆盖率提高到了90%，我的任务达到了。

文献参考

1.spring boot整合Websocket笔记

https://yq.aliyun.com/articles/637898?spm=a2c4e.11153940.0.0.537f1d36hwhtLi

SpringBoot开发系列（7）-开发WebSocket的一点经验