Spring Boot で WebSocket (STOMPではなく、TextWebSocketHandler)を試してみる

[programming] Java

こないだSTOMP over WebSocketを試してみましたが、今度はSTOMPを使わずに、TextWebSocketHandlerを使ったWebSocketを試してみます。題材も同じくチャットです。

ルームの情報をどのように渡そうか悩みました。コードを書き始めるまでは、WebSocketで接続する先にヘッダとかで渡せばいいかなと思っていましたが、調べてみたら渡せなかったので、やもえずURLのクエリパラメータとして渡すようにしました。

コード

テキストでのやり取りを行うので、TextWebSocketHandler を利用します。TextWebSocketHandlerの各メソッドをoverrideして必要な処理を実装するだけです。

ルームの情報は、URLのクエリとしてクライアントから送っているので、接続が確立したタイミング(afterConnectionEstablished)にて、ルーム毎にWebSocketSessionを保持するようにします。 メッセージを受け取ったら、自分のルームと同じWebSocketSessionに対して、メッセージを送るだけです。

package com.example;

import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

import org.springframework.stereotype.Component;
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;

@Component
public class ChatHandler extends TextWebSocketHandler {

    private ConcurrentHashMap<String, Set<WebSocketSession>> roomSessionPool = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {

        String roomName = session.getUri().getQuery();

        roomSessionPool.compute(roomName, (key, sessions) -> {

            if (sessions == null) {
                sessions = new CopyOnWriteArraySet<>();
            }
            sessions.add(session);

            return sessions;
        });
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {

        String roomName = session.getUri().getQuery();

        for (WebSocketSession roomSession : roomSessionPool.get(roomName)) {
            roomSession.sendMessage(message);
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {

        String roomName = session.getUri().getQuery();

        roomSessionPool.compute(roomName, (key, sessions) -> {

            sessions.remove(session);
            if (sessions.isEmpty()) {
                // 1件もない場合はMapからクリア
                sessions = null;
            }

            return sessions;
        });
    }
}

WebSocketの設定として、URLとHandlerを紐付けます。

package com.example;

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;

import lombok.AllArgsConstructor;

@Configuration
@EnableWebSocket
@AllArgsConstructor
public class WebSocketConfig implements WebSocketConfigurer {

    private final ChatHandler chatHandler;

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        registry.addHandler(chatHandler, "/endpoint");
    }

}

クライアント側では、下記のようなコードになりました。

<!DOCTYPE html>
<html lang="ja">
<head>
<meta charset="UTF-8" />
<title>チャット</title>
<link rel="stylesheet" href="/webjars/bootstrap/3.3.7/css/bootstrap.min.css" />
<link rel="stylesheet" href="/webjars/bootstrap/3.3.7/css/bootstrap-theme.min.css" />
</head>
<body>
  <div class="container">
    <h2>チャット</h2>
    <div class="form-horizontal">
      <div class="form-group">
        <label for="roomName" class="col-sm-2 control-label">ルーム</label>
        <div class="col-sm-2">
          <input id="roomName" type="text" class="form-control" value="example" />
        </div>
        <div class="col-sm-3">
          <button id="connectButton" type="button" class="btn btn-default">接続</button>
          <button id="disconnectButton" class="btn btn-default">切断</button>
        </div>
      </div>
      <div class="form-group">
        <label for="message" class="col-sm-2 control-label">メッセージ</label>
        <div class="col-sm-4">
          <input id="message" type="text" class="form-control" />
        </div>
        <div class="col-sm-2">
          <button id="sendButton" type="button" class="btn btn-default">送信</button>
        </div>
      </div>
      <div class="row">
        <div class="col-sm-4 col-sm-offset-2">
          <ul id="messageList" class="list-unstyled">
          </ul>
        </div>
      </div>
    </div>
  </div>
  <script src="/webjars/jquery/1.12.4/jquery.min.js"></script>
  <script src="/webjars/bootstrap/3.3.7/js/bootstrap.min.js"></script>
  <script>
    $(function() {
      var endpoint = 'ws://' + location.host + '/endpoint';
      var webSocket = null;

      $('#connectButton').click(function() {

        $("#messageList").empty();

        webSocket = new WebSocket(endpoint + '?' + encodeURIComponent($('#roomName').val()));
        webSocket.onopen = function() {
          $('#roomName').prop('disabled', true);
          $('#connectButton').prop('disabled', true);
          $('#disconnectButton').prop('disabled', false);
        };
        webSocket.onclose = function() {
        };
        webSocket.onmessage = function(message) {
          $('#messageList').prepend($('<li>').text(message.data));
        };
        webSocket.onerror = function() {
          alert('エラーが発生しました。');
        };
      });

      $('#disconnectButton').click(function() {

        webSocket.close();
        webSocket = null;

        $('#roomName').prop('disabled', false);
        $('#connectButton').prop('disabled', false);
        $('#disconnectButton').prop('disabled', true);
      });

      $('#sendButton').click(function() {
        if (!webSocket) {
          alert('未接続です。');
          return;
        }

        webSocket.send($('#message').val());
      });
    });
  </script>
</body>
</html>

STOMPの時のほうが、いろいろシンプルに書けるので、ブラウザがクライアントならば、STOMPを使わない理由はないかなと思っています。