springboot中使用websocket

springboot中websocket的使用

1. 添加依赖

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

2. 添加配置类

// 通过@EnableWebSocket注解，启动spring-boot-starter-websocket的自动化配置
@Configuration
@EnableWebSocket
public class WebSocketConfig implements WebSocketConfigurer {

    @Bean
    public WebSocketHandler sysWebSocketHandler(){
        return new WebSocketHandler();
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 注册自定义的WebSocketHandler的实现类,并指定类对应的websocket访问的ServerEndpoint为//websocket/*
        registry.addHandler(sysWebSocketHandler(), Constants.WebSocket.PATH_PREFIX + "/*")
        // 设置允许跨域访问
                .setAllowedOrigins("*");
    }
}

3. 添加消息处理器

@Component
public class WebSocketHandler implements WebSocketHandler {

    private static final Logger log = LoggerFactory.getLogger(WebSocketHandler.class);

    @Resource
    private IWebSocketService websocketService;

    /**
     * 建立连接
     * @param session
     * @throws Exception
     */
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        websocketService.handleOpen(session);
    }

    /**
     * 接收消息
     * @param session
     * @param message
     * @throws Exception
     */
    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        if (message instanceof TextMessage){
            TextMessage textMessage = (TextMessage) message;
            websocketService.handleMessage(session,textMessage.getPayload());
        }
    }

    /**
     * 发生错误
     * @param session
     * @param exception
     * @throws Exception
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        websocketService.handleError(session,exception);
    }

    /**
     * 关闭连接
     * @param session
     * @param closeStatus
     * @throws Exception
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        websocketService.handleClose(session);
    }

    /**
     * 是否支持发送部分消息
     * @return
     */
    @Override
    public boolean supportsPartialMessages() {
        return false;
    }
}

4.websocket service接口

package com.fuj.wms.common.core.service;

import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.util.Map;

/**
 * websocket service接口
 */
public interface IWebSocketService {

    /**
     * 会话开始回调
     *
     * @param session 会话
     */
    void handleOpen(WebSocketSession session);

    /**
     * 会话结束回调
     *
     * @param session 会话
     */
    void handleClose(WebSocketSession session);

    /**
     * 处理消息
     *
     * @param session 会话
     * @param message 接收的消息
     */
    void handleMessage(WebSocketSession session, String message);

    /**
     * 发送消息
     *
     * @param session 当前会话
     * @param message 要发送的消息
     */
    void sendMessage(WebSocketSession session, String message);

    /**
     * 发送消息
     *
     * @param sessionId  用户id
     * @param message 要发送的消息
     */
    void sendMessage(String sessionId, TextMessage message);

    /**
     * 发送消息
     *
     * @param sessionId  用户id
     * @param message 要发送的消息
     */
    void sendMessage(String sessionId, String message);

    /**
     * 发送消息
     *
     * @param session 当前会话
     * @param message 要发送的消息
     */
    void sendMessage(WebSocketSession session, TextMessage message);

    /**
     * 广播
     *
     * @param message 字符串消息
     */
    void broadCast(String message);

    /**
     * 广播
     *
     * @param message 文本消息
     */
    void broadCast(TextMessage message);

    /**
     * 排除sessionIds广播
     * @param message 字符串消息
     * @param sessionIds 排除的session
     */
    void broadCastExclude(String message,String... sessionIds);

    /**
     * 排除sessionIds广播
     * @param message 文本消息
     * @param sessionIds 排除的session
     */
    void broadCastExclude(TextMessage message,String... sessionIds);

    /**
     * 处理会话异常
     *
     * @param session 会话
     * @param error   异常
     */
    void handleError(WebSocketSession session, Throwable error);

    /**
     * 主动关闭websocket
     * @param session 会话
     */
    void closeWebSocket(WebSocketSession session);

    /**
     * 主动关闭websocket
     * @param sessionId 会话id
     */
    void closeWebSocket(String sessionId);

    /**
     * 获得所有的 websocket 会话
     *
     * @return 所有 websocket 会话
     */
    Map<String, WebSocketSession> getSessions();

    /**
     * 得到当前连接数
     *
     * @return 连接数
     */
    int getConnectionCount();

    /**
     * 获取请求头
     * @param session
     * @return
     */
    String getRequest(WebSocketSession session);

    /**
     * 获取客户端IP地址
     * @param session
     * @return
     */
    String getHostAddress(WebSocketSession session);

    /**
     * 获取客户端电脑名称
     * @param session
     * @return
     */
    String getHostName(WebSocketSession session);
}

5.websocket service接口实现类

package com.fuj.wms.common.core.service;

import com.fuj.wms.common.constant.CacheConstants;
import com.fuj.wms.common.constant.Constants;
import com.fuj.wms.common.core.redis.RedisCache;
import com.fuj.wms.common.utils.StringUtils;
import com.fuj.wms.common.utils.spring.SpringUtils;
import com.fuj.wms.common.core.domain.model.FujWebSocket;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.ConcurrentWebSocketSessionDecorator;

import java.io.IOException;
import java.util.Arrays;
import java.util.Date;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * websocket 处理业务基础实现
 */
public abstract class WebSocketService implements IWebSocketService {

    private static final Logger log = LoggerFactory.getLogger(WebSocketService.class);

    /**
     * 静态变量，用来记录当前在线连接数，线程安全的类。
     */
    protected static AtomicInteger connectionCount = new AtomicInteger(0);

    /**
     * 存放所有在线的客户端
     */
    protected static Map<String, WebSocketSession> sessionsMap = new ConcurrentHashMap<>();

    /**
     * 发送时间限制(单位:毫秒)
     */
    protected final int sendTimeLimit = 10 * 1000 * 1000;

    /**
     * 发送长度限制(单位：bytes)
     */
    protected final int bufferSizeLimit = 1024 * 1024 * 1024;

    @Override
    public void handleOpen(WebSocketSession session) {
        log.info("WebSocket连接建立中 ==> session_id = {},IP地址 = {},PC名称 = {},uri = {}", session.getId(), getHostAddress(session), getHostName(session), getRequest(session));
        sessionsMap.put(session.getId(), new ConcurrentWebSocketSessionDecorator(session, sendTimeLimit, bufferSizeLimit));
        connectionCount.incrementAndGet();
//        this.sendMessage(session, "连接成功");
        log.info("WebSocket连接建立成功,当前在线数为: {} ==> 开始监听新连接：session_id = {}。", connectionCount.get(), session.getId());

        //存入redis
        //存map
        Map<String, FujWebSocket> webSocketMap = SpringUtils.getBean(RedisCache.class).getCacheMap(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_MAP_KEY);
        if (StringUtils.isEmpty(webSocketMap)) {
            webSocketMap = new ConcurrentHashMap<>();
            SpringUtils.getBean(RedisCache.class).setCacheMap(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_MAP_KEY, webSocketMap);
        }
        SpringUtils.getBean(RedisCache.class).setCacheMapValue(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_MAP_KEY, session.getId(), createFujWebSocket(session));

        //存数量
        Object count = SpringUtils.getBean(RedisCache.class).getCacheObject(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_COUNT_KEY);
        if (StringUtils.isNull(count)) {
            SpringUtils.getBean(RedisCache.class).setCacheObject(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_COUNT_KEY, Constants.Number.INTEGER_DEFAULT_VALUE);
        }
        SpringUtils.getBean(RedisCache.class).incr(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_COUNT_KEY);
    }

    @Override
    public void handleClose(WebSocketSession session) {
        log.info("WebSocket连接关闭中 ==> session_id = {},IP地址 = {},PC名称 = {},接口 = {}", session.getId(), getHostAddress(session), getHostName(session), getRequest(session));
        sessionsMap.remove(session.getId());
        connectionCount.decrementAndGet();
        log.info("WebSocket连接关闭成功,当前在线数为: {} ==> 已关闭连接：session_id = {}。", connectionCount.get(), session.getId());

        //从redis中删除
        //删除map
        SpringUtils.getBean(RedisCache.class).deleteCacheMapValue(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_MAP_KEY, session.getId());
        //数量-1
        SpringUtils.getBean(RedisCache.class).decr(CacheConstants.WMS_KEY + CacheConstants.Monitor.WEBSOCKET_COUNT_KEY);
    }

    @Override
    public void handleMessage(WebSocketSession session, String message) {
        log.info("WebSocket服务端收到客户端消息 ==> session_id = {},IP地址 = {},PC名称 = {},接口 = {},message = {}", session.getId(), getHostAddress(session), getHostName(session), getRequest(session), message);
//        this.sendMessage(session, "收到消息");
    }

    @Override
    public void sendMessage(WebSocketSession session, String message) {
        try {
            log.info("WebSocket发送消息 ==> session_id = {},IP地址 = {},PC名称 = {},接口 = {},message = {}", session.getId(), getHostAddress(session), getHostName(session), getRequest(session), message);
            session.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void sendMessage(String sessionId, TextMessage message) {
        WebSocketSession webSocketSession = sessionsMap.get(sessionId);
        this.sendMessage(webSocketSession, message);
    }

    @Override
    public void sendMessage(String sessionId, String message) {
        WebSocketSession webSocketSession = sessionsMap.get(sessionId);
        this.sendMessage(webSocketSession, message);
    }

    @Override
    public void sendMessage(WebSocketSession session, TextMessage message) {
        try {
            log.info("WebSocket发送消息 ==> session_id = {},IP地址 = {},PC名称 = {},接口 = {},message = {}", session.getId(), getHostAddress(session), getHostName(session), getRequest(session), message.getPayload());
            session.sendMessage(message);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void broadCast(String message) {
        sessionsMap.forEach((sessionId, toSession) -> {
            if (toSession.isOpen()) {
                this.sendMessage(toSession, message);
            }
        });
    }

    @Override
    public void broadCast(TextMessage message) {
        sessionsMap.forEach((sessionId, toSession) -> {
            if (toSession.isOpen()) {
                this.sendMessage(toSession, message);
            }
        });
    }

    @Override
    public void broadCastExclude(String message, String... sessionIds) {
        sessionsMap.forEach((sessionId, toSession) -> {
            if (!Arrays.asList(sessionIds).contains(sessionId)){
                if (toSession.isOpen()) {
                    this.sendMessage(toSession, message);
                }
            }
        });
    }

    @Override
    public void broadCastExclude(TextMessage message, String... sessionIds) {
        sessionsMap.forEach((sessionId, toSession) -> {
            if (!Arrays.asList(sessionIds).contains(sessionId)){
                if (toSession.isOpen()) {
                    this.sendMessage(toSession, message);
                }
            }
        });
    }

    @Override
    public void handleError(WebSocketSession session, Throwable error) {
        log.error("WebSocket发生错误 ==> session_id = {},IP地址 = {},PC名称 = {},接口 = {},错误信息为 = {}", session.getId(), getHostAddress(session), getHostName(session), getRequest(session), error.getMessage());
    }

    @Override
    public void closeWebSocket(WebSocketSession session) {
        try {
            if (session.isOpen()) {
                session.close();
            }
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    @Override
    public void closeWebSocket(String sessionId) {
        WebSocketSession webSocketSession = sessionsMap.get(sessionId);
        this.closeWebSocket(webSocketSession);
    }

    @Override
    public Map<String, WebSocketSession> getSessions() {
        return sessionsMap;
    }

    @Override
    public int getConnectionCount() {
        return connectionCount.get();
    }

    @Override
    public String getRequest(WebSocketSession session) {
        String path = session.getUri().getPath();
        return StringUtils.replace(path, Constants.WebSocket.PATH_PREFIX + "/", "").trim();
    }

    @Override
    public String getHostAddress(WebSocketSession session) {
        return session.getRemoteAddress().getAddress().getHostAddress();
    }

    @Override
    public String getHostName(WebSocketSession session) {
        return session.getRemoteAddress().getHostName();
    }

    /**
     * 创建FujWebSocket实体 存入redis
     *
     * @param session
     * @return
     */
    private FujWebSocket createFujWebSocket(WebSocketSession session) {
        FujWebSocket fujWebSocket = new FujWebSocket(session.getId(), getRequest(session), getHostAddress(session), getHostName(session), session.getUri().toString());
        fujWebSocket.setCreateTime(new Date());
        fujWebSocket.setCreateBy(Constants.User.WMS);
        return fujWebSocket;
    }
}