Java 实现 websocket

websocket 实现类

package com.jm.websocket;

import cn.hutool.core.util.ObjectUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArraySet;


/**
 * 与页面交互的ws
 */
@Component
@Slf4j
@ServerEndpoint("/record/websocket/{userId}") //此注解相当于设置访问URL
public class RecordWebSocket {
    private Session session;

    final static CopyOnWriteArraySet<RecordWebSocket> aiWebSockets = new CopyOnWriteArraySet<>();

    final static Map<String, Session> aiSessionPool = new HashMap<String, Session>();

    final static List<String> list = new ArrayList<>();

    public List<String> getList() {
        return list;
    }

    private static ApplicationContext applicationContext;
    public static void setApplicationContext(ApplicationContext applicationContext) {
        RecordWebSocket.applicationContext = applicationContext;
    }

    @OnOpen
    public void onOpen(Session session, @PathParam(value = "userId") String userId, EndpointConfig config) {
        try {
            this.session = session;
            aiWebSockets.add(this);
            aiSessionPool.put(userId, session);


//            log.info("【websocket消息】有新的连接,总数为:" + aiWebSockets.size());
        } catch (Exception e) {
            log.error("WebSocket 连接异常: " + e.getMessage());
        }
    }

    @OnClose
    public void onClose() {
        try {
            aiWebSockets.remove(this);
//            log.info("【websocket消息】连接断开,总数为:" + aiWebSockets.size());
        } catch (Exception e) {
            log.error("WebSocket 连接异常: " + e.getMessage());
        }
    }

    /**
     * 找到该用户对应的连接,设置强提醒
     * @param map
     * @param value
     * @return
     */
    public static List<String> getKey(Map map, Object value) {
        List<String> keyList = new ArrayList<>();
        for (Object temp : map.keySet()) {
            String key = temp.toString();
            if (map.get(key).equals(value)) {
                keyList.add(key);
            }
        }
        return keyList;
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        try {
            JSONObject jsonArray = JSONUtil.parseObj(message);
            String o = (String)jsonArray.get("userId");
            if (ObjectUtil.isNotNull(o)){
//                System.out.println("接受到消息"+o+"ccc"+message);
                sendOneMessage(o,"heartbeat");
            }

        } catch (Exception e) {
            log.error("接受消息失败={}", e);
        }
    }

    private void reSendStrongReminderMessage() {

    }

    // 此为广播消息
    public void sendAllMessage(String message) {
        for (RecordWebSocket webSocket : aiWebSockets) {
            try {
                if (webSocket.session.isOpen()) {
                    webSocket.session.getAsyncRemote().sendText(message);
                }
            } catch (Exception e) {
                log.error(e.getMessage());
            }
        }
    }

    // 此为单点消息
    public void sendOneMessage(String userId, String message) {
        Session session = aiSessionPool.get(userId);
        if (session != null && session.isOpen()) {
            try {
//                log.info("【websocket消息】 单点消息:" + message);
                session.getAsyncRemote().sendText(message);
            } catch (Exception e) {
                log.error(e.getMessage());
            }
        }
    }

    // 此为单点消息(多人)
    public void sendMoreMessage(String[] userIds, String message) {
        for (String userId : userIds) {
            Session session = aiSessionPool.get(userId);
            if (session != null && session.isOpen()) {
                try {
//                    log.info("【websocket消息】 单点消息:" + message);
                    session.getAsyncRemote().sendText(message);
                } catch (Exception e) {
                    log.error(e.getMessage());
                }
            }
        }

    }


    public static void sendMsgToAllOnline(String msg) {
        if (aiSessionPool.isEmpty()) {
            return;
        }

        aiSessionPool.values().forEach(session -> {
            if (session != null && session.isOpen()) {
                try {
                    log.info("ws待推送向前端的消息:" + msg);
                    session.getAsyncRemote().sendText(msg);
                } catch (Exception e) {
                    log.error("ws向前端推送消息错误,消息内容={},错误原因={}", msg, e);
                }
            }
        });
    }

}

实现推送数据，

@Resource
private RecordWebSocket recordWebSocket;

public String acceptAiMsg(AcceptAuditResultDto dto) {
   recordWebSocket.sendOneMessage(dto.getWebSocketId(), dto.getAuditContent());
   return "ok";
}