两种SpringBoot整合Websocket的方式

Websocket协议已经出来很久了，所以网上的教程也很多，我这里大概就是整理了两种比较好用的，大家可以根据自己的业务场景去选择性的使用

首先是引入需要的依赖：

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

第一种我们使用注解的形式：

配置文件

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

然后就是核心代码了：

@Component
@ServerEndpoint("/websocket/message/{token}")
public class WebSocketServer
{
    /**
     * WebSocketServer 日志控制器
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketServer.class);


    /**
     * 默认最多允许同时在线人数100
     */
    public static int socketMaxOnlineCount = 100;

    private static Semaphore socketSemaphore = new Semaphore(socketMaxOnlineCount);

    /**
     * 连接建立成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session,@PathParam("token") String token) throws Exception
    {
        boolean semaphoreFlag = false;
        //身份验证
        if(!StringUtils.isNotNull(token)){
            session.close();
            return;
        }

        TokenService tokenService = SpringUtils.getBean(TokenService.class);

        LoginUser user = tokenService.getUserByToken(token);
        if(!StringUtils.isNotNull(user)){
            session.close();
            return;
        }

        // 尝试获取信号量
        semaphoreFlag = SemaphoreUtils.tryAcquire(socketSemaphore);
        if (!semaphoreFlag)
        {
            // 未获取到信号量
            LOGGER.error("\n 当前在线人数超过限制数- {}", socketMaxOnlineCount);
            WebSocketUsers.sendMessageToUserByText(session, "当前在线人数超过限制数：" + socketMaxOnlineCount);
            session.close();
        }
        else
        {
            // 添加用户
            WebSocketUsers.put(user.getUsername(), session);
            LOGGER.info("\n 建立连接 - {}", session);
            LOGGER.info("\n 当前人数 - {}", WebSocketUsers.getUsers().size());
            WebSocketUsers.sendMessageToUserByText(session, "连接成功");
        }
    }

    /**
     * 连接关闭时处理
     */
    @OnClose
    public void onClose(Session session)
    {
        LOGGER.info("\n 关闭连接 - {}", session);
        // 移除用户
        WebSocketUsers.remove(session);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }

    /**
     * 抛出异常时处理
     */
    @OnError
    public void onError(Session session, Throwable exception) throws Exception
    {
        if (session.isOpen())
        {
            // 关闭连接
            session.close();
        }
        String sessionId = session.getId();
        LOGGER.info("\n 连接异常 - {}", sessionId);
        LOGGER.info("\n 异常信息 - {}", exception);
        // 移出用户
        WebSocketUsers.remove(session);
        // 获取到信号量则需释放
        SemaphoreUtils.release(socketSemaphore);
    }

    /**
     * 服务器接收到客户端消息时调用的方法
     */
    @OnMessage
    public void onMessage(String message, Session session)
    {
        if(!UserConstants.WEBSOCKET_HEARTBEAT.equals(message)){
            try{
                SysMessage msg = JSON.parseObject(message, new TypeReference<SysMessage>(){});
                if(StringUtils.isNotNull(msg.getRecipientName())){
                    //这里必须传递username
                    WebSocketUsers.sendMesssageToUserByName(msg.getRecipientName(),message);
                }
            }catch (Exception e){
                LOGGER.error("\n 错误的websocket信息格式 - {}", message);
            }
        }
        LOGGER.debug("\n 收到客户端发送的消息 - {}", message);
    }

}

public class WebSocketUsers
{
    /**
     * WebSocketUsers 日志控制器
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(WebSocketUsers.class);

    /**
     * 用户集
     */
    private static Map<String, Session> USERS = new ConcurrentHashMap<String, Session>();

    /**
     * 存储用户
     *
     * @param key 唯一键
     * @param session 用户信息
     */
    public static void put(String key, Session session)
    {
        USERS.put(key, session);
    }

    /**
     * 移除用户
     *
     * @param session 用户信息
     *
     * @return 移除结果
     */
    public static boolean remove(Session session)
    {
        String key = null;
        boolean flag = USERS.containsValue(session);
        if (flag)
        {
            Set<Map.Entry<String, Session>> entries = USERS.entrySet();
            for (Map.Entry<String, Session> entry : entries)
            {
                Session value = entry.getValue();
                if (value.equals(session))
                {
                    key = entry.getKey();
                    break;
                }
            }
        }
        else
        {
            return true;
        }
        return remove(key);
    }

    /**
     * 移出用户
     *
     * @param key 键
     */
    public static boolean remove(String key)
    {
        LOGGER.info("\n 正在移出用户 - {}", key);
        Session remove = USERS.remove(key);
        if (remove != null)
        {
            boolean containsValue = USERS.containsValue(remove);
            LOGGER.info("\n 移出结果 - {}", containsValue ? "失败" : "成功");
            return containsValue;
        }
        else
        {
            return true;
        }
    }

    /**
     * 获取在线用户列表
     *
     * @return 返回用户集合
     */
    public static Map<String, Session> getUsers()
    {
        return USERS;
    }

    /**
     * 群发消息文本消息
     *
     * @param message 消息内容
     */
    public static void sendMessageToUsersByText(String message)
    {
        Collection<Session> values = USERS.values();
        for (Session value : values)
        {
            sendMessageToUserByText(value, message);
        }
    }

    /**
     * 发送文本消息
     *
     * @param session 自己的用户名
     * @param message 消息内容
     */
    public static void sendMessageToUserByText(Session session, String message)
    {
        if (session != null)
        {
            try
            {
                session.getBasicRemote().sendText(message);
            }
            catch (IOException e)
            {
                LOGGER.error("\n[发送消息异常]", e);
            }
        }
        else
        {
            LOGGER.info("\n[你已离线]");
        }
    }

    public static void sendMesssageToUserByName(String username,String message){
        Session session = USERS.get(username);
        if(StringUtils.isNotNull(session)){
            try
            {
                session.getBasicRemote().sendText(message);
            }
            catch (IOException e)
            {
                LOGGER.error("\n[发送消息异常]", e);
            }
        }
    }
}

使用到的工具类：

/**
 * 信号量相关处理
 * 
 * @author ruoyi
 */
public class SemaphoreUtils
{
    /**
     * SemaphoreUtils 日志控制器
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(SemaphoreUtils.class);

    /**
     * 获取信号量
     * 
     * @param semaphore
     * @return
     */
    public static boolean tryAcquire(Semaphore semaphore)
    {
        boolean flag = false;

        try
        {
            flag = semaphore.tryAcquire();
        }
        catch (Exception e)
        {
            LOGGER.error("获取信号量异常", e);
        }

        return flag;
    }

    /**
     * 释放信号量
     * 
     * @param semaphore
     */
    public static void release(Semaphore semaphore)
    {

        try
        {
            semaphore.release();
        }
        catch (Exception e)
        {
            LOGGER.error("释放信号量异常", e);
        }
    }
}

这里使用了，Semaphore来限制访问特定资源的并发线程数。

Semaphore（信号量）是一个线程同步结构，它通过计数器来控制对共享资源的访问。Semaphore还可以在线程之间发送信号，以避免丢失信号。Semaphore是一种变量，用于管理并发进程并同步它们

就好像自己创建一个变量 int = 0 ，然后等到int =100 的时候就不允许有人访问了，Semaphore的功能更强悍一些。

第二种就是

实现WebSocketConfigurer接口。通过重写registerWebSocketHandlers方法来注册自定义的WebSocketHandler，并指定类对应的WebSocket访问的ServerEndpoint。

配置类

/**
 * websocket配置信息，配置连接地址
 *
 */
@EnableWebSocket
@Configuration
public class WebSocketConfig implements WebSocketConfigurer {

    @Resource
    private MyWebsocketServer myWebsocketServer;

    /**
     * Register {@link WebSocketHandler WebSocketHandlers} including SockJS fallback options if desired.
     *
     * @param registry
     */
    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {

        WebSocketHandlerRegistration webSocketHandlerRegistration = registry.addHandler(myWebsocketServer, "/ws/{visionId}/{token}").setAllowedOrigins("*");

    }
}

核心代码

@Slf4j
@Component
@RequiredArgsConstructor
public class MyWebsocketServer implements WebSocketHandler {

    protected static final CopyOnWriteArrayList<WebSocketSession> WEB_SOCKET_SESSIONS = new CopyOnWriteArrayList<>();


    @Autowired
    private TokenService tokenService;


    /**
     * 建立连接后操作
     *
     * @param session 连接session信息
     * @throws Exception exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
          URI uri = session.getUri();
          String[] split = String.valueOf(uri.getPath()).split("/");
          String token= split[3];
         if(!StringUtils.isNotNull(token)){
            session.close();
            return;
           }

        LoginUser user = tokenService.getUserByToken(token);
        if(!StringUtils.isNotNull(user)){
            session.close();
            return;
         }
         WEB_SOCKET_SESSIONS.add(session);
    }

    /**
     * 接收到消息后的处理
     * @param session 连接session信息
     * @param message 信息
     * @throws Exception exception
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        Map<String,Object> resultMap = new HashMap<>();
        resultMap.put("msgType","100");
        sendMessage(resultMap);
    }

    /**
     * ws连接出错时调用
     *
     * @param session   session连接信息
     * @param exception exception
     * @throws Exception exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        if (session.isOpen()) {
            Map<String,Object> resultMap = new HashMap<>();
            resultMap.put("error", "ws连接出错，即将关闭此session,sessionId=【" + session.getId() + "】");
            sendMessage(resultMap);
            session.close();
        }
        WEB_SOCKET_SESSIONS.remove(session);
    }

    /**
     * 连接关闭后调用
     *
     * @param session     session连接信息
     * @param closeStatus 关闭状态
     * @throws Exception exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        if (session.isOpen()) {
            Map<String,Object> resultMap = new HashMap<>();
            resultMap.put("error", "ws连接出错，即将关闭此session,sessionId=【" + session.getId() + "】");
            sendMessage(resultMap);
            session.close();
        }
        WEB_SOCKET_SESSIONS.remove(session);
    }

    /**
     * 是否支持分片消息
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    /**
     * 群发发送消息
     *
     * @param message 消息
     * @throws IOException ioException
     */
    public void sendMessage(Map<String,Object> message) throws IOException {
        if (CollUtil.isNotEmpty(WEB_SOCKET_SESSIONS)) {
            for (WebSocketSession webSocketSession : WEB_SOCKET_SESSIONS) {
                ObjectMapper objectMapper = new ObjectMapper();
                AjaxResult success = AjaxResult.success(message);
                String resultString = objectMapper.writeValueAsString(success);
                webSocketSession.sendMessage(new TextMessage(resultString));
            }
        }
    }

    /**
     * 发给指定连接消息
     *
     * @param message 消息
     * @throws IOException ioException
     */
    public void sendMessage(WebSocketMsg message, String sessionId) throws IOException {
        if (CollUtil.isNotEmpty(WEB_SOCKET_SESSIONS)) {
            for (WebSocketSession webSocketSession : WEB_SOCKET_SESSIONS) {
                if (sessionId.equals(webSocketSession.getId())) {
                    ObjectMapper objectMapper = new ObjectMapper();
                    AjaxResult success = AjaxResult.success(message);
                    String resultJson = objectMapper.writeValueAsString(success);
                    webSocketSession.sendMessage(new TextMessage(resultJson));
                }
            }
        }
    }
}