websocket实现消息推送

闲来无事，写一篇关于websocket推送消息的文章，概念我们就不多讲了，直接上干货吧！！！

我们要做的就是：使用websocket来推送消息，但是如果说接收消息的那个用户没有连接上websocket，那么他就会接收不到我们的消息，这时候我们需要将消息存到数据库，等用户上线时自己查询数据库里面自己的消息，假如用户在线，我们就会将消息直接推送给他并且也保存在数据库，这时我们的前端未读消息数量应该+1，这个我们可以使用一个定时任务，比如隔5s我就像全部用户都推送消息，让前端触发消息回调的方法，我们可以在那个方法里面做手脚，就是去查询一遍我们的数据库中当前用户的未读消息，来达到消息数量的增减。

1.引入websocket依赖和一些其他依赖

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

<dependency>
    <groupId>org.projectlombok</groupId>
    <artifactId>lombok</artifactId>
    <version>1.18.16</version>
    <scope>compile</scope>
</dependency>

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

2.编写websocket配置类

@Component
public class WebSocketConfig {

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

3.编写websocket连接类

@Component
@ServerEndpoint("/webSocket/{userId}")
@Slf4j
public class WebSocket {
    private String userId;
    private Session session;
    private static Map<String, Session> sessionPool = new HashMap<String, Session>();
    private static CopyOnWriteArraySet<WebSocket> webSocketSet = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session,@PathParam(value = "userId") String userId) {
        this.session = session;
        this.userId = userId;
        webSocketSet.add(this);
        sessionPool.put(userId, session);
        log.info("【websocket消息】有新的连接, 总数:{}", webSocketSet.size());
    }

    @OnClose
    public void onClose() {
        webSocketSet.remove(this);
        sessionPool.remove(this.userId);
        log.info("【websocket消息】连接断开, 总数:{}", webSocketSet.size());
    }

    @OnMessage
    public void onMessage(String message) {
        log.info("【websocket消息】收到客户端发来的消息:{}", message);
    }


    /**
     * 服务端推送消息（单点发送）
     *
     * @param userId
     * @param message
     */
    public void pushMessage(String userId, String message) {
        Session session = sessionPool.get(userId);
        if (session != null && session.isOpen()) {
            try {
                log.info("【websocket消息】 单点消息:" + message);
                session.getAsyncRemote().sendText(message);
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }

    /**
     * 服务器端推送消息（群发）
     */
    public void pushMessage(String message) {
        try {
            webSocketSet.forEach(ws -> ws.session.getAsyncRemote().sendText(message));
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

4.编写两个控制器

@RequestMapping("/demo")
@RestController
public class DemoContrlller {
    @Autowired
    WebSocket webSocket;
    private Integer num = 1;

    /**
     * 单点发送
     * @param userId
     * @param message
     */
    @RequestMapping("/sendMessage")
    public void sendMessage(String userId, String message) {
        webSocket.pushMessage(userId, message);
    }

    /**
     * 群发
     * @param message
     */
    @RequestMapping("/send")
    public void send(String message) {
        webSocket.pushMessage(message);
    }

    /**
     * 模拟消息增加
     * @return
     */
    @RequestMapping("/add")
    public Integer num() {
        num++;
        return num;
    }

    /**
     * 模拟消息减少
     * @return
     */
    @RequestMapping("/reduce")
    public Integer jian() {
        num--;
        return num;
    }


}


// 用于访问到对应的视图
@Controller
@RequestMapping("/test")
public class TestController {
    @RequestMapping("/index")
    public String index() {
        return "index";
    }
}

前端代码

<!DOCTYPE html>
<html lang="en" xmlns:th="http://www.thymeleaf.org">
<head>
    <meta charset="UTF-8">
    <title>Title</title>
    <!-- 最新版本的 Bootstrap 核心 CSS 文件 -->
    <link rel="stylesheet" href="https://stackpath.bootstrapcdn.com/bootstrap/3.4.1/css/bootstrap.min.css" integrity="sha384-HSMxcRTRxnN+Bdg0JdbxYKrThecOKuH5zCYotlSAcp1+c8xmyTe9GYg1l9a69psu" crossorigin="anonymous">

    <!-- 可选的 Bootstrap 主题文件（一般不用引入） -->
    <link rel="stylesheet" href="https://stackpath.bootstrapcdn.com/bootstrap/3.4.1/css/bootstrap-theme.min.css" integrity="sha384-6pzBo3FDv/PJ8r2KRkGHifhEocL+1X2rVCTTkUfGk7/0pbek5mMa1upzvWbrUbOZ" crossorigin="anonymous">

    <!-- 最新的 Bootstrap 核心 JavaScript 文件 -->
    <script src="https://stackpath.bootstrapcdn.com/bootstrap/3.4.1/js/bootstrap.min.js" integrity="sha384-aJ21OjlMXNL5UyIl/XNwTMqvzeRMZH2w8c5cRVpzpU8Y5bApTppSuUkhZXN0VxHd" crossorigin="anonymous"></script>
</head>
<body>

<div class="btn btn-default">
    <span class="glyphicon glyphicon-envelope" aria-hidden="true"></span>
    <span  class="badge badge-danger" style="background-color:#FF0000;color:#FFFFFF " id ="ms"></span>
</div>
<button id="add" onclick="add()">添加</button>
<button id="reduce" onclick="reduce()">减少</button>
<div id="message"></div>
<script src="https://cdn.bootcss.com/jquery/1.12.4/jquery.min.js"></script>
<script src="https://cdn.bootcss.com/bootstrap/3.3.5/js/bootstrap.min.js"></script>
<script>
    var websocket = null;
    if ('WebSocket' in window) {
        websocket = new WebSocket('ws://192.168.1.189:8001/webSocket/1');
    } else {
        alert('该浏览器不支持websocket!');
    }

    websocket.onopen = function (event) {
        console.log('建立连接');
    }

    websocket.onclose = function (event) {
        console.log('连接关闭');
    }

    websocket.onmessage = function (event) {
        // 收到消息
        console.log('收到消息:' + event.data)
        setMessageInnerHTML(event.data);
        f();
    }

    websocket.onerror = function () {
        alert('websocket通信发生错误！');
    }

    window.onbeforeunload = function () {
        websocket.close();
    }
    function setMessageInnerHTML(innerHTML) {
        document.getElementById('message').innerHTML = innerHTML + '<br/>';
    }
    function f() {
        $.ajax({
            type: 'post',
            dataType: 'text',
            url: '/demo/num',
            data: {},
            cache: false,
            async: true,
            success: function (data) {
                var data = eval('(' + data + ')');
                console.log('数量:' + data)
                if(data>0){
                    document.getElementById('ms').innerHTML = data + '<br/>';
                }

            }
        });
    }
    function add() {
        $.ajax({
            type: 'post',
            dataType: 'text',
            url: '/demo/add',
            data: {},
            cache: false,
            async: true,
            success: function (data) {
                if(data>0){
                    document.getElementById('ms').innerHTML = data + '<br/>';
                }
            }
        });
    }
    function reduce() {
        $.ajax({
            type: 'post',
            dataType: 'text',
            url: '/demo/reduce',
            data: {},
            cache: false,
            async: true,
            success: function (data) {
                if(data>0){
                    document.getElementById('ms').innerHTML = data + '<br/>';
                }else{
                    document.getElementById('ms').innerHTML ='';
                }
            }
        });

    }
</script>
</body>
</html>