spring boot 支持websocket

通信流程

pom依赖

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

websocket 配置

@Configuration
public class WebsocketConfig {

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

endpoint接收session

@ServerEndpoint(value = "/websocket/{category)/userName}")
@Component
@Scope
public class WebsocketEndpoint {

    private static final Logger LOGGER = LoggerFactory.getLogger(WebsocketEndpoint.class);
    private Session session;

    private String user;

    private String category;

    @OnOpen
    public void onOpen(@PathParam("category") String category, @PathParam("userName") String userName,
            Session session) {
        this.session = session;
        this.category = category;
        this.user = userName;
        WebsocketHandler.put(this);
        LOGGER.info("WebSocketopen category:{), user:{", category, userName);
    }

    @OnClose
    public void onClose() {
        WebsocketHandler.remove(this.category, this.user);
        LOGGER.info("WebSocketclose category.{), user:{", category, user);
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        LOGGER.debug("on message:", message);
    }

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

    public void sendMessage(String message) throws IOException {
        LOGGER.debug("Send message:{", message);
        this.session.getBasicRemote().sendText(message);
    }

    public void sendobject(Object message) throws IOException {
        ObjectMapper objectMapper = new ObjectMapper();
        String jsonString = objectMapper.writeValueAsString(message);
        LOGGER.debug("Send message:{", jsonString);
        this.session.getBasicRemote().sendText(jsonString);
    }

    public Session getSession() {
        return session;
    }

    public void setSession(Session session) {
        this.session = session;
    }

    public String getUser() {
        return user;
    }

    public void setUser(String user) {
        this.user = user;
    }

    public String getCategory() {
        return category;
    }

    public void setCategory(String category) {
        this.category = category;
    }
    
}

endpoint session管理中心

public class WebsocketHandler {
    private static final Logger LOGGER = LoggerFactory.getLogger(WebsocketHandler.class);

    private static ConcurrentHashMap<String, Map<String, WebsocketEndpoint>> websocketServers;

    private static ConcurrentHashMap<String, Future<String>> workerFutures;

    private static ConcurrentHashMap<String, WebsocketWorker> workers;

    static {
        websocketServers = new ConcurrentHashMap<String, Map<String, WebsocketEndpoint>>();
        workerFutures = new ConcurrentHashMap<String, Future<String>>();
        workers = new ConcurrentHashMap<String, WebsocketWorker>();
    }

    public static WebsocketEndpoint getWebSocketServer(String category, String user) {
        Map<String, WebsocketEndpoint> userWebsocketMap = websocketServers.get(category);
        if (userWebsocketMap != null) {

            return userWebsocketMap.get(user);
        } else {
            return null;

        }
    }

    public static void put(WebsocketEndpoint webSocketServer) {
        LOGGER.info("add websocket category:{}, user:{}", webSocketServer.getCategory(), webSocketServer.getUser());
        synchronized (websocketServers) {
            if (!websocketServers.contains(webSocketServer.getCategory())) {
                websocketServers.put(webSocketServer.getCategory(), new ConcurrentHashMap<String, WebsocketEndpoint>());
            }
            websocketServers.get(webSocketServer.getCategory()).put(webSocketServer.getUser(), webSocketServer);
            if (!workers.contains(webSocketServer.getCategory())) {
                WebsocketWorker webSocketStatsWorker = new WebsocketWorker(webSocketServer.getCategory());
                workers.put(webSocketServer.getCategory(), webSocketStatsWorker);

                workerFutures.put(webSocketServer.getCategory(),
                        ThreadPoolFactory.getPool().submit(webSocketStatsWorker));
            }
        }
    }

    public static void remove(String category, String user) {
        LOGGER.info("remove websocket category:{}, user:{}", category, user);
        synchronized (websocketServers) {
            if (websocketServers.contains(category)) {
                Map<String, WebsocketEndpoint> userWebSocketMap = websocketServers.get(category);
                userWebSocketMap.remove(user);
                if (userWebSocketMap.isEmpty()) {
                    WebsocketWorker webSocketStatsWorker = workers.remove(category);
                    if (webSocketStatsWorker != null) {
                        webSocketStatsWorker.setClose(true);
                    }
                    Future<String> future = workerFutures.remove(category);
                    if (future != null) {
                        future.cancel(true);
                    }
                }
            }
        }
    }

    public static void sendMessage(String category, String user, String message) throws IOException {
        WebsocketEndpoint webSocketServer = getWebSocketServer(category, user);
        if (webSocketServer != null) {
            webSocketServer.sendMessage(message);
        }
    }

    public static void sendObject(String category, String user, Object message) throws IOException {
        WebsocketEndpoint webSocketServer = getWebSocketServer(category, user);
        if (webSocketServer != null) {

            webSocketServer.sendobject(message);
        }
    }

    public static void sendObject(String category, Object message) throws IOException {
        Map<String, WebsocketEndpoint> userWebsocketMap = websocketServers.get(category);
        if (userWebsocketMap != null) {
            for (WebsocketEndpoint userWebSocket : userWebsocketMap.values()) {
                LOGGER.info("send message to category:{}, user:{", userWebSocket.getCategory(),
                        userWebSocket.getUser());
                try {
                    userWebSocket.sendobject(message);
                } catch (Exception e) {
                    LOGGER.error("Failed to send message to category:{}, user: {l", userWebSocket.getCategory(),
                            userWebSocket.getUser());
                }
            }
        }
    }
}

定时发送消息线程

public class WebsocketWorker implements Callable<String> {
    private static final Logger LOG = LoggerFactory.getLogger(WebsocketWorker.class);

    private String category;

    private boolean close = false;

    public WebsocketWorker(String category) {
        this.category = category;
    }

    @Override
    public String call() throws Exception {
        while (!close) {
            WebsocketHandler.sendObject(category, "new message");
            try {
                Thread.sleep(10000);
            } catch (Exception e) {
                LOG.error("", e);
            }
        }
        return "close";
    }

    public boolean isClose() {
        return close;
    }

    public void setClose(boolean close) {
        this.close = close;
    }

}