springboot中使用websocket

1. 引入maven

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

2.配置类

@EnableWebSocket
@Configuration
public class WebSocketConfig {

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

3.消息功能模块

由于在该类中无法使用注解引入spring类，所以使用spring的工具类引入bean.

@Slf4j
@Component
@ServerEndpoint("/wbs/{userId}")
public class ChatSocketMessage {

    /**
     * 静态变量，用来记录当前在线连接数。应该把它设计成线程安全的。
     */
    private static AtomicInteger onlineCount = new AtomicInteger(0);
    /**
     * concurrent包的线程安全Set，用来存放每个客户端对应的WebSocket对象。
     */
    private static ConcurrentHashMap<Integer, Session> webSocketMap = new ConcurrentHashMap<>();

    @OnMessage
    public void handleMessage(@PathParam("userId") Integer userId, String msg) {
        log.info("socket sendMsg ,userId={},msg={}", userId, msg);
        try {
            ChatMessageVo message = JSON.parseObject(msg, ChatMessageVo.class);
            SpringContextHolder.getBean(ChatMessageService.class).saveMessage(message);
        } catch (Exception e) {
            log.error("socket sendMsg error ,userId={},msg={},ex:{}", userId, msg, e);
        }
    }

    @OnOpen
    public void onOpen(Session session, @PathParam("userId") Integer userId) {
        webSocketMap.put(userId, session);
        onlineCount.addAndGet(1);
        log.info("socket create,userId={}", userId);
    }

    @OnClose
    public void onClose(Session session, @PathParam("userId") Integer userId) {
        webSocketMap.remove(userId);
        log.info("socket close,userId={}", userId);
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("WebSocket发生错误,{}", error);
    }

    public void sendMessage(Integer userId, Integer contactId, String message) {
        try {
            Session session = this.webSocketMap.get(userId);
            if (Objects.nonNull(session)) {
                session.getBasicRemote().sendText(message);
            } else {
                SpringContextHolder.getBean(ChatUserService.class).updateUnread(userId, contactId);
            }
        } catch (IOException e) {
            log.error("sendMessage ex:{}", e);
        }
    }

}

spring 工具类：

@Component
public class SpringContextHolder implements ApplicationContextAware {

    private static ApplicationContext applicationContext;

    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    public static ApplicationContext getApplicationContext() {
        return applicationContext;
    }

    /**
     * 通过name获取 Bean.
     *
     * @param name
     * @return
     */
    public static Object getBean(String name) {
        return getApplicationContext().getBean(name);
    }

    /**
     * 通过class获取Bean.
     *
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> T getBean(Class<T> clazz) {
        return getApplicationContext().getBean(clazz);
    }

    /**
     * 通过name,以及Clazz返回指定的Bean
     *
     * @param name
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> T getBean(String name, Class<T> clazz) {
        return getApplicationContext().getBean(name, clazz);
    }
}

4. 前端代码

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

    // 判断当前浏览器是否支持WebSocket
    if('WebSocket' in window){
        // 为了方便测试，故将链接写死
        websocket = new WebSocket("ws://localhost:8088/wbs/1");
    } else{
        alert('Not support websocket')
    }

    // 连接发生错误的回调方法
    websocket.onerror = function(){
        console.log("发生错误");
    };

    // 连接成功建立的回调方法
    websocket.onopen = function(event){
        console.log("打开连接");
    }

    // 接收到消息的回调方法
    websocket.onmessage = function(event){
        console.log(event.data)
    }

    // 连接关闭的回调方法
    websocket.onclose = function(){
        console.log("关闭连接")
    }

    // 监听窗口关闭事件，当窗口关闭时，主动去关闭websocket连接，防止连接还没断开就关闭窗口，server端会抛异常
    window.onbeforeunload = function(){
        websocket.close();
    }

    // 关闭连接
    function closeWebSocket(){
        websocket.close();
    }

    // 发送消息
    function send(){
        websocket.send(message);
    }
</script>