bms-java/src/main/java/com/zy/bms/websocket/WebSocketServer.java

package com.zy.bms.websocket;

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.LinkedList;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/webSocket/{deviceId}")
public class WebSocketServer {

    /**
     * 存放每个客户端对应的MyWebSocket对象。
     * key : 设备ID
     * value : 实时接收该设备的web端session
     */
    private static ConcurrentHashMap<String, List<Session>> webSocketMap = new ConcurrentHashMap<>();

    /**
     * 连接成功
     */
    @OnOpen
    public void onOpen(Session session, @PathParam("deviceId") String deviceId) {
        //如果该设备没有web端查看,则初始化
        if (!webSocketMap.containsKey(deviceId)) webSocketMap.put(deviceId, new LinkedList<>());
        //添加到集合中
        webSocketMap.get(deviceId).add(session);
    }

    /**
     * 连接关闭
     */
    @OnClose
    public void onClose(Session session, @PathParam("deviceId") String deviceId) {
        webSocketMap.get(deviceId).remove(session);
    }

    /**
     * 连接错误
     */
    @OnError
    public void onError(Session session, Throwable error, @PathParam("deviceId") String deviceId) {
        webSocketMap.get(deviceId).remove(session);
    }

    /**
     * 群发消息
     */
    public void massMessage(String deviceId, String message) {
        List<Session> sessions = webSocketMap.get(deviceId);
        if (sessions == null || sessions.isEmpty()) return;
        // 遍历客户端 发送消息
        for (Session session : sessions) {
            sendMessage(session, message);
        }
    }

    /**
     * 实现服务器主动推送
     */
    private void sendMessage(Session session, String message) {
        try {
            session.getBasicRemote().sendText(message);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}