Spring系列（7）-開發WebSocket的一點經驗

一、前言

在某些項目場景中，WebSocket是個利器，但畢竟常規應用場景很少。趁如今還記得些，把一些開發過程當中總結的一些經驗記下來，以避免過個一年半載再次須要用到時忘卻了。以前已經寫過一篇《WebSocket，再也不輪詢》，講了一些WebSocket的概念和應用場景，而本文此次偏實戰，講解的代碼會比較多一些。

代碼包括WebSocket的服務端和客戶端，以及如何寫WebSocket的單元測試。其中還會針對一些 「坑」 ，作重點分析。

二、WebSocket服務端

WebSocket服務端，即提供WebSocket服務的程序。SpringBoot開發WebSocket，常規有兩種方式 - 申明式和編程式，前者最簡單，我用的就是申明式。

2.一、pom.xml

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

2.二、註冊Bean

在包含@Configuration類（啓動類也包含該註解）中配置ServerEndpointExporter，配置後會自動註冊全部「@ServerEndpoint」註解聲明的Websocket Endpoint。

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

2.三、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.四、單例與多例的衝突

上述代碼中有一步是要調用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的實例都爲空。

三、WebSocket客戶端

通常不多有人在SpringBoot裏面寫WebSocket的客戶端，一般都是後端提供服務，前端來做爲客戶端通信。可是若是你的應用場景是後端之間的長鏈接交互，仍是會用到的。或者，當你須要給你的服務端寫單元測試時，這個後面再說。

3.一、pom.xml

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

3.二、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);
        }
    }

}

四、WebSocket單元測試

客戶要求咱們的SpringBoot程序發佈前，要經過sonar的質量檢查，其中有一項就是 「保證單元測試的覆蓋率超過50%」 。普通Http接口的單元測試咱們都知道，實在不會也能夠百度出來。但是你很難百度出來，WebSocket接口如何作單元測試？

後來我想，單元測試嘛，無非就是監聽後端服務的路由，調用一下程序的方法。那我能不能寫一個測試類，經過建立WebSocket客戶端的方式，模擬前端來測試服務端的邏輯？實際上我研究 「三、WebSocket客戶端「 ，就是爲了提升這個單元測試的覆蓋率。

4.一、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有下列四種：

1. MOCK（默認） ： 加載一個WebApplicationContext並提供一個模擬servlet環境。嵌入式servlet容器在使用此註釋時不會啓動。若是servlet API不在你的類路徑上，這個模式將透明地回退到建立一個常規的非web應用程序上下文。能夠與@AutoConfigureMockMvc結合使用，用於基於MockMvc的應用程序測試。
2. RANDOM_PORT ： 加載一個EmbeddedWebApplicationContext並提供一個真正的servlet環境。嵌入式servlet容器啓動並在隨機端口上偵聽。
3. DEFINED_PORT ： 加載一個EmbeddedWebApplicationContext並提供一個真正的servlet環境。嵌入式servlet容器啓動並監聽定義的端口（即從application.properties或默認端口8080）。
4. NONE ： 使用SpringApplication加載ApplicationContext，但不提供任何servlet環境（模擬或其餘）。

咱們在測試使用websocket時是須要完整的容器，因此能夠選 RANDOM_PORT或DEFINED_PORT。

4.二、測試類

爲了方便測試咱們使用 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%，個人任務達到了。