Java WebSocket实现前后端消息推送

WebSocket的代码编写会根据业务逻辑而进行变化，需要去理解编写思路，这样才能在工作中使用得游刃有余。

1. 引入依赖

<!-- websocket -->
<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-websocket</artifactId>
</dependency>

2.  编写WebSocketConfig配置类

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.server.standard.ServerEndpointExporter;

@Configuration
public class WebSocketConfig {

    @Bean
    public ServerEndpointExporter getServerEndpointExporter() {
        return new ServerEndpointExporter();
    }

}

3. 编写WebSocket服务类

下面的服务类中，可以编写多个sendMeg方法(写法比较多样化)，作用是发送消息回前端，使用方式就是你在自己的业务代码中自行调用，例（serviceImpl中调用）：

WebSocketServer.sendCountToUaa(ownerId, unreadNewsCount);

import com.sms.service.InternalNewsService;
import com.sms.service.impl.InternalNewsServiceImpl;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.context.ApplicationContext;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Websocket服务,推送信息
 */

@Slf4j
@Component
@ServerEndpoint("/webSocket/sms")
public class WebSocketServer {

    /**
     * 定义集合，用于存储连接信息
     * 根据业务需求自定义，数据格式可随意变化
     */
    //浏览器端连接，这种数据格式是因为考虑到同一个用户账号可以在不同的浏览器登录，所以一个用户需要保存多个连接
    private static final ConcurrentHashMap<String, ArrayList<Session>> sessionMap = new ConcurrentHashMap<>();
    //微信端连接
    private static final List<Session> wxSessionList = new ArrayList<>();

    /**
     * 作用：
     * 这里使用@Autowared是无法正常注入对象的，需要注入外部对象就需要使用这种方式
     * 需在启动类也就是main方法中加入：WebSocketServer.setApplicationContext(context);
     * context参数是启动类中run方法的返回值
     * 使用：
     * 在需要对象的地方使用
     * 例：
     * UserServiceImpl userServiceImpl = applicationContext.getBean(UserService.class);
     */
    private static ApplicationContext applicationContext;

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

    /**
     * 前端关闭页面或者主动关闭websocket链接，都会执行@OnClose
     */
    @OnClose
    public void close(Session session) {
        if (null == session) {
            return;
        }
        /**
         * 以下代码都是根据业务逻辑编写，写法并不固定
         */
        //关闭微信端中的session
        ListIterator<Session> sessionListIterator = wxSessionList.listIterator();
        while (sessionListIterator.hasNext()) {
            Session next = sessionListIterator.next();
            if (session == next) {
                sessionListIterator.remove();
                return;
            }
        }
        //关闭浏览器端中的session
        Iterator<Map.Entry<String, ArrayList<Session>>> iterator = sessionMap.entrySet().iterator();
        while (iterator.hasNext()) {
            ArrayList<Session> values = iterator.next().getValue();
            if (null == values || 1 > values.size()) {
                continue;
            }
            for (int i = 0; i < values.size(); i++) {
                Session sess = values.get(i);
                if (null == sess) {
                    continue;
                }
                if (session == sess) {
                    values.remove(i--);
                    return;
                }
            }
        }
    }

    /**
     * 前端连接后端socket时执行@OnMessage，前后端建立起连接，后端保存连接
     */
    @OnMessage
    public void onMessage(String userId, Session session) {
        if (StringUtils.isEmpty(userId) || null == session) {
            return;
        }
        if ("ykj".equals(userId)) {
            //微信端消息标识，自定义，没有微信端可直接将这个if删掉
            wxSessionList.add(session);
        } else {
            //接收到消息后，找到对应的session
            ArrayList<Session> sessions = sessionMap.get(userId);
            if (null == sessions) {
                sessions = new ArrayList<>();
            } else {
                //遍历，看该session是否存在，如果存在代表是在心跳检测操作
                for (int i = 0; i < sessions.size(); i++) {
                    if (sessions.get(i) == session) {
                        try {
                            //发送消息回前端
                            session.getAsyncRemote().sendText("{\"heartbeat\":\"socket心跳检测成功！！！\"}");
                            return;
                        } catch (Exception e) {
                            e.printStackTrace();
                        }
                    }
                }
            }
            //说明session不存在，添加到列表
            sessions.add(session);
            sessionMap.put(userId, sessions);

            //查询当前用户未读消息条数（业务代码，查询数据库）
            InternalNewsService internalNewsService = applicationContext.getBean(InternalNewsService.class);
            Integer newsCount = internalNewsService.selectUnreadNewsCount(userId, 0, 4);

            if (newsCount == null) {
                session.getAsyncRemote().sendText("0");//发送消息回前端
            } else {
                session.getAsyncRemote().sendText(newsCount.toString());//发送消息回前端
            }
            return;
        }
        try {
            session.getAsyncRemote().sendText("200");//发送消息回前端
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    /**
     * 发送数据到前端
     *
     * @param msg
     */
    public static void sendMsg(String userIds, String msg) {
        if (StringUtils.isEmpty(userIds) || StringUtils.isEmpty(msg)) {
            return;
        }
        String[] user_id = userIds.split(",");
        int uis = user_id.length;
        if (1 > uis) {
            return;
        }
        for (int ii = 0; ii < uis; ii++) {
            String userId = user_id[ii];
            if (StringUtils.isEmpty(userId)) {
                continue;
            }
            ArrayList<Session> sessions = sessionMap.get(userId);
            if (null == sessions || 1 > sessions.size()) {
                continue;
            }
            Iterator<Session> iterator = sessions.iterator();
            while (iterator.hasNext()) {
                Session next = iterator.next();
                if (null == next || !next.isOpen()) {
                    continue;
                }
                try {
                    next.getBasicRemote().sendText(msg);
                    //System.out.println("推送给socket：" + next);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * 发送数据到前端uaa端
     */
    public static void sendCountToUaa(String userId, Integer count) {
        if (StringUtils.isEmpty(userId)) {
            return;
        }
        ArrayList<Session> sessions = sessionMap.get(userId);
        if (null == sessions || 1 > sessions.size()) {
            return;
        }
        Iterator<Session> iterator = sessions.iterator();
        while (iterator.hasNext()) {
            Session next = iterator.next();
            if (null == next || !next.isOpen()) {
                continue;
            }
            try {
                //next.getBasicRemote().sendText(count.toString());
                next.getAsyncRemote().sendText(count.toString());
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 给微信端用户推送消息
     *
     * @param msg
     */
    public static void sendToVXMsg(String msg) {
        ListIterator<Session> sessionListIterator = wxSessionList.listIterator();
        while (sessionListIterator.hasNext()) {
            Session session = sessionListIterator.next();
            try {
                session.getBasicRemote().sendText(msg);
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

}