Springboot 集成websocket 并支持服务集群

1、新增配置类声明

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

@Configuration
public class WebsocketConfig {
    /**
     * 如果单元测试报错，请在类上加上以下注解内容
     * @SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
     * @return
     */
    @Bean
    public ServerEndpointExporter serverEndpointExporter(){
        return new ServerEndpointExporter();
    }

}

2、新建websocket连接类

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;

/**
 * html页面与之关联的接口
 * var reqUrl = "http://localhost:8080/ws/device/{deviceType}/{deviceAddress}";
 * socket = new WebSocket(reqUrl.replace("http", "ws"));
 */
@Slf4j
@Component
@ServerEndpoint("ws/device/{deviceType}/{deviceAddress}")
public class OwonWebSocketServerEndpoint {

    private String KEY;
    private String DEVICE_TYPE;

    @OnOpen
    public void onOpen(Session session, @PathParam("deviceType") String deviceType, @PathParam("deviceAddress") String deviceAddress) {
        log.info("发现设备连接，deviceType：" + deviceType + "，deviceAddress：" + deviceAddress);
        DeviceDO userDevice = SpringContextUtils.getBean(rDeviceMapper.class).findByDeviceTypeAndDeviceAddress(deviceType, deviceAddress);
        if (userDevice == null) {
            try {
                session.close();
            } catch (IOException e) {
                // ignore
            }
            return;
        }
        this.KEY = deviceType + WsSessionManager.SPLIT + deviceAddress;
        this.DEVICE_TYPE = deviceType;
        SpringContextUtils.getBean(WsSessionManager.class).add(KEY, session);
        log.info(String.format("成功建立连接, key:" + this.KEY));
    }

    @OnClose
    public void onClose() {
        SpringContextUtils.getBean(WsSessionManager.class).remove(this.KEY);
        log.info("成功关闭连接, key:" + KEY);
    }

    @OnMessage
    public void onMessage(Session session, String message) {
        log.info("收到消息, message:" + message);
    }

    @OnError
    public void onError(Session session, Throwable error) {
        log.info("发生错误:" + this.KEY);
        error.printStackTrace();
    }

    /**
     * 指定发消息
     *
     * @param message
     */
    public void sendMessage(String deviceType, String deviceAddress, String message) {
        SpringContextUtils.getBean(MessageSendManager.class).sendMessage(deviceType, deviceAddress, message);
    }

3、新建session管理类

import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;

import javax.websocket.Session;
import java.io.IOException;
import java.util.Collections;
import java.util.Enumeration;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;


@Slf4j
@Service
public class WsSessionManager {

    @Value("${eureka.instance.instance-id}")
    private String instanceId;
    @Autowired
    private JedisUtil jedisUtil;

    public static final String SPLIT = "@#@";

    /**
     * 保存连接 session 的地方
     */
    private static ConcurrentHashMap<String, Session> SESSION_POOL = new ConcurrentHashMap<>();

    /**
     * 添加 session
     *
     * @param key
     * @param session
     */
    public void add(String key, Session session) {
        // 添加 session
        SESSION_POOL.put(key, session);
        // 记录session所在的机器
        String url = getLocalServerUrl();
        jedisUtil.set(key, url, 24 * 3600);
    }

    /**
     * 设置链接机器地址
     *
     * @param key
     */
    public void setServerUrl(String key) {
        // 记录session所在的机器
        String url = getLocalServerUrl();
        jedisUtil.set(key, url, 24 * 3600);
    }

    /**
     * 删除 session,会返回删除的 session
     *
     * @param key
     * @return
     */
    public Session remove(String key) {

        // 删除 session
        Session session = SESSION_POOL.remove(key);
        // 删除记录的机器地址
        jedisUtil.del(key);
        return session;
    }

    /**
     * 删除并同步关闭连接
     *
     * @param key
     */
    public void removeAndClose(String key) {

        Session session = remove(key);
        if (session != null) {
            try {
                // 关闭连接
                session.close();
            } catch (IOException e) {
                // todo: 关闭出现异常处理
                e.printStackTrace();
            }
        }
    }

    /**
     * 获得 session
     *
     * @param key
     * @return
     */
    public Session get(String key) {

        // 获得 session
        return SESSION_POOL.get(key);
    }

    /**
     * 获取本机的地址
     *
     * @return
     */
    public String getLocalServerUrl() {
        return "http://" + instanceId;
    }

    /**
     * 组装session key
     *
     * @param deviceType 设备类型
     * @param devId  设备id
     * @return
     */
    public String getSessionKey(String deviceType, String devId) {
        return deviceType + SPLIT + devId;
    }

    /**
     * 获取redis 里面存储的链接地址
     *
     * @param sessionKey
     * @return
     */
    public String getServerUrl(String sessionKey) {
        return jedisUtil.get(sessionKey);
    }

    /**
     * 获取所有sessionKey
     *
     * @return
     */
    public List<String> getAllSessionKeys() {
        Enumeration<String> keys = SESSION_POOL.keys();
        if (keys.hasMoreElements()) {
            return Collections.list(keys);
        }
        return Lists.newArrayList();
    }

}

4、新建消息发送类

import cn.hutool.http.HttpUtil;
import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.websocket.Session;

@Slf4j
@Component
public class MessageSendManager {

    @Autowired
    private WsSessionManager wsSessionManager;

    /**
     * 发送消息，先看本机，本机没有链接就转发
     * @param deviceType
     * @param deviceAddress
     * @param message
     * @return
     */
    public Boolean sendMessage(String deviceType, String deviceAddress, String message) {
        //先尝试找本机
        Session session = wsSessionManager.get(wsSessionManager.getSessionKey(deviceType, deviceAddress));
        if (null != session) {
            synchronized (session) {
                try {
                    session.getAsyncRemote().sendText(message);
                    log.info("MessageSendManager sendMsg 消息发送成功 deviceType={},deviceAddress={},payload={}", deviceType, deviceAddress, JSON.toJSONString(message));
                    return true;
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return false;
            }
        } else {
            // 转发到链接所在的机器
            String url = wsSessionManager.getServerUrl(wsSessionManager.getSessionKey(deviceType, deviceAddress));
            if (StringUtils.isBlank(url)) {
                log.info("MessageSendManager sendMsg 找不到链接地址 deviceType={},deviceAddress={}", deviceType, deviceAddress);
                return false;
            }
            // 本机地址
            String localUrl = wsSessionManager.getLocalServerUrl();
            if (StringUtils.equals(url, localUrl)) {
                log.info("MessageSendManager sendMsg 本机找不到 deviceType={},deviceAddress={}", deviceType, deviceAddress);
                return false;
            }
            // 转发到其他机器
            transferByMsg(url, deviceType, deviceAddress, message);
            return true;
        }
    }

    /**
     * 发送消息，本机
     * @param message
     * @return
     */
    public Boolean sendMessageByKey(String key, String message) {
        //先尝试找本机
        Session session = wsSessionManager.get(key);
        if (null != session) {
            synchronized (session) {
                try {
                    session.getAsyncRemote().sendText(message);
                    log.info("MessageSendManager sendMsg 消息发送成功 key={},payload={}", key, JSON.toJSONString(message));
                    return true;
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return false;
            }
        }
        return false;
    }

    /**
     * 发送消息，先看本机，本机没有链接就转发
     * @param deviceType
     * @param deviceAddress
     * @param message
     * @return
     */
    public Boolean sendMsgToClient(String deviceType, String deviceAddress, String message) {
        //先尝试找本机
        Session session = wsSessionManager.get(wsSessionManager.getSessionKey(deviceType, deviceAddress));
        if (null != session) {
            synchronized (session) {
                try {
                    session.getAsyncRemote().sendText(message);
                    log.info("MessageSendManager sendMsg 消息发送成功 deviceType={},deviceAddress={},payload={}", deviceType, deviceAddress, JSON.toJSONString(message));
                    return true;
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return false;
            }
        }
        return false;
    }

    private void transferByMsg(String url, String deviceType, String deviceAddress, String message) {
        String urlString = url + "/device/msg/dispatch";
        HttpUtil.post(urlString, JSON.toJSONString(new WsMsgDispatchDTO(deviceType, deviceAddress, message)));
    }

}

5、新建转发消息接收类

@RestController
@RequestMapping("/device")
public class DeviceMsgDispatchController {

    @Autowired
    private MessageSendManager manager;

    /**
     * 消息转发处理
     * @param dto
     * @return
     */
    @RequestMapping(value = "/msg/dispatch", method = RequestMethod.POST, consumes = MediaType.APPLICATION_JSON_UTF8_VALUE)
    public BaseResult<Boolean> dispatch(@RequestBody WsMsgDispatchDTO dto) {
        return BaseResult.success(manager.sendMsgToClient(dto.getDeviceType(), dto.getDeviceAddress(), dto.getMessage()));
    }

}