如何编写websocket

小白一个，就是不会

已于 2024-05-11 16:14:14 修改

阅读量213

点赞数 4

文章标签： websocket 网络协议网络

于 2024-05-11 16:09:28 首次发布

本文链接：https://blog.csdn.net/A789AAA/article/details/138719805

版权

实体类

package com.fsmer.infra.project.infra.ws.vo;

import lombok.Data;

/**
 * @date 2023-12-27 15:54
 */
@Data
public class  CommWebSocketVO {
    /**
     * 1 监听心跳  2消息发送
     */
    private String msgType;
    /**
     * 1.试试定位 2预警
     */
    private String  businessType;
    /**
     * 发送信息内容
     */
    private String msgContent;

    /**
     * 发送信息内容
     */
    private String warnContent;

    private CommWebSocketVO(){}

    private CommWebSocketVO(String msgType,String businessType,String msgContent){
        this.msgType=msgType;
        this.businessType=businessType;
        this.msgContent=msgContent;
    }
    public CommWebSocketVO(String msgType, String msgContent){
        this.msgType=msgType;
        this.msgContent=msgContent;
    }
}

ws

package com.fsmer.infra.project.infra.ws;

import cn.hutool.core.util.ObjUtil;
import cn.hutool.core.util.StrUtil;
import com.alibaba.fastjson.JSON;
import com.fsmer.feign.system.api.user.AdminUserApi;
import com.fsmer.feign.system.api.user.dto.AdminUserRespDTO;
import com.fsmer.framework.common.pojo.CommonResult;
import com.fsmer.infra.project.infra.ws.vo.CommRedisWebSocketVO;
import com.fsmer.infra.project.infra.ws.vo.CommWebSocketVO;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnError;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 消息推送
 * @author Tracy
 */
@ServerEndpoint("/infra/ws/messageWebsocket/{accessToken}")
@Component
public class MessageWebsocket {
    //改bean是获取用户信息
    private static AdminUserApi adminUserApi;
    private Session session;
    private static final Logger log = LoggerFactory.getLogger(MessageWebsocket.class);
    private static int onlineCount = 0;
    private static ConcurrentHashMap<String, MessageWebsocket> webSocketMap = new ConcurrentHashMap<>();
    private static ConcurrentHashMap<String, List<Session>> sessionMap = new ConcurrentHashMap<>();
    private String accessToken = "";
    private String userId;
    private static ApplicationContext applicationContext;
    @Autowired
    public void setInjectBean(AdminUserApi adminUserApi) {
        MessageWebsocket.adminUserApi = adminUserApi;
    }

    @OnOpen
    public void onOpen(Session session, @PathParam("accessToken") String accessToken) throws IOException {
        this.session = session;
        this.accessToken = accessToken;
        AdminUserRespDTO userInfo = new AdminUserRespDTO();
        CommonResult<AdminUserRespDTO> commonResult = adminUserApi.getToken(accessToken);
        if (commonResult.getData() != null) {
            userInfo = commonResult.getData();
        }
            if (userInfo == null) {
                log.info("accessToken" + accessToken + "无效");
                String message=JSON.toJSONString(  new CommWebSocketVO("1","用户token无效或无匹配用户，请检查用户信息是否有效，链接自动关闭"));
                //无效时返回消息给连接方
                sendMessage(message);
                log.info("返回消息:" + userId + ",报文:" + message);
                this.onClose();
            }
        if (userInfo != null) {
            // 使用userId+sessionId作为key
            userId = userInfo.getId() + "_" + session.getId();
            if (StringUtils.isNotBlank(userId)){
                if (webSocketMap.containsKey(userId)) {
                    webSocketMap.remove(userId);
                    webSocketMap.put(userId, this);
                } else {
                    webSocketMap.put(userId, this);
                    addOnlineCount();
                }
                if(sessionMap.containsKey(userId)){
                    sessionMap.get(userId).add(session);
                }else{
                   sessionMap.put(userId, new ArrayList<Session>(){{ add(session); }});
                }
            }
        }
        log.info("用户连接:" + userId + ",当前在线人数为:" + getOnlineCount());
    }


    public static synchronized int getOnlineCount() {
        return onlineCount;
    }

    public static synchronized void addOnlineCount() {
        MessageWebsocket.onlineCount++;
    }

    public static synchronized void subOnlineCount() {
        MessageWebsocket.onlineCount--;
    }
    public static  void setApplicationContext(ApplicationContext tx){
        MessageWebsocket.applicationContext=tx;
    }


    @OnClose
    public void onClose() {
        if (StringUtils.isNotBlank(userId)){
            if (webSocketMap.containsKey(userId)) {
                webSocketMap.remove(userId);
                subOnlineCount();
            }
        }

        log.info("用户退出:" + userId + ",当前在线人数为:" + getOnlineCount());
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.error("用户错误:" + this.accessToken + ",原因:" + error.getMessage());
        error.printStackTrace();
    }

    @OnMessage
    public void onMessage( String message, Session session) throws IOException {
        if (StringUtils.isEmpty(message)){
            return;
        }
        //返回主叫方实体
        log.info("返回消息:" + userId + ",报文:" + JSON.toJSONString(message));
        //返回消息给发送方
//        CommWebSocketVO msg= JSON.parseObject(message,CommWebSocketVO.class);
//        sendMessage(msg,this.userId);
        session.getAsyncRemote().sendText(message);
    }

    /**
     * 发送到某个客户端
     * @param message 消息内容
     * @param userId
     * @throws IOException
     */
    public static void sendMessage(CommWebSocketVO message, String userId)  {
        try{
            String sendMsg= JSON.toJSONString(message);
            log.info("发送消息到:" + userId + "，报文:" + JSON.toJSONString(message));
            // 发送给该用户对应的所有客户端
            ConcurrentHashMap<String, MessageWebsocket> newMap = new ConcurrentHashMap<>();
            for (Map.Entry<String, MessageWebsocket> item : webSocketMap.entrySet()){
                // 判断item的key是否包含该用户，如果包含，存入newMap
                String pattern = userId + "_";
                if (StringUtils.isNotEmpty(item.getKey()) && item.getKey().contains(pattern)){
                    newMap.put(item.getKey(),item.getValue());
                }
            }
            // newMap中有值表示该用户在客户端有登录
            if (StringUtils.isNotBlank(userId) && !newMap.isEmpty()) {
                // 遍历newMap发送到对应客户端
                for (Map.Entry<String, MessageWebsocket> item : newMap.entrySet()) {
                    try {
                        item.getValue().sendMessage(sendMsg);
                    } catch (IOException e) {
                        continue;
                    }
                }
                // webSocketMap.get(userId).sendMessage(sendMsg);
            } else {
                log.error("用户" + userId + ",不在线！");
            }
        }catch (Exception e){
              e.printStackTrace();
        }
    }


    /**
     * 群发消息
     *
     * @param message
     * @throws IOException
     */
    public static void sendMesage(List<String> users, CommWebSocketVO message) {
        for (Map.Entry<String, MessageWebsocket> item : webSocketMap.entrySet()) {
            try {
               String sendMsg= JSON.toJSONString(message);
                if (users.contains(item.getKey())) {
                    item.getValue().sendMessage(sendMsg);
                    log.info("cabinet发送消息到:" + item.getKey() + "，报文:" + sendMsg);
                }

            } catch (IOException e) {
                continue;
            }
        }
    }

    public  synchronized void sendMessage(String message) throws IOException {
        this.session.getBasicRemote().sendText(message);
    }
    public  static  synchronized void sendMessage(CommWebSocketVO message)  {
        try{
            String sendMsg= JSON.toJSONString(message);
            for (Map.Entry<String, MessageWebsocket> item : webSocketMap.entrySet()) {
                try {
                    item.getValue().sendMessage(sendMsg);
                } catch (IOException e) {
                    continue;
                }
            }
        }catch (Exception e ){
            e.printStackTrace();
        }
    }

    //自定义群发消息
    public void sendInfo(String message) {

        for (Map.Entry<String, MessageWebsocket> item : webSocketMap.entrySet()) {
            try {
                item.getValue().sendMessage(message);
            } catch (IOException e) {
                continue;
            }
        }
    }

    //群推消息
    public static void sendMessage(CommWebSocketVO message , String... userIds) {
         List<String> userIdList = Arrays.asList(userIds);
         webSocketMap.forEach((userId,session)->{
             log.info("当前登录用户id" + userId);
             boolean flag = true;
             if (userIds == null || userIds.length == 0) {
             } else {
                 flag = userIdList.contains(userId);
             }
             if (flag) {
                 log.info("服务端给客户端[{}]发送消息{}",JSON.toJSONString(message));
                 try {
                     sendMessage(message, userId);
                 } catch (Exception e) {
                     e.printStackTrace();
                     log.error("pushMessage error", e);
                 }
             }
         });
    }
}