websocket实现单发，群发

package com.yzd.server.websocket;

import com.yzd.server.dao.mapper.EventMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * @Author: liury
 * @Description: webSocket
 * @DateTime: 2022/6/19 20:54
 **/
@Slf4j
@Component
@ServerEndpoint("/webSocket/{id}")
public class WebSocket {
    private static int onlineCount = 0;

    @Resource
    private EventMapper eventMapper;
    // 这里用ConcurrentHashMap 因为他是一个线程安全的Map
    private static ConcurrentHashMap<Long, WebSocket> websockets = new ConcurrentHashMap<>();

    private Session session;


    @OnOpen
    public void onOpen(@PathParam("id") Long id, Session session) {  // 接收到前端传来的用户ID
        this.session = session;
        websockets.put(id, this);  //将ID作为key，当前的对象作为Value
        log.info("【建立连接】 用户为：" + this.session);
        log.info("【建立连接】 用户Id为：" + id);
        log.info("【建立连接】 总数为：" + websockets.size());
    }


    /**
     * 有用户连接断开时候触发该方法
     */
    @OnClose
    public void onClose() {
        websockets.remove(this); // 将当前的对象从集合中删除
        log.info("【连接断开】 用户为：" + this.session);
        log.info("【连接断开】 总数为：" + websockets.size());
    }


    @OnMessage

    public void onMessage(String message) throws IOException {
        if ("ping".equals(message)) {
            log.info("--------------------------webSocket心跳检测：" + message);
        } else {
            log.info("【收到客户端发的消息】：" + message);
        }


    }


    @OnError

    public void onError(Session session, Throwable error) {
        error.printStackTrace();

    }


    public static void sendMessageAll(String message) throws IOException {

//        if (sessions.size() != 0) {
//            for (Session s : sessions) {
//                if (s != null) {
//                    s.getAsyncRemote().sendText(message);
//                }
//            }
//        }


    }


    /**
     * 发送消息方法  【为了方便大家理解，我这里直接不封装了】
     *
     * @param message 消息
     * @param userId  用户ID
     */

    public void sendMessage(String message, Long userId) {
        if (userId == null) {  // 如果等于null则证明是群发
            // 获取当前Map的一个迭代器，遍历Map的方式有很多种，看着来
            Iterator<Map.Entry<Long, WebSocket>> iterator = websockets.entrySet().iterator();
            List<Long> userIds = eventMapper.selectUserId();
            // 这个就是遍历这个集合的过程....
            while (iterator.hasNext()) {
                // 获取每一个Entry实例
                Map.Entry<Long, WebSocket> entry = iterator.next();
                // 获取每一个Value，而这个Value就是WebSocket的实例
                WebSocket webSocket = entry.getValue();
                // 获取每一个Key，这个Key就是用户ID
                Long key = entry.getKey();
                // 接下来就是遍历群发

                for (Long id : userIds) {
                    if (id.equals(key)) {
                        try {
                            if (webSocket.session.isOpen()) {
                                log.info("广播消息 【给用户】 ：" + webSocket + "发送消息" + "【" + message + "】");
                                webSocket.session.getBasicRemote().sendText(message); // 发送！！！！！！！！！
                            }
                        } catch (IOException e) {
                            throw new RuntimeException(e);
                        }
                    }
                }
            }
        } else {  // 如果不是群发，则判断ID，其余步骤一致
            // 获取当前Map的一个迭代器，遍历Map的方式有很多种，看着来
            Iterator<Map.Entry<Long, WebSocket>> iterator2 = websockets.entrySet().iterator();
            // 这个就是遍历这个集合的过程....
            while (iterator2.hasNext()) {
                // 获取每一个Entry实例
                Map.Entry<Long, WebSocket> entry = iterator2.next();
                // 获取每一个Value，而这个Value就是WebSocket的实例
                WebSocket webSocket = entry.getValue();
                // 获取每一个Key，这个Key就是用户ID
                Long key = entry.getKey();
                // 判断用户ID与当前的Key相等
                if (userId.equals(key)) {

                    try {
                        if (webSocket.session.isOpen()) {
                            log.info("广播消息 【给用户】 ：" + key + "发送消息" + "【" + message + "】"); // 打印
                            webSocket.session.getBasicRemote().sendText(message);  // 则发送给当前的用户即可
                        }

                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                }

            }
        }
    }


    public static synchronized void addOnlineCount() {

        WebSocket.onlineCount++;

    }


    public static synchronized void subOnlineCount() {

        WebSocket.onlineCount--;

    }


}