Mediation-Backend/ruoyi-common/src/main/java/com/ruoyi/common/websocket/MyWebSocketHandler.java

package com.ruoyi.common.websocket;


import com.ruoyi.common.exception.ServiceException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @author wangqiong
 * @description
 * @date 2023-11-25 15:16
 */
@Component
@Slf4j
@ServerEndpoint("/websocket/{userName}")
public class MyWebSocketHandler {
    // 接口路径 ws://127.0.0.1:9000/websocket;
    private Session session;

    //concurrent包的线程安全Set，用来存放每个客户端对应的WebSocket对象。
    private static CopyOnWriteArraySet<Session> sessions  = new CopyOnWriteArraySet<>();
    // 用来存在线连接数
    private static Map<String, Session> sessionPool = new HashMap<>();

    /**
     * 链接成功调用的方法
     */
    @OnOpen
    public void onOpen(Session session, @PathParam(value = "userName") String userName) {
        try {
            sessions.add(session);
            sessionPool.put(userName, session);

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 链接关闭调用的方法
     */
    @OnClose
    public void onClose(Session session,@PathParam(value = "userName") String userName) {
        try {
            if(session!=null && session.isOpen()){
                session.close();
                sessions.remove(session);
                sessionPool.remove(userName);
            }

            log.info("【websocket消息】连接断开，总数为:" + sessions.size());

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 收到客户端消息后调用的方法
     *
     * @param message
     * @param
     */
    @OnMessage
    public void onMessage(@PathParam(value = "userName") String userName, String message) {

        System.out.println("【websocket消息】收到客户端消息:" + message);
        // 将消息广播给其它用户
        for (Session session : sessions) {
            if(session!=null && session.isOpen()){
                try {
                    session.getBasicRemote().sendText(message);
                }catch (Exception e){
                    throw new ServiceException(e.getMessage());
                }
            }else {
                try {
                   session.close();
                } catch (IOException e) {
                    throw new ServiceException(e.getMessage());
                }
               sessions.remove(session);
                sessionPool.remove(userName);
            }
        }

    }

    /**
     * 发送错误时的处理
     *
     * @param session
     * @param error
     */
    @OnError
    public void onError(Session session, Throwable error) {

        throw new ServiceException(error.getMessage());
    }


    /**
     * 推消息给前端
     *
     * @param userId
     * @param message
     * @return
     */
    public static Runnable sendOneMessage(String userId, String message) {
        Session session = sessionPool.get(userId);
        if (session != null && session.isOpen()) {
            try {
                log.info("【推给前端消息】 :" + message);

                //高并发下，防止session占用期间，被其他线程调用
                synchronized (session) {
                    session.getBasicRemote().sendText(message);
                }

            } catch (Exception e) {
                e.printStackTrace();
            }
        }
        return null;
    }





}