springboot整合websocket实现点对点和全体发消息

背景 最近项目里有个类似备忘录的功能，每个用户管理多个客户，每个客户都有一个备忘录，一旦员工管理的客户多了，备忘录自然也多了，往往这时员工就搞不清啥时候该干啥了，所以要做一个类似小秘书的功能，按照备忘录里设置执行时间前30分钟通知一下员工，当然其他的消息都可以通知。
功能介绍 员工会针对某个客户添加一个执行计划，哪天几点上门/电话/短信拜访客户或其他事项。如果说我们设置在执行时间前30分钟后台主动推送消息到前台通知员工的话这时需要用到的技术是：延时队列+websocket
延时队列 我这边用的是rabbitmq，添加执行计划成功后，把执行时间减半小时压入延时队列里，等待到期消费，到期后调用websocket向前台发消息通知员工。
当然也可以写定时任务1min轮训一次，但这样太low。
**websocket：**springboot整合websocket
由于本篇主要讲解springboot整合websocket，所以延时队列怎么实现暂时不做赘述：

1. maven先加入依赖：

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

2. 写个websocket配置类：

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

/**
 * WebSocket配置
 * @author zhanghang
 * @date 2019/1/3 17:53
 */
@Component
public class WebSocketConfig {

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

3. 由于我们的websocket对象中肯定需要用到用户信息，但是WsSession终是没有的，只能从HttpSession中拿，所以写一个获取HttpSession的配置类

import javax.servlet.http.HttpSession;
import javax.websocket.HandshakeResponse;
import javax.websocket.server.HandshakeRequest;
import javax.websocket.server.ServerEndpointConfig;

/**
 * 协助server获取http session
 * @author zhanghang
 * @date 2019/1/4 13:58
 */
public class HttpSessionWSHelper extends ServerEndpointConfig.Configurator {
    @Override
    public void modifyHandshake(ServerEndpointConfig sec,HandshakeRequest request, HandshakeResponse response) {
        HttpSession httpSession=(HttpSession) request.getHttpSession();
        sec.getUserProperties().put(HttpSession.class.getName(),httpSession);
    }
}

4. 然后再写websocket的具体实现类

import com.ccq.calling.config.shiro.ShiroUser;
import com.ccq.calling.utils.MyApp;
import lombok.extern.slf4j.Slf4j;
import org.apache.shiro.subject.SimplePrincipalCollection;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpSession;
import javax.websocket.*;
import javax.websocket.server.ServerEndpoint;
import java.util.concurrent.CopyOnWriteArraySet;

/**
 * websocket实现
 * @author zhanghang
 * @date 2019/1/3 17:54
 */
@Component
@ServerEndpoint(value = "/websocket",configurator = HttpSessionWSHelper.class)
@Slf4j
public class WebSocket {
    private Integer userId;
    private Session session;
    private static CopyOnWriteArraySet<WebSocket> webSockets = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session) {
        HttpSession httpSession = (HttpSession) session.getUserProperties().get(HttpSession.class.getName());
        ShiroUser shiroUser = (ShiroUser)((SimplePrincipalCollection) httpSession.getAttribute(MyApp.HTTPSESSION_USER_ATTR)).getPrimaryPrincipal();
        this.userId = shiroUser.getId();
        this.session = session;
        webSockets.add(this);
        log.info("【websocket消息】有新的连接, 总数:{}", webSockets.size());
    }


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

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

    /**
     * 点对点广播消息
     * @param userId 要广播的用户id
     * @param msg 消息
     */
    public void sendP2PMessage(Integer userId,String msg){
        for (WebSocket webSocket: webSockets) {
            log.info("【websocket消息】点对点广播消息, userId={},message={}", userId,msg);
            try {
                if (webSocket.userId.intValue() == userId.intValue()) {
                    webSocket.session.getBasicRemote().sendText(msg);
                }
            } catch (Exception e) {
                log.error("=================websocket点对点广播消息出错===================");
                log.error(e.getMessage(),e);
            }
        }
    }

    /**
     * 全体广播消息
     * @param msg 消息
     */
    public void sendMessage(String msg){
        for (WebSocket webSocket: webSockets) {
            log.info("【websocket消息】全体广播消息,message={}",msg);
            try {
                webSocket.session.getBasicRemote().sendText(msg);
            } catch (Exception e) {
                log.error("=================websocket全体广播消息出错===================");
                log.error(e.getMessage(),e);
            }
        }
    }
}

5. 前端代码：

	<script>
        var websocket = null;
        if('WebSocket' in window){
            websocket = new WebSocket('ws://localhost:8888/websocket');
        }else{
            alert("该浏览器不支持WebSocket");
        }

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

        websocket.onclose = function (event) {
            console.log("断开连接");
        }

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

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

        window.onbeforeunload = function (event) {
            websocket.close();
        }
    </script>

6. 然后我们再写个测试类去测试一下：

import com.ccq.calling.config.websocket.WebSocket;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PathVariable;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

/**
 * @author zhanghang
 * @date 2019/1/3 17:57
 */
@RestController
@RequestMapping("/test")
public class TestController {

    @Autowired
    private WebSocket webSocket;

    @GetMapping("/sendP2PMessage/{userId}/{msg}")
    public void sendP2PMessage(@PathVariable("userId") Integer userId,@PathVariable("msg") String msg){
        webSocket.sendP2PMessage(userId,msg);
    }

    @GetMapping("/sendMessage/{msg}")
    public void sendMessage(@PathVariable("msg") String msg){
        webSocket.sendMessage(msg);
    }
}