消息推送 bug修改

src/main/java/com/ai/da/common/websocket/NotificationConnection.java

@@ -4,40 +4,92 @@ 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.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 
-@ServerEndpoint(value = "/notification")
+@ServerEndpoint(value = "/notification/{id}")
 @Component
 @Slf4j
 public class NotificationConnection {
 
     static Map<String,Session> sessionMap = new ConcurrentHashMap<>();
+
+    private Session session;
+    // 这里用ConcurrentHashMap 因为他是一个线程安全的Map
+    private static ConcurrentHashMap<Long, NotificationConnection> websockets = new ConcurrentHashMap<>();
     //连接建立时执行的操作
-    @OnOpen
+    /*@OnOpen
     public void onOpen(Session session){
         sessionMap.put(session.getId(),session);
-        log.info("websocket is open");
+        log.info("websocket is open, sessionId: {}",session.getId());
+    }*/
+
+    @OnOpen
+    public void onOpen(Session session, @PathParam("id") String id) {  // 接收到前端传来的用户ID
+        this.session = session;
+        websockets.put(Long.parseLong(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 text){
         log.info("收到了一条消息："+text);
 //        return "已收到你的消息";
     }
-    //连接关闭的时候执行的操作
+    /*//连接关闭的时候执行的操作
     @OnClose
     public void onClose(Session session){
         sessionMap.remove(session.getId());
-        log.info("websocket is close");
+        log.info("websocket is close, sessionId: {}",session.getId());
     }
 
     public void sendMsg(String message) throws IOException {
         for(String key:sessionMap.keySet()){
             sessionMap.get(key).getBasicRemote().sendText(message);
         }
+    }*/
+
+    public void sendMsg(String message, Long userId) throws IOException {
+        if (userId == null) {  // 如果等于null则证明是群发
+            // 获取当前Map的一个迭代器，遍历Map的方式有很多种，看着来
+            // 这个就是遍历这个集合的过程....
+            for (Map.Entry<Long, NotificationConnection> entry : websockets.entrySet()) {
+                // 获取每一个Entry实例
+                // 获取每一个Value，而这个Value就是WebSocket的实例
+                NotificationConnection webSocket = entry.getValue();
+                // 接下来就是遍历群发
+                log.info("广播消息 【给用户】 ：{}发送消息【{}】", webSocket, message);
+                webSocket.session.getBasicRemote().sendText(message); // 发送！！！！！！！！！
+            }
+        } else {  // 如果不是群发，则判断ID，其余步骤一致
+            // 获取当前Map的一个迭代器，遍历Map的方式有很多种，看着来
+            // 这个就是遍历这个集合的过程....
+            for (Map.Entry<Long, NotificationConnection> entry : websockets.entrySet()) {
+                // 获取每一个Entry实例
+                // 获取每一个Value，而这个Value就是WebSocket的实例
+                NotificationConnection webSocket = entry.getValue();
+                // 获取每一个Key，这个Key就是用户ID
+                Long key = entry.getKey();
+                // 判断用户ID与当前的Key相等
+                if (userId.equals(key)) {
+                    log.info("私发消息 【给用户】 ：{}发送消息【{}】", key, message); // 打印
+                    webSocket.session.getBasicRemote().sendText(message);  // 则发送给当前的用户即可
+                }
+            }
+        }
     }
 }