SpringBoot中使用websocket.md

最近有這樣一個需求，網關廠家將物聯設備接入我司雲平臺的時候，希望能看到上報設備數據的關鍵日誌，以方便調試。
首先想到的就是使用websocket推送。瀏覽器發起websocket連接，發送訂閱消息，然後往這個連接session中推送日誌。

整個設計流程如下圖：

1.實現

我們設計兩個類，一個類命名爲WebSocketServer 用來管理websocket連接以及發送消息；另一個類命名爲WebSocketBus用來管理WebSocketServer 對象，以及接收設備日誌後匹配對應的WebSocketServer 對象，將日誌信息推送到瀏覽器。

Talk is cheap， show me the code!

@ServerEndpoint(value="/websocket/message")
@Component
public class WebSocketServer {
    
    private Logger log = Logger.getLogger("WebSocket");
    
    private WebSocketBus webSocketBus;

    
    private Session session;//與某個客戶端的連接會話，需要通過它來給客戶端發送數據
    
    private String key;//訂閱日誌的標識
    
    
    

    /**
     * 連接建立成功調用的方法*/
    @OnOpen
    public void onOpen(Session session) {
        this.session = session;  
         //在線數加1
   
        String success = "websocket連接成功！";
        try {
            sendMessage(success);
        } catch (IOException e) {
            log.error(e,e);
        }
    }

    /**
     * 連接關閉調用的方法
     */
    @OnClose
    public void onClose() {
        webSocketBus.removeServer(key,this);  //從set中刪除
        log.info("============= 有一連接關閉！key=" + key+",sessionId="+session.getId());
    }

    /**
     * 收到客戶端消息後調用的方法
     *
     * @param message 客戶端發送過來的消息*/
    @OnMessage
    public void onMessage(String message, Session session) {
        if(StringUtil.isEmpty(message)) {
            return;
        }
        HashMap<String,String> javaObject = JacksonUtil.toJavaObject(message, HashMap.class);
        if(javaObject == null){
             log.info("unknown message: " + message);
             return ;
         }
        String operation = javaObject.get("operation");
        String key = javaObject.get("key");
        String uuid=(String)javaObject.get("uuid");
        if(StringUtil.isEmpty(operation)){
            log.info("unknown operation : " + message);
            return ;
        }
        
        if("register".equals(operation)){
            this.key = key;
            getWebSocketBus().addServer(key, this);
    
            log.info("[+]     register key="+key+",account="+",uuid="+uuid);
        }else if("unRegister".equals(operation)){
            webSocketBus.removeServer(key, this); 
            log.info("[+]     unregister key="+key+",account="+",uuid="+uuid);
        }
             
        
    }

    private WebSocketBus getWebSocketBus() {
        // TODO Auto-generated method stub
        if(this.webSocketBus == null) {
            webSocketBus =  (WebSocketBus) SpringContextUtil.getBean("webSocketBus");
        }
        return webSocketBus;
        
    }

    /**
     * 發生錯誤時調用
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error(error,error);
    }


    public void sendMessage(String message) throws IOException {
//        this.session.getBasicRemote().sendText(message);
        this.session.getAsyncRemote().sendText(message);
    }


    public String getkey() {
        return key;
    } 
}

@Component
public class WebSocketBus {
    private Logger log = Logger.getLogger(this.getClass());
    
    @Autowired
    private KafkaTemplate<String,String> kafkaTemplate;
    
        
    @Value(value="${LogRegistTopic}")
    private String logRegister;//訂閱|取消訂閱topic
    
    private static AtomicInteger onlineCount = new AtomicInteger(0);
    
    private Map<String,Set<WebSocketServer>> sessionCache = new ConcurrentHashMap<>();
    private Gson gson = new GsonBuilder().create();//線程安全的，大膽用吧
    
    
    @KafkaListener(id = "log",topics = {"${LogTopic}"})
    public void listen(ConsumerRecord<String, ?> record) {
        Optional kafkaMessage = Optional.ofNullable(record.value());
        Optional<String> kafkaKey = Optional.ofNullable(record.key());
        if (kafkaKey.isPresent()) {
            Object value = kafkaMessage.get();
            String key = kafkaKey.get();
            GatewayFormatLog gatewayLog = gson.fromJson((String)value, GatewayFormatLog.class);
            if(sessionCache.containsKey(key)) {
                Set<WebSocketServer> set = sessionCache.get(key);
                for(WebSocketServer server :set) {
                    try {
                            server.sendMessage(gatewayLog.getMessage().toString());
                    } catch (IOException e) {
                        log.error(e,e);
                    }
                }
            }
        }
    }
    
    
    
    public void addServer(String key,WebSocketServer server) {
        boolean notRegisted = true;
        if(sessionCache.containsKey(key)) {
            Set<WebSocketServer> set = sessionCache.get(key);
            if(set.contains(server)){
                notRegisted = false;
            }else {
                set.add(server);
            }
            
        }else {
            Set<WebSocketServer> set = new  CopyOnWriteArraySet<WebSocketServer>();
            set.add(server);
            sessionCache.put(key, set);
        }
        if(notRegisted) {
            kafkaTemplate.send(logRegister, gson.toJson(new DataEvent("register",key)));
        }
        
    }
    
    public void removeServer(String key,WebSocketServer server) {
        if(sessionCache.containsKey(key)) {
            Set<WebSocketServer> set = sessionCache.get(key);
            set.remove(server);
            kafkaTemplate.send(logRegister, gson.toJson(new DataEvent("unregister",key)));
        }
    }
    
    public static int getOnlineCount() {
        return onlineCount.get();
    }

    public static void addOnlineCount() {
        onlineCount.incrementAndGet();
    }

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

}

WebSocketServer負責建立連接，以及收發websocket消息。瀏覽器每發起一個連接，都對應一個WebSocketServer對象。
WebSocketBus管理多個WebSocketServer對象。執行過程如下：
1.WebSocketServer的onMessage接收訂閱消息，並將訂閱消息交給WebSocketBus處理
2.1 如果是訂閱消息，執行WebSocketBus.addServer，同一個key對應一個或多個websocket連接（就是多個客戶端訂閱了同一個key）。我們用CopyOnWriteArraySet存放WebSocketServer，如果是同一個對象，則不會重複添加。如果這個key沒有被訂閱過，就往kafka中發一條訂閱消息。設備服務消費
2.2 如果是取消訂閱服務，過程類似，往kafka中發一條消息訂閱消息。

3.設備服務會將訂閱的日誌發送到Kafka中。WebSocketBus的listen來消費，根據key匹配到多個WebSocketServer，向每一個WebSocketServer推送消息。

下面是前端的測試代碼：

<!DOCTYPE html>
<html>
<head>
    <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
    <title>WebSocket/SockJS)</title>

    <script src="http://cdn.sockjs.org/sockjs-0.3.min.js"></script>

    <script type="text/javascript">
        var websocket = null;

        //判斷當前瀏覽器是否支持WebSocket
        if('WebSocket' in window){
            websocket = new WebSocket("ws://localhost:9100/websocket/message");
        }
        else{
            alert('Not support websocket')
        }

        //連接發生錯誤的回調方法
        websocket.onerror = function(){
            setMessageInnerHTML("error");
        };

        //連接成功建立的回調方法
        websocket.onopen = function(event){
            setMessageInnerHTML("open");
        }

        //接收到消息的回調方法
        websocket.onmessage = function(){
            setMessageInnerHTML(event.data);
        }

        //連接關閉的回調方法
        websocket.onclose = function(){
            setMessageInnerHTML("close");
        }

        //監聽窗口關閉事件，當窗口關閉時，主動去關閉websocket連接，防止連接還沒斷開就關閉窗口，server端會拋異常。
        window.onbeforeunload = function(){
            websocket.close();
        }

        //將消息顯示在網頁上
        function setMessageInnerHTML(innerHTML){
            document.getElementById('message').innerHTML += innerHTML + '<br/>';
        }

        //關閉連接
        function closeWebSocket(){
            websocket.close();
        }

        //發送消息
        function send(){
            var message = document.getElementById('text').value;
            websocket.send(message);
        }

    </script>
</head>
<body>
Welcome<br/>
<input id="text" type="text" />
<button onclick="send()">Send</button>
<button onclick="closeWebSocket()">Close</button>
<div id="message"></div>
</body>
</html>

2. 使用@ServerEndpoint無法注入Bean

你可能注意到了，WebSocketServer中使用WebSocketBus時，並沒有使用@Autowired，爲什麼呢？實際上使用@Autowired注入之後，沒有注入成功，使用時webSocketBus還是爲null。
我們在WebSocketServer類上使用了@Component註解。雖然@Component默認是單例模式的，但springboot還是會爲每個websocket連接初始化一個bean。
查了一下源碼：

public class DefaultServerEndpointConfigurator
        extends ServerEndpointConfig.Configurator {

    @Override
    public <T> T getEndpointInstance(Class<T> clazz)
            throws InstantiationException {
        try {
            return clazz.getConstructor().newInstance();
        } catch (InstantiationException e) {
            throw e;
        } catch (ReflectiveOperationException e) {
            InstantiationException ie = new InstantiationException();
            ie.initCause(e);
            throw ie;
        }
    }
}

使用@ServerEndpoint註解之後，無法自動注入Bean。每次創建一個新的連接之後，都是用反射創建一個對象，中間沒有從Sprin容器中找相應的Bean。
所以我們要麼自己獲取Bean，要麼將注入的Bean設置爲static，讓其注入到類上。
工具類獲取Spring Bean，只需實現ApplicationContextAware 接口即可。

@Component
public class SpringContextUtil implements ApplicationContextAware {
     private static ApplicationContext applicationContext;     //Spring應用上下文環境
     private static Properties properties=new Properties();
      /**
      * 實現ApplicationContextAware接口的回調方法，設置上下文環境   
      * @param applicationContext
      * @throws BeansException
      */
      public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
          SpringContextUtil.applicationContext = applicationContext;
      }
     
      /**
      * @return ApplicationContext
      */
      public static ApplicationContext getApplicationContext() {
        return applicationContext;
      }
      
      * 獲取對象   
      * @param name
      * @return Object 一個以所給名字註冊的bean的實例
      * @throws BeansException
      */
      public static Object getBean(String name) throws BeansException {
          if(applicationContext==null) {
              return null;
          }
        return applicationContext.getBean(name);
      }
}