WebSocket整理

WebSocket整理

增加maven導入包

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>
<dependency>
    <groupId>org.apache.tomcat.embed</groupId>
    <artifactId>tomcat-embed-websocket</artifactId>
    <version>9.0.12</version>
</dependency>

增加開啓配置文件WebSocketConfig

當前文件使用於在idea開發工具中使用，若部署到tomcat中時，請將@Bean行註釋

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

/**
 * webSocket配置類
 *
 * @author zichen
 */
@Configuration
public class WebSocketConfig {

    /**
     * 注入ServerEndpointExporter，
     * 這個bean會自動註冊使用了@ServerEndpoint註解聲明的Websocket endpoint
     * @return ServerEndpointExporter對象
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter() {
        return new ServerEndpointExporter();
    }

}

具體實現

1. 服務接口實現調用
2. 可使用@OnMessage註解方法接收和返回的String字符串中將要傳輸的數據使用JSONString的方式接收和輸出

import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.zc.model.dto.ChatDTO;
import com.zc.model.po.chat.UserGroup;
import com.zc.service.UserGroupService;
import com.zc.service.UserService;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;


/**
 * 自定義webSocket配置
 *
 * @author zichen
 */
@Slf4j
@ServerEndpoint("/websocket/{userId}")
@Component
public class CustomWebSocket {

    /**
     * 與某個客戶端的連接會話，需要通過它來給客戶端發送數據
     */
    private Session session;

    /**
     * concurrent包的線程安全Set，用來存放每個客戶端對應的CumWebSocket對象。
     */
    private static CopyOnWriteArraySet<CustomWebSocket> webSockets = new CopyOnWriteArraySet<>();

    /**
     * 會話連接池
     */
    private static ConcurrentHashMap<Integer, Session> sessionPool = new ConcurrentHashMap<>();

    /**
     * 在工具類中注入Service方式
     * 1.需創建靜態的當前類
     * 2.使用@Autowired注入Service
     * 3.使用@PostConstruct註解init方法
     */
    private static CustomWebSocket customWebSocket;

    @Autowired
    private UserGroupService userGroupService;

    @Autowired
    private UserService userService;

    @PostConstruct
    public void init() {
        customWebSocket = this;
        customWebSocket.userService = userService;
        customWebSocket.userGroupService = userGroupService;
    }

    /**
     * 連接建立成功調用的方法
     *
     * @param session 會話信息
     * @param userId  用戶主鍵
     */
    @OnOpen
    public void onOpen(Session session, @PathParam(value = "userId") Integer userId) {
        this.session = session;
        if (StrUtil.isNotEmpty(openId)) {
            userId = customWebSocket.userService.pullId(openId);
        }
        webSockets.add(this);
        sessionPool.put(userId, session);
        session.getAsyncRemote().sendText("連接成功");
        log.info("【websocket消息】有新的連接，登錄用戶ID爲" + userId + "，連接總數爲:" + webSockets.size());
    }

    /**
     * 連接關閉調用的方法
     */
    @OnClose
    public void onClose() {
        webSockets.remove(this);
        log.info("【websocket消息】連接斷開，總數爲:" + webSockets.size());
    }

    /**
     * 收到客戶端消息後調用
     * @type single：單點發送，group：羣聊
     * @param message 消息信息
     */
    @OnMessage
    public void onMessage(String message) {
        log.info("【websocket消息】收到客戶端消息:" + message);
        ChatDTO chatDTO = JSON.parseObject(message, ChatDTO.class);
        System.err.println(chatDTO);
        switch (chatDTO.getType().toLowerCase()) {
            case "single":
                sendSingleMessage(Integer.parseInt(chatDTO.getUserId()), chatDTO.getMessage());
                break;
            case "group":
                sendGroupMessage(Integer.parseInt(chatDTO.getGroupId()), chatDTO.getMessage());
                break;
            default:
                sendAllMessage(chatDTO.getMessage());
                break;
        }
    }

    /**
     * 發生錯誤時調用
     *
     * @param session 會話信息
     * @param error   錯誤信息
     */
    @OnError
    public void onError(Session session, Throwable error) {
        log.error(session.getId() + "----websocket-------有異常");
        error.printStackTrace();
    }

    /**
     * 單點發送信息
     *
     * @param message 消息信息
     */
    public void sendSingleMessage(Integer userId, String message) {
        Session session = sessionPool.get(userId);
        if (session != null) {
            try {
                session.getAsyncRemote().sendText(message);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 羣消息
     * @param groupId
     * @param message
     */
    public void sendGroupMessage(Integer groupId, String message) {
        List<UserGroup> userGroupList = customWebSocket.userGroupService.pull(groupId);
        userGroupList.forEach(userGroup -> {
            Session session = sessionPool.get(userGroup.getUserId());
            if (session != null) {
                try {
                    session.getAsyncRemote().sendText(message);
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }
        });
    }

    /**
     * 廣播發送消息
     *
     * @param message 消息信息
     */
    public void sendAllMessage(String message) {
        webSockets.forEach(webSocket -> {
            try {
                webSocket.session.getAsyncRemote().sendText(message);
            } catch (Exception e) {
                e.printStackTrace();
            }
        });
    }

}